diff --git a/internal/redaction/betterleaks.go b/internal/redaction/betterleaks.go index 85464be7a..92d02a700 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,149 @@ 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) { + // 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 { + chunks = byteSpans(text, d.chunks) + } else { + chunks = splitLines(text, d.base.chunkSize, chunkOverlapLines) + d.chunks, d.lines = lineSpans(text, chunks), lines + } + + // Hashing the whole document is serial work on every pass otherwise. + digests := make([][sha256.Size]byte, len(chunks)) + var g errgroup.Group + g.SetLimit(runtime.GOMAXPROCS(0)) + for i, c := range chunks { + 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 + } + 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..7de8a4737 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,6 +564,202 @@ 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 { + 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") + } + 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") +} + +// 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) { + doc := sessionDoc(200, func(i int) string { + switch i % 50 { + case 3: + return "git clone https://oauth2:" + fakeGitHubPAT + "@github.com/example/repo.git" + case 17: + return "ANTHROPIC_API_KEY=" + fakeAnthropicKey + case 31: + return "export " + awsPair + } + return "" + }) + + 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 bf6fa8e0c..516355718 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) } @@ -353,10 +368,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,7 +380,11 @@ 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 @@ -383,6 +399,9 @@ func (w *rewriter) redactLeaf(s string) string { if c == 0 { continue } + if n == 0 && !w.pathFilter(path) { + return s + } body = strings.ReplaceAll(body, sr.secret, w.placeholder(sr.ruleID)) n += c w.byRule[sr.ruleID] += c @@ -517,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) diff --git a/internal/redaction/redaction_test.go b/internal/redaction/redaction_test.go index 16e246cb8..ec19ee4be 100644 --- a/internal/redaction/redaction_test.go +++ b/internal/redaction/redaction_test.go @@ -52,6 +52,40 @@ func (f *fakeScanner) Scan(ctx context.Context, text string) ([]Finding, error) return out, nil } +// 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 + bound []*fakeScanner +} + +func (f *fakeDocumentScanner) forDocument() Scanner { + s := &fakeScanner{findings: f.findings, requirePresent: f.requirePresent} + f.bound = append(f.bound, s) + return s +} + +// 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, 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) { testCases := []struct { name string @@ -126,6 +160,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 +304,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"}]}}}`,