ericchiang · May 10, 2022 · Apr 15, 2022
diff --git a/jar/walker.go b/jar/walker.go
@@ -62,12 +62,20 @@ type Walker struct {
 	HandleReport func(path string, r *Report)
 	// HandleRewrite is called when a JAR is rewritten successfully.
 	HandleRewrite func(path string, r *Report)
+	// Parser will be used when checking JARs, if provided. If
+	// unset, a Parser with sensible defaults will be created.
+	Parser *Parser
 }
 
 // Walk attempts to scan a directory for vulnerable JARs.
 func (w *Walker) Walk(dir string) error {
+	p := w.Parser
+	if p == nil {
+		p = &Parser{}
+	}
+
 	fsys := os.DirFS(dir)
-	wk := walker{w, fsys, dir}
+	wk := walker{w, fsys, dir, p}
 
 	return fs.WalkDir(fsys, ".", func(p string, d fs.DirEntry, err error) error {
 		if err != nil {
@@ -88,6 +96,8 @@ type walker struct {
 	*Walker
 	fs  fs.FS
 	dir string
+	// p is the Parser to use for this walk. p is guaranteed to be non-nil.
+	p *Parser
 }
 
 func (w *walker) filepath(path string) string {
@@ -154,7 +164,7 @@ func (w *walker) visit(p string, d fs.DirEntry) error {
 	if !IsJAR(zr) {
 		return nil
 	}
-	r, err := Parse(zr)
+	r, err := w.p.Parse(zr)
 	if err != nil {
 		return fmt.Errorf("scanning jar: %v", err)
 	}

diff --git a/jar/walker_test.go b/jar/walker_test.go
@@ -16,6 +16,7 @@ package jar
 
 import (
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -177,3 +178,59 @@ func TestWalkerRewrite(t *testing.T) {
 		t.Errorf("walking filesystem after rewrite returned diff (-want, +got): %s", diff)
 	}
 }
+
+// TestNonDefaultParser verifies that Walker can be configured with a
+// non-default Parser by scanning a large JAR file with two
+// configurations: one where Parser.MaxBytes is larger than the file
+// size and one where Parser.MaxBytes is smaller than the file
+// size. It ensures that the first case succeeds and the second fails.
+func TestNonDefaultParser(t *testing.T) {
+	jar := "400mb_jar_in_jar.jar"
+
+	tempDir := t.TempDir()
+	src := testdataPath(jar)
+	dest := filepath.Join(tempDir, jar)
+	cpFile(t, dest, src)
+
+	tests := []struct {
+		desc     string
+		maxBytes int64
+		wantErr  bool
+	}{
+		{
+			desc:     "MaxBytes > JAR size",
+			maxBytes: 4 << 30, // 4GiB
+			wantErr:  false,
+		},
+		{
+			desc:     "MaxBytes < JAR size",
+			maxBytes: 4 << 20, // 4MiB
+			wantErr:  true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			var gotErr error
+
+			p := &Parser{MaxBytes: tc.maxBytes}
+			w := &Walker{
+				Parser: p,
+				HandleError: func(path string, err error) {
+					if err != nil && strings.HasSuffix(path, filepath.FromSlash("/"+jar)) {
+						gotErr = err
+					}
+				},
+			}
+			if err := w.Walk(tempDir); err != nil {
+				t.Errorf("Walk returned unexpected error: %v", err)
+			}
+
+			if tc.wantErr && gotErr == nil {
+				t.Error("Parser failed to generate expected error when scanning JARs > MaxBytes, got nil, want error")
+			} else if !tc.wantErr && gotErr != nil {
+				t.Errorf("Parser generated unexpected error when scanning JARs <= MaxBytes, got %v, want nil error", gotErr)
+			}
+		})
+	}
+}