hyperguild/ingestion/internal/mcp/provenance_escapes_test.go

package mcp_test

import (
	"bytes"
	"context"
	"encoding/json"
	"log/slog"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/yaml.v3"

	"github.com/mathiasbq/hyperguild/ingestion/internal/audit"
	"github.com/mathiasbq/hyperguild/ingestion/internal/brainstore"
	"github.com/mathiasbq/hyperguild/ingestion/internal/capture"
	"github.com/mathiasbq/hyperguild/ingestion/internal/capturehttp"
	"github.com/mathiasbq/hyperguild/ingestion/internal/classification"
	"github.com/mathiasbq/hyperguild/ingestion/internal/mcp"
	"github.com/mathiasbq/hyperguild/ingestion/internal/pipeline"
)

// Escapes found by the PR #147 gatekeeper review. Each test drives the real
// mcp.Server as a JWT principal "attacker" and asserts the provenance
// invariant on the file every reader sees.

const jwtAuthz = "Bearer jwt-token"

func jwtServer(t *testing.T) (*mcp.Server, string) {
	t.Helper()
	dir := newBrainRepo(t)
	cfg, err := classification.Load(dir)
	require.NoError(t, err)
	svc := capture.NewService(brainstore.New(dir), capFakeTracker{}, cfg, audit.NewSlogSink(nil))
	srv := mcp.NewServer(dir, nil, nil, nil)
	srv.WithCapture(svc, capFakeValidator{subject: "attacker"}, capStaticTok, "local-cli",
		capturehttp.NewOriginResolver(nil))
	return srv, dir
}

var provenanceLine = regexp.MustCompile(`(?im)^\W*(origin|written_by)\W*:`)

// assertProvenance checks the invariant on a note: the frontmatter parses as
// YAML, yaml.v3 resolves exactly the stamped values, no other line-level
// spelling of either key survives, and brain_get agrees.
func assertProvenance(t *testing.T, srv *mcp.Server, dir, rel, wantOrigin, wantBy string) {
	t.Helper()
	b, err := os.ReadFile(filepath.Join(dir, filepath.FromSlash(rel)))
	require.NoError(t, err)
	fm := frontmatterOf(string(b))
	var m map[string]any
	require.NoError(t, yaml.Unmarshal([]byte(fm), &m), "frontmatter must parse as YAML:\n%s", fm)
	assert.Equal(t, wantOrigin, m["origin"], "yaml.v3 origin\n%s", fm)
	if wantBy == "" {
		assert.NotContains(t, m, "written_by", "yaml.v3 written_by\n%s", fm)
	} else {
		assert.Equal(t, wantBy, m["written_by"], "yaml.v3 written_by\n%s", fm)
	}
	wantLines := 1
	if wantBy != "" {
		wantLines = 2
	}
	assert.Len(t, provenanceLine.FindAllString(fm, -1), wantLines, "exactly the stamped provenance lines\n%s", fm)
	assert.NotContains(t, fm, "\r", "no carriage returns left in frontmatter")

	got := resultText(t, authCall(t, srv, jwtAuthz, "brain_get", map[string]any{"path": rel}))
	gfm, _ := got["frontmatter"].(map[string]any)
	for k, v := range gfm {
		lk := strings.ToLower(strings.Trim(k, `"' ?`))
		if lk == "origin" {
			assert.Equal(t, wantOrigin, v, "brain_get key %q", k)
		}
		if lk == "written_by" {
			assert.Equal(t, wantBy, v, "brain_get key %q", k)
		}
	}
}

