From 31856f67a6a0e998cae35be621bbd41916571e48 Mon Sep 17 00:00:00 2001 From: Miguel Martinez Trivino Date: Wed, 30 Sep 2026 14:56:25 +0200 Subject: [PATCH 1/4] perf(redaction): consult the path filter only for leaves with a secret The rewriter called the path filter for every string leaf on every pass. The AI coding session filter splits the leaf path for each protected pattern, so on a large session this was most of the time spent rewriting. The filter now runs only for a leaf that holds a pending secret, and nothing is recorded for a leaf it protects. Refs: #3498 Assisted-by: Claude Code Signed-off-by: Miguel Martinez Trivino Chainloop-Trace-Sessions: ba0ac486-c03f-4584-99da-b10a229310bc --- internal/redaction/redaction.go | 35 +++++++++++++++++----------- internal/redaction/redaction_test.go | 34 +++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 14 deletions(-) diff --git a/internal/redaction/redaction.go b/internal/redaction/redaction.go index bf6fa8e0c..90a4414d7 100644 --- a/internal/redaction/redaction.go +++ b/internal/redaction/redaction.go @@ -353,10 +353,7 @@ func (w *rewriter) rewrite(node any, path string) any { } return v case string: - if !w.pathFilter(path) { - return v - } - return w.redactLeaf(v) + return w.redactLeaf(v, path) default: // Numbers, booleans and null cannot carry a secret the scanner reported // as a string. @@ -368,30 +365,40 @@ func (w *rewriter) rewrite(node any, path string) any { // against the leaf's JSON-encoded form, because that is the text the scanner // saw: a secret containing a newline, for instance, reaches us as the two // characters `\n`. -func (w *rewriter) redactLeaf(s string) string { +// +// The path filter is consulted only once the leaf is known to hold a secret, and +// nothing is recorded for a leaf it protects. Nearly every leaf holds none, and +// the filter runs for each leaf on every pass otherwise. +func (w *rewriter) redactLeaf(s, path string) string { body, err := encodeStringBody(s) if err != nil { return s } - var ( - n int - lastRule string - ) + type hit struct { + sr secretRule + count int + } + var hits []hit for _, sr := range w.secrets { c := strings.Count(body, sr.secret) if c == 0 { continue } body = strings.ReplaceAll(body, sr.secret, w.placeholder(sr.ruleID)) - n += c - w.byRule[sr.ruleID] += c - w.located[sr.secret] = struct{}{} - lastRule = sr.ruleID + hits = append(hits, hit{sr: sr, count: c}) } - if n == 0 { + if len(hits) == 0 || !w.pathFilter(path) { return s } + + var n int + for _, h := range hits { + n += h.count + w.byRule[h.sr.ruleID] += h.count + w.located[h.sr.secret] = struct{}{} + } + lastRule := hits[len(hits)-1].sr.ruleID w.count += n var out string diff --git a/internal/redaction/redaction_test.go b/internal/redaction/redaction_test.go index 16e246cb8..a19f8d2ac 100644 --- a/internal/redaction/redaction_test.go +++ b/internal/redaction/redaction_test.go @@ -126,6 +126,20 @@ func TestRedact(t *testing.T) { wantUnchanged: true, wantUnlocated: map[string]int{"r1": 1}, }, + { + name: "a secret in a protected and an eligible leaf is replaced only where eligible", + doc: `{"keepme":"SEC","other":"x SEC"}`, + findings: []Finding{{RuleID: "r1", Secret: "SEC"}}, + opts: []Option{WithPathFilter(func(p string) bool { + return p != "/keepme" + })}, + wantReplacements: 1, + wantByRule: map[string]int{"r1": 1}, + // The copy left in the protected leaf is still found on the next + // pass, where no eligible leaf holds it any more. + wantUnlocated: map[string]int{"r1": 1}, + mustContain: []string{`"keepme":"SEC"`, `"other":"x [REDACTED:r1]"`}, + }, { name: "finding present nowhere is classified as an artifact", doc: `{"a":"plain"}`, @@ -256,6 +270,26 @@ func TestRedact(t *testing.T) { } } +// TestRedactFiltersOnlyMatchingLeaves pins the cost of the path filter. It is +// consulted for the leaves that hold a pending secret, not for every leaf on +// every pass: a transcript has millions of leaves and only a handful of secrets, +// and the AI coding session filter is not cheap. +func TestRedactFiltersOnlyMatchingLeaves(t *testing.T) { + var consulted []string + filter := func(p string) bool { + consulted = append(consulted, p) + return true + } + + doc := `{"a":"plain","b":{"c":"x SEC x","d":["plain","plain"]},"e":"plain"}` + scanner := &fakeScanner{findings: []Finding{{RuleID: "r1", Secret: "SEC"}}, requirePresent: true} + + _, report, err := New(scanner, WithPathFilter(filter)).Redact(context.Background(), []byte(doc)) + require.NoError(t, err) + require.Equal(t, 1, report.Replacements) + assert.Equal(t, []string{"/b/c"}, consulted) +} + func TestRedactIsIdempotent(t *testing.T) { docs := []string{ `{"data":{"raw_session":{"main":[{"content":"FAKE-AWS-KEY-NOT-A-REAL-PATTERN"}]}}}`, From 665212f8ecff7989f60173e6b830321a31ffb425 Mon Sep 17 00:00:00 2001 From: Miguel Martinez Trivino Date: Wed, 30 Sep 2026 15:01:03 +0200 Subject: [PATCH 2/4] perf(redaction): rescan only the chunks that changed between passes Every redaction pass scanned the whole document again, although a pass changes only the leaves that held a secret. Redaction replaces text inside string leaves and never adds or removes a line of the rendered document. So the chunks of the first pass are kept as ranges of lines, and a later pass scans only the chunks whose text changed. The findings of the other chunks are reused, keyed by the SHA-256 digest of the chunk text. A version with a different number of lines is cut afresh. The Redactor binds a scanner to each document through an optional interface, so a scanner that remembers earlier passes is never shared between documents. Refs: #3498 Assisted-by: Claude Code Signed-off-by: Miguel Martinez Trivino Chainloop-Trace-Sessions: ba0ac486-c03f-4584-99da-b10a229310bc --- internal/redaction/betterleaks.go | 137 +++++++++++++++++++- internal/redaction/betterleaks_test.go | 165 +++++++++++++++++++++++++ internal/redaction/redaction.go | 17 ++- internal/redaction/redaction_test.go | 31 +++++ 4 files changed, 348 insertions(+), 2 deletions(-) diff --git a/internal/redaction/betterleaks.go b/internal/redaction/betterleaks.go index 85464be7a..c27ec0f5c 100644 --- a/internal/redaction/betterleaks.go +++ b/internal/redaction/betterleaks.go @@ -17,6 +17,7 @@ package redaction import ( "context" + "crypto/sha256" "fmt" "runtime" "slices" @@ -124,6 +125,16 @@ func newBetterleaksScanner() (*betterleaksScanner, error) { } func (s *betterleaksScanner) Scan(ctx context.Context, text string) ([]Finding, error) { + perChunk, err := s.scanChunks(ctx, text, splitLines(text, s.chunkSize, chunkOverlapLines)) + if err != nil { + return nil, err + } + return slices.Concat(perChunk...), nil +} + +// scanChunks scans the given chunks of text and returns the findings of each +// one, in the order of chunks. +func (s *betterleaksScanner) scanChunks(ctx context.Context, text string, chunks []span) ([][]Finding, error) { if err := ctx.Err(); err != nil { return nil, err } @@ -131,7 +142,7 @@ func (s *betterleaksScanner) Scan(ctx context.Context, text string) ([]Finding, s.mu.Lock() defer s.mu.Unlock() - source := chunkedSource{text: text, chunks: splitLines(text, s.chunkSize, chunkOverlapLines)} + source := chunkedSource{text: text, chunks: chunks} // Chunks are scanned concurrently, so findings arrive interleaved. They are // collected per chunk and joined in chunk order, which keeps the result @@ -175,9 +186,133 @@ func (s *betterleaksScanner) Scan(ctx context.Context, text string) ([]Finding, return nil, err } + return perChunk, nil +} + +// forDocument returns a scanner for the passes over a single document. +func (s *betterleaksScanner) forDocument() Scanner { + return &documentScanner{base: s} +} + +// documentScanner scans successive versions of one document, and scans again +// only the chunks whose text changed since the previous version. +// +// Redaction replaces text inside string leaves, which never adds or removes a +// line of the rendered document. So after the first pass the chunks are fixed as +// ranges of lines rather than of bytes, and a chunk far from any replacement has +// exactly the same text as before. The detector is deterministic, so its +// findings for that text are the ones found last time. +type documentScanner struct { + base *betterleaksScanner + // chunks are the line ranges the document was cut into on its first scan. + chunks []lineSpan + // lines is the number of lines the document had then. A version with a + // different count is cut afresh. + lines int + // findings holds the findings of each chunk of the previous version, by the + // digest of its text. + findings map[[sha256.Size]byte][]Finding +} + +// lineSpan is a half-open range [start, end) of line indices. +type lineSpan struct { + start, end int +} + +func (d *documentScanner) Scan(ctx context.Context, text string) ([]Finding, error) { + lines := strings.Count(text, "\n") + 1 + var chunks []span + if d.chunks != nil && lines == d.lines { + chunks = byteSpans(text, d.chunks) + } else { + chunks = splitLines(text, d.base.chunkSize, chunkOverlapLines) + d.chunks, d.lines = lineSpans(text, chunks), lines + } + + perChunk := make([][]Finding, len(chunks)) + digests := make([][sha256.Size]byte, len(chunks)) + var stale []int + for i, c := range chunks { + digests[i] = sha256.Sum256([]byte(text[c.start:c.end])) + if found, ok := d.findings[digests[i]]; ok { + perChunk[i] = found + continue + } + stale = append(stale, i) + } + + if len(stale) > 0 { + toScan := make([]span, len(stale)) + for j, i := range stale { + toScan[j] = chunks[i] + } + scanned, err := d.base.scanChunks(ctx, text, toScan) + if err != nil { + return nil, err + } + for j, i := range stale { + perChunk[i] = scanned[j] + } + } + + // Only the previous version is ever compared against, so older findings + // are dropped rather than kept for the life of the document. + d.findings = make(map[[sha256.Size]byte][]Finding, len(chunks)) + for i := range chunks { + d.findings[digests[i]] = perChunk[i] + } + return slices.Concat(perChunk...), nil } +// lineSpans converts byte spans of text, each starting at a line start, into +// line ranges. +func lineSpans(text string, spans []span) []lineSpan { + out := make([]lineSpan, len(spans)) + // Spans are ordered and may overlap, so the line count is carried forward + // from the previous span's start rather than recounted from the beginning. + offset, line := 0, 0 + for i, s := range spans { + line += strings.Count(text[offset:s.start], "\n") + offset = s.start + end := line + strings.Count(text[s.start:s.end], "\n") + if s.end == len(text) && !strings.HasSuffix(text, "\n") { + // The final line has no newline to count. + end++ + } + out[i] = lineSpan{start: line, end: end} + } + return out +} + +// byteSpans converts line ranges back into byte spans of text. A range whose +// last line has no newline runs to the end of the text. +func byteSpans(text string, ranges []lineSpan) []span { + out := make([]span, len(ranges)) + offset, line := 0, 0 + // advance moves offset to the start of the given line. + advance := func(to int) { + for ; line < to; line++ { + nl := strings.IndexByte(text[offset:], '\n') + if nl < 0 { + offset = len(text) + line = to + return + } + offset += nl + 1 + } + } + for i, r := range ranges { + advance(r.start) + start := offset + advance(r.end) + out[i] = span{start: start, end: offset} + // The next range starts at or before this end, so rewind. + line, offset = r.start, start + } + return out +} + // appendSecret records a locatable secret. A finding without one cannot be // searched for in the document, so it is dropped rather than reported. func appendSecret(dst []Finding, ruleID, secret string) []Finding { diff --git a/internal/redaction/betterleaks_test.go b/internal/redaction/betterleaks_test.go index 418347440..bd30410c8 100644 --- a/internal/redaction/betterleaks_test.go +++ b/internal/redaction/betterleaks_test.go @@ -572,6 +572,171 @@ func TestChunkedRedact(t *testing.T) { assert.Equal(t, 1, report.ByRule["anthropic-api-key"]) } +// multiChunkText is indented-JSON-shaped text spanning many 256-byte chunks, with +// a distinct token on every tenth line. +func multiChunkText() string { + var sb strings.Builder + for i := range 300 { + if i%10 == 0 { + pat := fakeGitHubPAT[:len(fakeGitHubPAT)-2] + string(rune('A'+i/10%26)) + "x" + sb.WriteString(` "repository": "https://oauth2:` + pat + `@github.com/example/repo.git",` + "\n") + continue + } + sb.WriteString(` "content": "ordinary transcript line ` + strconv.Itoa(i) + `",` + "\n") + } + return sb.String() +} + +// uniqueFindings is the set of findings, which is what must not change however +// the text is chunked. +func uniqueFindings(findings []Finding) []Finding { + seen := make(map[Finding]struct{}, len(findings)) + var out []Finding + for _, f := range findings { + if _, dup := seen[f]; !dup { + seen[f] = struct{}{} + out = append(out, f) + } + } + return out +} + +// TestLineSpansRoundTrip checks that chunks converted to line ranges and back +// are the same byte ranges, which is what lets a later pass reuse them. +func TestLineSpansRoundTrip(t *testing.T) { + testCases := []struct { + name string + text string + size int + overlap int + }{ + {name: "empty", text: "", size: 4}, + {name: "single chunk", text: "a\nb\n", size: 16}, + {name: "trailing newline", text: multiChunkText(), size: 256, overlap: chunkOverlapLines}, + {name: "no trailing newline", text: strings.TrimSuffix(multiChunkText(), "\n"), size: 256, overlap: chunkOverlapLines}, + {name: "no overlap", text: multiChunkText(), size: 256}, + {name: "long lines", text: "aaaaaaaa\nbbbbbbbb\nc", size: 4, overlap: 2}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + spans := splitLines(tc.text, tc.size, tc.overlap) + assert.Equal(t, spans, byteSpans(tc.text, lineSpans(tc.text, spans))) + }) + } +} + +// TestDocumentScannerMatchesFreshScan is the safety property of rescanning +// incrementally: after the document changes, the scanner bound to it reports the +// same secrets as a scan from scratch would. +func TestDocumentScannerMatchesFreshScan(t *testing.T) { + replaceFirst := func(text string) string { + i := strings.Index(text, "ghp_") + return text[:i] + "[REDACTED:github-pat]" + text[i+len(fakeGitHubPAT):] + } + + testCases := []struct { + name string + edit func(string) string + }{ + {name: "unchanged", edit: func(s string) string { return s }}, + {name: "one secret replaced in place", edit: replaceFirst}, + {name: "every secret replaced in place", edit: func(s string) string { + for strings.Contains(s, "ghp_") { + s = replaceFirst(s) + } + return s + }}, + {name: "a line inserted", edit: func(s string) string { + return ` "content": "ANTHROPIC_API_KEY=` + fakeAnthropicKey + `",` + "\n" + s + }}, + {name: "a line removed", edit: func(s string) string { + return s[strings.IndexByte(s, '\n')+1:] + }}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + scanner := chunkedScanner(t, 256) + doc := scanner.forDocument() + text := multiChunkText() + + first, err := doc.Scan(context.Background(), text) + require.NoError(t, err) + fresh, err := scanner.Scan(context.Background(), text) + require.NoError(t, err) + require.Equal(t, fresh, first, "the first scan of a document is a scan from scratch") + + edited := tc.edit(text) + got, err := doc.Scan(context.Background(), edited) + require.NoError(t, err) + want, err := scanner.Scan(context.Background(), edited) + require.NoError(t, err) + assert.ElementsMatch(t, uniqueFindings(want), uniqueFindings(got)) + }) + } +} + +// TestDocumentScannerRescansOnlyChangedChunks pins the point of scanning +// incrementally: a pass after a few in-place replacements scans the chunks that +// changed, not the whole document again. +func TestDocumentScannerRescansOnlyChangedChunks(t *testing.T) { + scanner := chunkedScanner(t, 256) + doc := scanner.forDocument() + text := multiChunkText() + + _, err := doc.Scan(context.Background(), text) + require.NoError(t, err) + + i := strings.Index(text, "ghp_") + edited := text[:i] + "[REDACTED:github-pat]" + text[i+len(fakeGitHubPAT):] + + before := scanner.detector.TotalBytes.Load() + _, err = doc.Scan(context.Background(), edited) + require.NoError(t, err) + scanned := scanner.detector.TotalBytes.Load() - before + + assert.Positive(t, scanned, "the changed chunk has to be scanned again") + assert.Less(t, scanned, uint64(len(edited)/4), "only the chunks around the change are scanned again") +} + +// TestRedactIncrementalMatchesFullRescan runs whole redactions with and without +// incremental rescanning and requires the same result. +func TestRedactIncrementalMatchesFullRescan(t *testing.T) { + var sb strings.Builder + sb.WriteString(`{"data":{"raw_session":{"main":[`) + for i := range 200 { + if i > 0 { + sb.WriteString(",") + } + content := "ordinary transcript text " + strconv.Itoa(i) + switch i % 50 { + case 3: + content = "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" + case 17: + content = "ANTHROPIC_API_KEY=" + fakeAnthropicKey + case 31: + content = "export " + awsPair + } + sb.WriteString(`{"role":"user","content":"` + content + `"}`) + } + sb.WriteString(`]}}}`) + doc := []byte(sb.String()) + + scanner := chunkedScanner(t, 256) + // Embedding hides forDocument, so this one scans every pass from scratch. + full := struct{ Scanner }{scanner} + + want, wantReport, err := New(full).Redact(context.Background(), doc) + require.NoError(t, err) + got, gotReport, err := New(scanner).Redact(context.Background(), doc) + require.NoError(t, err) + + assert.Equal(t, string(want), string(got)) + assert.Equal(t, wantReport, gotReport) + assert.True(t, gotReport.Changed()) +} + func BenchmarkDefaultScannerInit(b *testing.B) { for b.Loop() { if _, err := newBetterleaksScanner(); err != nil { diff --git a/internal/redaction/redaction.go b/internal/redaction/redaction.go index 90a4414d7..433e3723b 100644 --- a/internal/redaction/redaction.go +++ b/internal/redaction/redaction.go @@ -75,6 +75,14 @@ type Scanner interface { Scan(ctx context.Context, text string) ([]Finding, error) } +// documentBinder is implemented by scanners that scan a document faster when +// they know that successive texts are passes over the same document. +type documentBinder interface { + // forDocument returns a scanner for the passes over a single document. It + // may remember earlier passes, so it is never shared between documents. + forDocument() Scanner +} + // PathFilter reports whether the string leaf at the given path may be rewritten. // Paths look like "/data/raw_session/main/0/content": a leading slash, object // keys and array indices separated by slashes. @@ -216,6 +224,13 @@ func (r *Redactor) Redact(ctx context.Context, doc []byte) ([]byte, *Report, err skip := make(map[string]struct{}) converged := false + // Every pass scans a new version of the same document, which a scanner that + // knows so can scan faster. + scanner := r.scanner + if b, ok := scanner.(documentBinder); ok { + scanner = b.forDocument() + } + for pass := 1; pass <= r.maxPasses; pass++ { report.Passes = pass @@ -227,7 +242,7 @@ func (r *Redactor) Redact(ctx context.Context, doc []byte) ([]byte, *Report, err return nil, nil, fmt.Errorf("rendering document: %w", err) } - findings, err := r.scanner.Scan(ctx, text) + findings, err := scanner.Scan(ctx, text) if err != nil { return nil, nil, fmt.Errorf("scanning for secrets: %w", err) } diff --git a/internal/redaction/redaction_test.go b/internal/redaction/redaction_test.go index a19f8d2ac..65b3e069b 100644 --- a/internal/redaction/redaction_test.go +++ b/internal/redaction/redaction_test.go @@ -52,6 +52,37 @@ func (f *fakeScanner) Scan(ctx context.Context, text string) ([]Finding, error) return out, nil } +// fakeDocumentScanner is a fakeScanner that also binds itself to a document, to +// check how Redact uses a scanner that can rescan incrementally. +type fakeDocumentScanner struct { + fakeScanner + documents int +} + +func (f *fakeDocumentScanner) forDocument() Scanner { + f.documents++ + return &f.fakeScanner +} + +// TestRedactBindsOneScannerPerDocument checks that every pass over a document +// goes through the scanner bound to it, and that each Redact binds a new one: +// what that scanner remembers belongs to a single document. +func TestRedactBindsOneScannerPerDocument(t *testing.T) { + scanner := &fakeDocumentScanner{fakeScanner: fakeScanner{ + findings: []Finding{{RuleID: "r1", Secret: "SEC"}}, + requirePresent: true, + }} + r := New(scanner) + + for i := 1; i <= 2; i++ { + _, report, err := r.Redact(context.Background(), []byte(`{"a":"x SEC x"}`)) + require.NoError(t, err) + assert.Equal(t, i, scanner.documents) + assert.Equal(t, 2*i, scanner.calls, "both passes use the bound scanner") + assert.Equal(t, 2, report.Passes) + } +} + func TestRedact(t *testing.T) { testCases := []struct { name string From 70d8a14e6d64a51c9bafe50c6e988378052971e6 Mon Sep 17 00:00:00 2001 From: Miguel Martinez Trivino Date: Wed, 30 Sep 2026 15:08:07 +0200 Subject: [PATCH 3/4] refactor(redaction): simplify the rewriter and hash chunks in parallel - Consult the path filter at the first match in the existing loop, instead of collecting matches and committing them in a second loop. - Compute the chunk digests concurrently. Hashing the whole document was serial work on every pass. - State at encode that rewriting a leaf keeps the line count, which the incremental rescan relies on. - Share the test builders for sessions and distinct tokens. The old multi-chunk builder repeated some of its tokens. Refs: #3498 Assisted-by: Claude Code Signed-off-by: Miguel Martinez Trivino Chainloop-Trace-Sessions: ba0ac486-c03f-4584-99da-b10a229310bc --- internal/redaction/betterleaks.go | 16 ++++-- internal/redaction/betterleaks_test.go | 70 ++++++++++++++------------ internal/redaction/redaction.go | 31 ++++++------ 3 files changed, 68 insertions(+), 49 deletions(-) diff --git a/internal/redaction/betterleaks.go b/internal/redaction/betterleaks.go index c27ec0f5c..e0982a71c 100644 --- a/internal/redaction/betterleaks.go +++ b/internal/redaction/betterleaks.go @@ -229,11 +229,21 @@ func (d *documentScanner) Scan(ctx context.Context, text string) ([]Finding, err d.chunks, d.lines = lineSpans(text, chunks), lines } - perChunk := make([][]Finding, len(chunks)) + // Hashing the whole document is serial work on every pass otherwise. digests := make([][sha256.Size]byte, len(chunks)) - var stale []int + var g errgroup.Group + g.SetLimit(runtime.GOMAXPROCS(0)) for i, c := range chunks { - digests[i] = sha256.Sum256([]byte(text[c.start:c.end])) + g.Go(func() error { + digests[i] = sha256.Sum256([]byte(text[c.start:c.end])) + return nil + }) + } + _ = g.Wait() + + perChunk := make([][]Finding, len(chunks)) + var stale []int + for i := range chunks { if found, ok := d.findings[digests[i]]; ok { perChunk[i] = found continue diff --git a/internal/redaction/betterleaks_test.go b/internal/redaction/betterleaks_test.go index bd30410c8..d12c6720d 100644 --- a/internal/redaction/betterleaks_test.go +++ b/internal/redaction/betterleaks_test.go @@ -525,8 +525,7 @@ func TestChunkedScanIsDeterministic(t *testing.T) { var sb strings.Builder const tokens = 60 for i := range tokens { - pat := fakeGitHubPAT[:len(fakeGitHubPAT)-2] + string(rune('A'+i%26)) + string(rune('a'+i/26)) - sb.WriteString(` "repository": "https://oauth2:` + pat + `@github.com/example/repo.git",` + "\n") + sb.WriteString(` "repository": "https://oauth2:` + distinctPAT(i) + `@github.com/example/repo.git",` + "\n") sb.WriteString(` "content": "` + strings.Repeat("ordinary transcript text ", 10) + `",` + "\n") } text := sb.String() @@ -547,24 +546,17 @@ func TestChunkedScanIsDeterministic(t *testing.T) { func TestChunkedRedact(t *testing.T) { scanner := chunkedScanner(t, 256) - var sb strings.Builder - sb.WriteString(`{"data":{"raw_session":{"main":[`) - for i := range 200 { - if i > 0 { - sb.WriteString(",") - } - content := "ordinary transcript text " + strconv.Itoa(i) + doc := sessionDoc(200, func(i int) string { switch i { case 3: - content = "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" + return "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" case 197: - content = "ANTHROPIC_API_KEY=" + fakeAnthropicKey + return "ANTHROPIC_API_KEY=" + fakeAnthropicKey } - sb.WriteString(`{"role":"user","content":"` + content + `"}`) - } - sb.WriteString(`]}}}`) + return "" + }) - out, report, err := New(scanner).Redact(context.Background(), []byte(sb.String())) + out, report, err := New(scanner).Redact(context.Background(), doc) require.NoError(t, err) assert.NotContains(t, string(out), fakeGitHubPAT) assert.NotContains(t, string(out), fakeAnthropicKey) @@ -572,14 +564,38 @@ func TestChunkedRedact(t *testing.T) { assert.Equal(t, 1, report.ByRule["anthropic-api-key"]) } +// distinctPAT returns a GitHub token of the same shape as fakeGitHubPAT, and a +// different one for every i below 26*26. +func distinctPAT(i int) string { + return fakeGitHubPAT[:len(fakeGitHubPAT)-2] + string(rune('A'+i%26)) + string(rune('a'+i/26)) +} + +// sessionDoc builds an AI coding session document with the given number of +// turns. content returns the text of turn i, or "" for ordinary text. +func sessionDoc(turns int, content func(i int) string) []byte { + var sb strings.Builder + sb.WriteString(`{"data":{"raw_session":{"main":[`) + for i := range turns { + if i > 0 { + sb.WriteString(",") + } + c := content(i) + if c == "" { + c = "ordinary transcript text " + strconv.Itoa(i) + } + sb.WriteString(`{"role":"user","content":"` + c + `"}`) + } + sb.WriteString(`]}}}`) + return []byte(sb.String()) +} + // multiChunkText is indented-JSON-shaped text spanning many 256-byte chunks, with // a distinct token on every tenth line. func multiChunkText() string { var sb strings.Builder for i := range 300 { if i%10 == 0 { - pat := fakeGitHubPAT[:len(fakeGitHubPAT)-2] + string(rune('A'+i/10%26)) + "x" - sb.WriteString(` "repository": "https://oauth2:` + pat + `@github.com/example/repo.git",` + "\n") + sb.WriteString(` "repository": "https://oauth2:` + distinctPAT(i/10) + `@github.com/example/repo.git",` + "\n") continue } sb.WriteString(` "content": "ordinary transcript line ` + strconv.Itoa(i) + `",` + "\n") @@ -703,25 +719,17 @@ func TestDocumentScannerRescansOnlyChangedChunks(t *testing.T) { // TestRedactIncrementalMatchesFullRescan runs whole redactions with and without // incremental rescanning and requires the same result. func TestRedactIncrementalMatchesFullRescan(t *testing.T) { - var sb strings.Builder - sb.WriteString(`{"data":{"raw_session":{"main":[`) - for i := range 200 { - if i > 0 { - sb.WriteString(",") - } - content := "ordinary transcript text " + strconv.Itoa(i) + doc := sessionDoc(200, func(i int) string { switch i % 50 { case 3: - content = "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" + return "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" case 17: - content = "ANTHROPIC_API_KEY=" + fakeAnthropicKey + return "ANTHROPIC_API_KEY=" + fakeAnthropicKey case 31: - content = "export " + awsPair + return "export " + awsPair } - sb.WriteString(`{"role":"user","content":"` + content + `"}`) - } - sb.WriteString(`]}}}`) - doc := []byte(sb.String()) + return "" + }) scanner := chunkedScanner(t, 256) // Embedding hides forDocument, so this one scans every pass from scratch. diff --git a/internal/redaction/redaction.go b/internal/redaction/redaction.go index 433e3723b..516355718 100644 --- a/internal/redaction/redaction.go +++ b/internal/redaction/redaction.go @@ -390,30 +390,27 @@ func (w *rewriter) redactLeaf(s, path string) string { return s } - type hit struct { - sr secretRule - count int - } - var hits []hit + var ( + n int + lastRule string + ) for _, sr := range w.secrets { c := strings.Count(body, sr.secret) if c == 0 { continue } + if n == 0 && !w.pathFilter(path) { + return s + } body = strings.ReplaceAll(body, sr.secret, w.placeholder(sr.ruleID)) - hits = append(hits, hit{sr: sr, count: c}) + n += c + w.byRule[sr.ruleID] += c + w.located[sr.secret] = struct{}{} + lastRule = sr.ruleID } - if len(hits) == 0 || !w.pathFilter(path) { + if n == 0 { return s } - - var n int - for _, h := range hits { - n += h.count - w.byRule[h.sr.ruleID] += h.count - w.located[h.sr.secret] = struct{}{} - } - lastRule := hits[len(hits)-1].sr.ruleID w.count += n var out string @@ -539,6 +536,10 @@ func rejectDuplicateKeys(doc []byte) error { // encode serialises v. HTML escaping is disabled so transcript text keeps its // angle brackets and ampersands instead of being mangled into \u sequences. +// +// Newlines inside strings are escaped, so rewriting a leaf never changes the +// number of lines in the indented rendering. The betterleaks scanner relies on +// that to rescan only the chunks that changed between passes. func encode(v any, indent bool) (string, error) { var buf bytes.Buffer enc := json.NewEncoder(&buf) From 94cb638c8e5b4c9bb46026c4c4776d750ce461dd Mon Sep 17 00:00:00 2001 From: Miguel Martinez Trivino Date: Wed, 30 Sep 2026 15:17:40 +0200 Subject: [PATCH 4/4] fix(redaction): honour cancellation when every chunk is unchanged A pass whose chunks were all unchanged returned the findings of the previous pass without reaching scanChunks, which checks the context. A cancelled redaction could therefore report success. The document scanner now checks the context itself. The document binding test now gives each document its own fake scanner, so it fails if Redact ignores the bound scanner or shares one between documents. Refs: #3498 Assisted-by: Claude Code Signed-off-by: Miguel Martinez Trivino Chainloop-Trace-Sessions: ba0ac486-c03f-4584-99da-b10a229310bc --- internal/redaction/betterleaks.go | 6 ++++++ internal/redaction/betterleaks_test.go | 15 +++++++++++++++ internal/redaction/redaction_test.go | 17 ++++++++++------- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/internal/redaction/betterleaks.go b/internal/redaction/betterleaks.go index e0982a71c..92d02a700 100644 --- a/internal/redaction/betterleaks.go +++ b/internal/redaction/betterleaks.go @@ -220,6 +220,12 @@ type lineSpan struct { } func (d *documentScanner) Scan(ctx context.Context, text string) ([]Finding, error) { + // Checked here as well as in scanChunks, which a pass whose chunks are all + // unchanged never reaches. + if err := ctx.Err(); err != nil { + return nil, err + } + lines := strings.Count(text, "\n") + 1 var chunks []span if d.chunks != nil && lines == d.lines { diff --git a/internal/redaction/betterleaks_test.go b/internal/redaction/betterleaks_test.go index d12c6720d..7de8a4737 100644 --- a/internal/redaction/betterleaks_test.go +++ b/internal/redaction/betterleaks_test.go @@ -716,6 +716,21 @@ func TestDocumentScannerRescansOnlyChangedChunks(t *testing.T) { assert.Less(t, scanned, uint64(len(edited)/4), "only the chunks around the change are scanned again") } +// TestDocumentScannerCancelledContext checks that cancellation is honoured +// even when every chunk could be answered from what the previous pass found. +func TestDocumentScannerCancelledContext(t *testing.T) { + doc := chunkedScanner(t, 256).forDocument() + text := multiChunkText() + + _, err := doc.Scan(context.Background(), text) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = doc.Scan(ctx, text) + require.ErrorIs(t, err, context.Canceled) +} + // TestRedactIncrementalMatchesFullRescan runs whole redactions with and without // incremental rescanning and requires the same result. func TestRedactIncrementalMatchesFullRescan(t *testing.T) { diff --git a/internal/redaction/redaction_test.go b/internal/redaction/redaction_test.go index 65b3e069b..ec19ee4be 100644 --- a/internal/redaction/redaction_test.go +++ b/internal/redaction/redaction_test.go @@ -52,16 +52,18 @@ func (f *fakeScanner) Scan(ctx context.Context, text string) ([]Finding, error) return out, nil } -// fakeDocumentScanner is a fakeScanner that also binds itself to a document, to -// check how Redact uses a scanner that can rescan incrementally. +// fakeDocumentScanner binds a new fakeScanner to every document, to check how +// Redact uses a scanner that can rescan incrementally. Its own Scan must never +// be called once Redact has bound a document scanner. type fakeDocumentScanner struct { fakeScanner - documents int + bound []*fakeScanner } func (f *fakeDocumentScanner) forDocument() Scanner { - f.documents++ - return &f.fakeScanner + s := &fakeScanner{findings: f.findings, requirePresent: f.requirePresent} + f.bound = append(f.bound, s) + return s } // TestRedactBindsOneScannerPerDocument checks that every pass over a document @@ -77,10 +79,11 @@ func TestRedactBindsOneScannerPerDocument(t *testing.T) { for i := 1; i <= 2; i++ { _, report, err := r.Redact(context.Background(), []byte(`{"a":"x SEC x"}`)) require.NoError(t, err) - assert.Equal(t, i, scanner.documents) - assert.Equal(t, 2*i, scanner.calls, "both passes use the bound scanner") assert.Equal(t, 2, report.Passes) + require.Len(t, scanner.bound, i, "each document gets its own scanner") + assert.Equal(t, 2, scanner.bound[i-1].calls, "both passes use the scanner bound to the document") } + assert.Zero(t, scanner.calls, "the unbound scanner is never used") } func TestRedact(t *testing.T) {