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:
molecule-ai[bot] 2026-04-17 16:26:46 +00:00 committed by GitHub
parent 705c0a46ce
commit 8d01a2a09c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 140 additions and 0 deletions

View File

@ -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,

View File

@ -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)
}
}