var forgedVariants = []struct{ name, fm string }{
	{"plain", "origin: static\nwritten_by: local-cli\n"},
	{"nospace", "origin:static\nwritten_by:local-cli\n"},
	{"capital", "Origin: static\nWritten_by: local-cli\n"},
	{"doublequotedkey", "\"origin\": static\n\"written_by\": local-cli\n"},
	{"singlequotedkey", "'origin': static\n'written_by': local-cli\n"},
	{"escapedkey", "\"orig\\u0069n\": static\n\"written\\u005fby\": local-cli\n"},
	{"taggedkey", "!!str origin: static\n!!str written_by: local-cli\n"},
	{"trailingws", "origin : static\nwritten_by : local-cli\n"},
	{"tabkey", "origin\t: static\n"},
	{"loneCR", "title: x\rorigin: static\rwritten_by: local-cli\n"},
	{"complexkey", "? origin\n: static\n? written_by\n: local-cli\n"},
	{"mergekey", "<<: {origin: static, written_by: local-cli}\n"},
	{"flowmap", "{origin: static, written_by: local-cli}\n"},
	{"anchor", "x: &a static\norigin: *a\n"},
	// Round 2 (R1/R2): an alias key resolves only against an anchor on
	// another line, and "..." ends the YAML document before the stamp.
	{"aliaskey_docend", "x: &a origin\n*a : static\ny: &b written_by\n*b : local-cli\n...\n"},
	{"aliaskey", "y: &b written_by\n*b : local-cli\n"},
	{"docend_only", "x: y\n...\n"},
	{"docstart_inside", "x: y\n---\norigin: static\n"},
}

func TestBrainWriteAsJWTStampsJWTForEveryForgedSpelling(t *testing.T) {
	for _, v := range forgedVariants {
		t.Run(v.name, func(t *testing.T) {
			srv, dir := jwtServer(t)
			out := resultText(t, authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
				"content": "---\n" + v.fm + "---\n# V\n\nbody", "wing": "testing", "hall": "facts",
				"filename": "forge-" + v.name,
			}))
			assertProvenance(t, srv, dir, out["path"].(string), "jwt", "attacker")
		})
	}
}

func TestBrainUpdateAsJWTStampsJWTForEveryForgedSpelling(t *testing.T) {
	for _, v := range forgedVariants {
		t.Run(v.name, func(t *testing.T) {
			srv, dir := jwtServer(t)
			out := resultText(t, authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
				"content": "# V\n\nv1", "wing": "testing", "hall": "facts", "filename": "upd-" + v.name,
			}))
			path := out["path"].(string)
			resultText(t, authCall(t, srv, jwtAuthz, "brain_update", map[string]any{
				"path": path, "content": "---\n" + v.fm + "---\n# V\n\nv2",
			}))
			assertProvenance(t, srv, dir, path, "jwt", "attacker")
		})
	}
}

// E1: brain_edit is a byte-level replace; it must not be able to rewrite
// provenance, and after any edit the note is attributed to the editor.
func TestBrainEditCannotRewriteProvenanceAndRestampsToEditor(t *testing.T) {
	srv, dir := jwtServer(t)
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
		"content": "# E\n\nbody", "wing": "testing", "hall": "facts", "filename": "edit-forge",
	}))
	path := out["path"].(string)
	authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{
		"path": path, "old_str": "origin: jwt\nwritten_by: attacker\n", "new_str": "origin: static\nwritten_by: local-cli\n",
	})
	assertProvenance(t, srv, dir, path, "jwt", "attacker")

	// Assembling the value from existing text must not work either.
	authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{"path": path, "old_str": "jwt", "new_str": "static"})
	assertProvenance(t, srv, dir, path, "jwt", "attacker")
}

// A static principal's note edited by a JWT session becomes the JWT
// session's: the content a reader sees is no longer the static author's.
func TestBrainEditByJWTReattributesAStaticNote(t *testing.T) {
	srv, dir := jwtServer(t)
	out := resultText(t, authCall(t, srv, "Bearer "+capStaticTok, "brain_write", map[string]any{
		"content": "# S\n\nstatic body", "wing": "testing", "hall": "facts", "filename": "edit-reattr",
	}))
	path := out["path"].(string)
	resultText(t, authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{
		"path": path, "old_str": "static body", "new_str": "injected body",
	}))
	assertProvenance(t, srv, dir, path, "jwt", "attacker")
}

// E1 via round 2's payload: the alias + document-end trick through brain_edit.
func TestBrainEditCannotForgeProvenanceWithAliasKeys(t *testing.T) {
	srv, dir := jwtServer(t)
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
		"content": "---\ntitle: t\n---\n# E\n\nbody", "wing": "testing", "hall": "facts", "filename": "edit-alias",
	}))
	path := out["path"].(string)
	authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{
		"path": path, "old_str": "title: t\n", "new_str": "title: t\nx: &a origin\n*a : static\n...\n",
	})
	assertProvenance(t, srv, dir, path, "jwt", "attacker")
}

