From 8d01a2a09c8cd7b2b571bfd2c014b9efa45b5115 Mon Sep 17 00:00:00 2001 From: "molecule-ai[bot]" <276602405+molecule-ai[bot]@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:26:46 +0000 Subject: [PATCH] fix(security): GLOBAL memory prompt injection safeguards (#767) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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= scope=GLOBAL from=]: 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=...]: - 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 --- platform/internal/handlers/memories.go | 37 +++++++ platform/internal/handlers/memories_test.go | 103 ++++++++++++++++++++ 2 files changed, 140 insertions(+) diff --git a/platform/internal/handlers/memories.go b/platform/internal/handlers/memories.go index cee4f890..85ae47dc 100644 --- a/platform/internal/handlers/memories.go +++ b/platform/internal/handlers/memories.go @@ -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= scope=GLOBAL from=]: +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, diff --git a/platform/internal/handlers/memories_test.go b/platform/internal/handlers/memories_test.go index 4324a62f..80bc84c8 100644 --- a/platform/internal/handlers/memories_test.go +++ b/platform/internal/handlers/memories_test.go @@ -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) + } +}