forked from molecule-ai/molecule-core
fix(security): GLOBAL memory prompt injection safeguards (#767)
Two defenses against GLOBAL-scope agent memory injection attacks:
1. Recall delimiter: Search() wraps every GLOBAL-scope memory value
with a non-instructable prefix before returning it to MCP clients:
[MEMORY id=<uuid> scope=GLOBAL from=<workspace_id>]: <value>
This prevents stored content (e.g. "IGNORE ALL PREVIOUS INSTRUCTIONS")
from being parsed as instructions in the agent's context window.
Raw DB content is unchanged — the wrapper is applied on read only.
2. Write audit log: Commit() writes an activity_log entry with
activity_type='memory_write_global' whenever a GLOBAL memory is
stored. The entry records a SHA-256 hash of the content (never
plaintext) alongside memory_id and namespace for forensic replay.
Audit failure is non-fatal — a logging error must not roll back
a successful write.
Tests:
- TestRecallMemory_GlobalScope_HasDelimiter — verifies exact delimiter
format [MEMORY id=... scope=GLOBAL from=...]: <value>
- TestCommitMemory_GlobalScope_AuditLogEntry — verifies activity_logs
INSERT fires on every GLOBAL write (via mock.ExpectationsWereMet)
- TestMemoriesCommit_Global_AsRoot — updated to expect the audit INSERT
All 16 Go test packages pass.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
705c0a46ce
commit
8d01a2a09c
@ -1,6 +1,9 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
@ -10,6 +13,12 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// globalMemoryDelimiter is the non-instructable prefix prepended to every
|
||||
// GLOBAL-scope memory value returned to MCP clients. Prevents stored content
|
||||
// from being parsed as LLM instructions in the agent's context window (#767).
|
||||
// Format: [MEMORY id=<uuid> scope=GLOBAL from=<workspace_id>]: <value>
|
||||
const globalMemoryDelimiter = "[MEMORY id=%s scope=GLOBAL from=%s]: %s"
|
||||
|
||||
// defaultMemoryNamespace is used when a caller omits the field on POST or
|
||||
// when querying for memories written before migration 017. Matches the
|
||||
// column default in platform/migrations/017_memories_fts_namespace.up.sql.
|
||||
@ -81,6 +90,26 @@ func (h *MemoriesHandler) Commit(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// #767 Audit: write a GLOBAL memory audit log entry for forensic replay.
|
||||
// Records a SHA-256 hash of the content — never plaintext — so the audit
|
||||
// trail can prove what was written without leaking sensitive values.
|
||||
// Failure is non-fatal: a logging error must not roll back a successful write.
|
||||
if body.Scope == "GLOBAL" {
|
||||
sum := sha256.Sum256([]byte(body.Content))
|
||||
auditBody, _ := json.Marshal(map[string]string{
|
||||
"memory_id": memoryID,
|
||||
"namespace": namespace,
|
||||
"content_sha256": hex.EncodeToString(sum[:]),
|
||||
})
|
||||
summary := "GLOBAL memory written: id=" + memoryID + " namespace=" + namespace
|
||||
if _, auditErr := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO activity_logs (workspace_id, activity_type, source_id, summary, request_body, status)
|
||||
VALUES ($1, $2, $3, $4, $5::jsonb, $6)
|
||||
`, workspaceID, "memory_write_global", workspaceID, summary, string(auditBody), "ok"); auditErr != nil {
|
||||
log.Printf("Commit: GLOBAL memory audit log failed for %s/%s: %v", workspaceID, memoryID, auditErr)
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusCreated, gin.H{"id": memoryID, "scope": body.Scope, "namespace": namespace})
|
||||
}
|
||||
|
||||
@ -211,6 +240,14 @@ func (h *MemoriesHandler) Search(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// #767: wrap GLOBAL-scope content with a non-instructable delimiter so
|
||||
// MCP tool outputs cannot be hijacked by stored prompt-injection payloads.
|
||||
// The raw content in the DB is unchanged — only the value returned to
|
||||
// callers is wrapped.
|
||||
if memScope == "GLOBAL" {
|
||||
content = fmt.Sprintf(globalMemoryDelimiter, id, wsID, content)
|
||||
}
|
||||
|
||||
memories = append(memories, map[string]interface{}{
|
||||
"id": id,
|
||||
"workspace_id": wsID,
|
||||
|
||||
@ -60,6 +60,10 @@ func TestMemoriesCommit_Global_AsRoot(t *testing.T) {
|
||||
WithArgs("root-ws", "global fact", "GLOBAL", "general").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("mem-global"))
|
||||
|
||||
// #767: GLOBAL writes always produce an audit log entry.
|
||||
mock.ExpectExec("INSERT INTO activity_logs").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "root-ws"}}
|
||||
@ -72,6 +76,9 @@ func TestMemoriesCommit_Global_AsRoot(t *testing.T) {
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Errorf("expected 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoriesCommit_Global_ForbiddenForChild(t *testing.T) {
|
||||
@ -605,3 +612,99 @@ func TestMemoriesSearch_LimitDefault_Is50(t *testing.T) {
|
||||
t.Errorf("sqlmock expectations not met (default limit should be 50): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- Issue #767: GLOBAL memory prompt injection safeguards ----------
|
||||
|
||||
// TestRecallMemory_GlobalScope_HasDelimiter verifies that GLOBAL-scope
|
||||
// memories returned by Search are wrapped with the non-instructable
|
||||
// [MEMORY id=... scope=GLOBAL from=...]: prefix. This prevents stored
|
||||
// content from being interpreted as LLM instructions by MCP tool outputs.
|
||||
func TestRecallMemory_GlobalScope_HasDelimiter(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
handler := NewMemoriesHandler()
|
||||
|
||||
// Parent lookup (needed by Search for access-control branching)
|
||||
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id").
|
||||
WithArgs("ws-reader").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(nil))
|
||||
|
||||
rows := sqlmock.NewRows([]string{"id", "workspace_id", "content", "scope", "namespace", "created_at"}).
|
||||
AddRow("mem-g1", "root-ws", "global knowledge", "GLOBAL", "general", "2024-01-01T00:00:00Z")
|
||||
|
||||
mock.ExpectQuery("SELECT id, workspace_id, content, scope, namespace, created_at FROM agent_memories WHERE scope = 'GLOBAL'").
|
||||
WillReturnRows(rows)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-reader"}}
|
||||
c.Request = httptest.NewRequest("GET", "/memories?scope=GLOBAL", nil)
|
||||
c.Request.URL.RawQuery = "scope=GLOBAL"
|
||||
|
||||
handler.Search(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var result []map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatalf("body not valid JSON: %v", err)
|
||||
}
|
||||
if len(result) != 1 {
|
||||
t.Fatalf("expected 1 memory in result, got %d", len(result))
|
||||
}
|
||||
|
||||
content, _ := result[0]["content"].(string)
|
||||
want := "[MEMORY id=mem-g1 scope=GLOBAL from=root-ws]: global knowledge"
|
||||
if content != want {
|
||||
t.Errorf("GLOBAL content delimiter missing or incorrect\ngot: %q\nwant: %q", content, want)
|
||||
}
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet sqlmock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCommitMemory_GlobalScope_AuditLogEntry verifies that writing a
|
||||
// GLOBAL-scope memory always produces an activity_log entry with
|
||||
// event_type='memory_write_global'. The audit entry stores the SHA-256
|
||||
// content hash (never plaintext) for forensic replay.
|
||||
func TestCommitMemory_GlobalScope_AuditLogEntry(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
handler := NewMemoriesHandler()
|
||||
|
||||
// Root workspace — allowed to write GLOBAL
|
||||
mock.ExpectQuery("SELECT parent_id FROM workspaces WHERE id").
|
||||
WithArgs("root-ws").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"parent_id"}).AddRow(nil))
|
||||
|
||||
mock.ExpectQuery("INSERT INTO agent_memories").
|
||||
WithArgs("root-ws", "sensitive global fact", "GLOBAL", "general").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("mem-audit"))
|
||||
|
||||
// KEY ASSERTION: GLOBAL write must produce an audit log entry.
|
||||
// We match on the SQL prefix; the exact arguments (content hash, etc.)
|
||||
// are validated by the implementation — here we verify the INSERT fires.
|
||||
mock.ExpectExec("INSERT INTO activity_logs").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "root-ws"}}
|
||||
body := `{"content":"sensitive global fact","scope":"GLOBAL"}`
|
||||
c.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
handler.Commit(c)
|
||||
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Errorf("expected 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
// ExpectationsWereMet fails if the audit INSERT was not called —
|
||||
// that's the primary assertion of this test.
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("GLOBAL memory write must produce audit log entry: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user