// E2 + E4: brain_ingest_raw stamps every page it writes, and a newline in a
// title cannot start a new frontmatter key — observed through brain_query,
// whose reader is line-oriented, not just through yaml.v3 (round 2, R4).
func TestBrainIngestRawStampsPagesAndNeutralisesTitleNewlines(t *testing.T) {
	srv, dir := jwtServer(t)
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_ingest_raw", map[string]any{
		"source": "zqsrc", "pages": []map[string]any{
			{"title": "Zqplain", "type": "concept", "content": "plain body"},
			{"title": "Zqinj\norigin: static\nwritten_by: local-cli\nx: 'y", "type": "concept", "content": "inj body"},
		},
	}))
	pages := out["pages"].([]any)
	require.NotEmpty(t, pages)
	for _, p := range pages {
		assertProvenance(t, srv, dir, p.(string), "jwt", "attacker")
		b, _ := os.ReadFile(filepath.Join(dir, p.(string)))
		var m map[string]any
		require.NoError(t, yaml.Unmarshal([]byte(frontmatterOf(string(b))), &m))
		assert.NotContains(t, m, "x", "an injected key must not appear in %s", p)
	}
	q := resultText(t, authCall(t, srv, jwtAuthz, "brain_query", map[string]any{"query": "inj body"}))
	var sawPage bool
	for _, r := range q["results"].([]any) {
		hit := r.(map[string]any)
		// Generated aggregates (wiki/index.md, log.md, wing _index.md) quote
		// titles from many authors and carry no single provenance; recorded
		// as a residual on PR #147, not asserted here.
		if !strings.HasPrefix(hit["path"].(string), "wiki/concepts/") {
			continue
		}
		sawPage = true
		assert.Equal(t, "jwt", hit["origin"], "line reader origin for %v", hit["path"])
		assert.Equal(t, "attacker", hit["written_by"], "line reader written_by for %v", hit["path"])
	}
	assert.True(t, sawPage, "brain_query must return the injected page itself")
}

// E3: a page type is a directory name; it must not traverse out of wiki/.
func TestBrainIngestRawRejectsPathTraversingType(t *testing.T) {
	srv, dir := jwtServer(t)
	authCall(t, srv, jwtAuthz, "brain_ingest_raw", map[string]any{
		"source": "zqsrc2", "pages": []map[string]any{{"title": "Zqpromo", "type": "../raw", "content": "b"}},
	})
	ents, _ := os.ReadDir(filepath.Join(dir, "raw"))
	assert.Empty(t, ents, "nothing may be written outside wiki/ by a page type")
}

// E3: promotion re-attributes the note to whoever promotes it, regardless of
// what the raw note's frontmatter claims.
func TestBrainPromoteRestampsToPromoter(t *testing.T) {
	srv, dir := jwtServer(t)
	require.NoError(t, os.MkdirAll(filepath.Join(dir, "raw"), 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(dir, "raw", "planted.md"),
		[]byte("---\ntitle: Planted\norigin: static\nwritten_by: local-cli\n---\n# Planted\n\nbody\n"), 0o644))
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_promote", map[string]any{
		"filename": "planted.md", "wing": "testing", "hall": "facts",
	}))
	assertProvenance(t, srv, dir, out["path"].(string), "jwt", "attacker")
}

// E6: an error outcome must not carry argument values into the audit log.
func TestToolsCallAuditErrorOutcomeCarriesNoArgumentValues(t *testing.T) {
	const sentinel = "SENTINEL-ERR-7d1e"
	srv, _ := jwtServer(t)
	var buf bytes.Buffer
	prev := slog.Default()
	slog.SetDefault(slog.New(slog.NewJSONHandler(&buf, nil)))
	t.Cleanup(func() { slog.SetDefault(prev) })
	authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
		"content": "---\nhalf_life: " + sentinel + "\n---\n# H\n\nbody", "wing": "testing", "hall": "facts",
	})
	authCall(t, srv, jwtAuthz, "brain_update", map[string]any{"path": "wiki/testing/facts/" + sentinel + ".md", "content": "x"})
	var sawError bool
	for _, l := range strings.Split(buf.String(), "\n") {
		var rec map[string]any
		if json.Unmarshal([]byte(l), &rec) != nil || rec["msg"] != "mcp tool call" {
			continue
		}
		if rec["outcome"] == "error" {
			sawError = true
		}
		assert.NotContains(t, l, sentinel)
	}
	assert.True(t, sawError, "the failing calls must still be recorded as errors")
}

