From c83a2cebfce9f276ba9eb3d05171d93d80f8fa67 Mon Sep 17 00:00:00 2001
From: alex <53851759+alxxjohn@users.noreply.github.com>
Date: Tue, 18 Aug 2026 18:50:09 -0400
Subject: [PATCH] fix(parser): bound tree-sitter scan resources

---
 internal/codeguard/runner/support/corpus.go   | 36 +++++++++++++-----
 .../runner/support/corpus_script_test.go      | 38 +++++++++++++++++++
 2 files changed, 65 insertions(+), 9 deletions(-)
 create mode 100644 internal/codeguard/runner/support/corpus_script_test.go

diff --git a/internal/codeguard/runner/support/corpus.go b/internal/codeguard/runner/support/corpus.go
index fca52012..404d9830 100644
--- a/internal/codeguard/runner/support/corpus.go
+++ b/internal/codeguard/runner/support/corpus.go
@@ -43,13 +43,22 @@ func readCappedFile(path string) ([]byte, error) {
 // Each cached slot carries its own sync.Once, so concurrent callers racing on a
 // cold slot compute it exactly once and every caller observes the same result.
 type fileCorpus struct {
-	mu      sync.Mutex
-	targets map[string]*targetListing
-	reads   map[string]*fileRead
-	asts    map[string]*goParse
-	scripts map[string]*scriptParse
+	mu          sync.Mutex
+	targets     map[string]*targetListing
+	reads       map[string]*fileRead
+	asts        map[string]*goParse
+	scripts     map[string]*scriptParse
+	scriptBytes int
+	scriptCount int
+	scriptParse chan struct{}
 }
 
+// maxTreeSitterScanBytes bounds the source represented by retained script
+// trees during one scan. Parsing is also serialized because the pure-Go
+// runtime's transient heap is much larger than its input.
+const maxTreeSitterScanBytes = 256 * 1024
+const maxTreeSitterScanFiles = 64
+
 type targetListing struct {
 	once  sync.Once
 	files []string
@@ -77,10 +86,11 @@ type scriptParse struct {
 
 func newFileCorpus() *fileCorpus {
 	return &fileCorpus{
-		targets: map[string]*targetListing{},
-		reads:   map[string]*fileRead{},
-		asts:    map[string]*goParse{},
-		scripts: map[string]*scriptParse{},
+		targets:     map[string]*targetListing{},
+		reads:       map[string]*fileRead{},
+		asts:        map[string]*goParse{},
+		scripts:     map[string]*scriptParse{},
+		scriptParse: make(chan struct{}, 1),
 	}
 }
 
@@ -153,12 +163,20 @@ func (c *fileCorpus) parseScript(path string, data []byte, lang checkSupport.Scr
 	c.mu.Lock()
 	entry, ok := c.scripts[key]
 	if !ok {
+		if c.scriptCount >= maxTreeSitterScanFiles || c.scriptBytes+len(data) > maxTreeSitterScanBytes {
+			c.mu.Unlock()
+			return nil, fmt.Errorf("tree-sitter scan budget of %d files or %d bytes exhausted", maxTreeSitterScanFiles, maxTreeSitterScanBytes)
+		}
 		entry = &scriptParse{}
 		c.scripts[key] = entry
+		c.scriptCount++
+		c.scriptBytes += len(data)
 	}
 	c.mu.Unlock()
 
 	entry.once.Do(func() {
+		c.scriptParse <- struct{}{}
+		defer func() { <-c.scriptParse }()
 		entry.tree, entry.err = checkSupport.ParseScriptSource(path, data, lang)
 	})
 	return entry.tree, entry.err
diff --git a/internal/codeguard/runner/support/corpus_script_test.go b/internal/codeguard/runner/support/corpus_script_test.go
new file mode 100644
index 00000000..e1f2b93a
--- /dev/null
+++ b/internal/codeguard/runner/support/corpus_script_test.go
@@ -0,0 +1,38 @@
+package support
+
+import (
+	"strings"
+	"testing"
+
+	checkSupport "github.com/devr-tools/codeguard/internal/codeguard/checks/support"
+)
+
+func TestParseScriptEnforcesAggregateBudget(t *testing.T) {
+	corpus := newFileCorpus()
+	data := []byte("const value = 1;\n")
+	corpus.scriptBytes = maxTreeSitterScanBytes - len(data)
+
+	if _, err := corpus.parseScript("first.ts", data, checkSupport.ScriptLangTypeScript); err != nil {
+		t.Fatalf("parse at budget boundary: %v", err)
+	}
+	if _, err := corpus.parseScript("second.ts", data, checkSupport.ScriptLangTypeScript); err == nil ||
+		!strings.Contains(err.Error(), "scan budget") {
+		t.Fatalf("parse beyond budget error = %v, want scan budget error", err)
+	}
+	if got := len(corpus.scripts); got != 1 {
+		t.Fatalf("cached script count = %d, want 1", got)
+	}
+}
+
+func TestParseScriptEnforcesFileCountBudget(t *testing.T) {
+	corpus := newFileCorpus()
+	corpus.scriptCount = maxTreeSitterScanFiles
+
+	if _, err := corpus.parseScript("extra.ts", nil, checkSupport.ScriptLangTypeScript); err == nil ||
+		!strings.Contains(err.Error(), "scan budget") {
+		t.Fatalf("parse beyond file budget error = %v, want scan budget error", err)
+	}
+	if got := len(corpus.scripts); got != 0 {
+		t.Fatalf("cached script count = %d, want 0", got)
+	}
+}