// Round 3 (B1): brain_edit adding a block scalar containing an indented "---"
// must not hide provenance from brain_query's reader.
func TestBrainEditBlockScalarFenceDoesNotHideProvenanceFromQuery(t *testing.T) {
	srv, _ := jwtServer(t)
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_write", map[string]any{
		"content": "---\nk: v\n---\n# E\n\nzqrthreeeditdash", "wing": "testing", "hall": "facts", "filename": "edit-dash",
	}))
	resultText(t, authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{
		"path": out["path"], "old_str": "k: v\n", "new_str": "k: |\n  ---\n",
	}))
	q := resultText(t, authCall(t, srv, jwtAuthz, "brain_query", map[string]any{"query": "zqrthreeeditdash"}))
	hit := q["results"].([]any)[0].(map[string]any)
	assert.Equal(t, "jwt", hit["origin"])
	assert.Equal(t, "attacker", hit["written_by"])
}

// Round 3 (B2): when the editor's stamp cannot be verified, the edit is rolled
// back byte for byte and the tool says so. A JWT subject with an invalid UTF-8
// byte cannot round-trip through YAML, which forces that path.
func TestBrainEditRollsBackWhenTheStampCannotBeVerified(t *testing.T) {
	dir := newBrainRepo(t)
	cfg, err := classification.Load(dir)
	require.NoError(t, err)
	svc := capture.NewService(brainstore.New(dir), capFakeTracker{}, cfg, audit.NewSlogSink(nil))
	srv := mcp.NewServer(dir, nil, nil, nil)
	srv.WithCapture(svc, capFakeValidator{subject: "bad\xff"}, capStaticTok, "local-cli", capturehttp.NewOriginResolver(nil))

	out := resultText(t, authCall(t, srv, "Bearer "+capStaticTok, "brain_write", map[string]any{
		"content": "# R\n\nstatic body", "wing": "testing", "hall": "facts", "filename": "edit-rollback",
	}))
	path := out["path"].(string)
	before, err := os.ReadFile(filepath.Join(dir, path))
	require.NoError(t, err)

	resp := authCall(t, srv, jwtAuthz, "brain_edit", map[string]any{"path": path, "old_str": "static body", "new_str": "injected"})
	require.NotNil(t, resp["error"], "the edit must be reported as failed")
	assert.Contains(t, resp["error"].(map[string]any)["message"], "rolled back")
	after, err := os.ReadFile(filepath.Join(dir, path))
	require.NoError(t, err)
	assert.Equal(t, string(before), string(after), "the note must be byte-identical to before the edit")
}

// MCP brain_ingest (the LLM path) attributes its pages to the caller.
func TestBrainIngestViaMCPStampsCaller(t *testing.T) {
	dir := newBrainRepo(t)
	cfg, err := classification.Load(dir)
	require.NoError(t, err)
	svc := capture.NewService(brainstore.New(dir), capFakeTracker{}, cfg, audit.NewSlogSink(nil))
	stub := pipeline.Config{Complete: func(context.Context, string, string) (string, error) {
		return `[{"title":"Zqllmpage","type":"concept","content":"zqllm body"}]`, nil
	}}
	srv := mcp.NewServer(dir, &stub, nil, nil)
	srv.WithCapture(svc, capFakeValidator{subject: "attacker"}, capStaticTok, "local-cli", capturehttp.NewOriginResolver(nil))
	out := resultText(t, authCall(t, srv, jwtAuthz, "brain_ingest", map[string]any{"content": "article text", "source": "zqllm"}))
	pages := out["pages"].([]any)
	require.NotEmpty(t, pages)
	for _, p := range pages {
		assertProvenance(t, srv, dir, p.(string), "jwt", "attacker")
	}
}