internal: properly fetch modules in source mode

Loads modules using go list instead of inferring through packages. This
fixes a few module-level edge cases where a vuln wouldn't be counted.

Change-Id: I24e0ffa73f47451806d88aa672ca8ef7a72fc2bb
Reviewed-on: https://go-review.googlesource.com/c/vuln/+/529278
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Zvonimir Pavlinovic <zpavlinovic@google.com>
diff --git a/cmd/govulncheck/testdata/source_multientry_json.ct b/cmd/govulncheck/testdata/source_multientry_json.ct
index 7de41ad..082f7ac 100644
--- a/cmd/govulncheck/testdata/source_multientry_json.ct
+++ b/cmd/govulncheck/testdata/source_multientry_json.ct
@@ -14,7 +14,7 @@
 }
 {
   "progress": {
-    "message": "Scanning your code and P packages across M dependent module for known vulnerabilities..."
+    "message": "Scanning your code and P packages across M dependent modules for known vulnerabilities..."
   }
 }
 {
diff --git a/cmd/govulncheck/testdata/source_multientry_text.ct b/cmd/govulncheck/testdata/source_multientry_text.ct
index 8bb2bcb..5918cd4 100644
--- a/cmd/govulncheck/testdata/source_multientry_text.ct
+++ b/cmd/govulncheck/testdata/source_multientry_text.ct
@@ -1,7 +1,7 @@
 #####
 # Test for multiple call stacks in source mode with expanded traces
 $ govulncheck -C ${moddir}/multientry . --> FAIL 3
-Scanning your code and P packages across M dependent module for known vulnerabilities...
+Scanning your code and P packages across M dependent modules for known vulnerabilities...
 
 Vulnerability #1: GO-2021-0113
     Due to improper index calculation, an incorrectly formatted language tag can
@@ -40,7 +40,7 @@
 #####
 # Test for multple call stacks in source mode with expanded traces
 $ govulncheck -C ${moddir}/multientry -show=traces ./... --> FAIL 3
-Scanning your code and P packages across M dependent module for known vulnerabilities...
+Scanning your code and P packages across M dependent modules for known vulnerabilities...
 
 Vulnerability #1: GO-2021-0113
     Due to improper index calculation, an incorrectly formatted language tag can
diff --git a/cmd/govulncheck/testdata/source_subdir_text.ct b/cmd/govulncheck/testdata/source_subdir_text.ct
index 1d52025..da6b08e 100644
--- a/cmd/govulncheck/testdata/source_subdir_text.ct
+++ b/cmd/govulncheck/testdata/source_subdir_text.ct
@@ -1,7 +1,7 @@
 #####
 # Test govulncheck runs on the subdirectory of a module
 $ govulncheck -C ${moddir}/vuln/subdir . --> FAIL 3
-Scanning your code and P packages across M dependent module for known vulnerabilities...
+Scanning your code and P packages across M dependent modules for known vulnerabilities...
 
 Vulnerability #1: GO-2021-0113
     Due to improper index calculation, an incorrectly formatted language tag can
@@ -19,7 +19,7 @@
 
 Found 0 vulnerabilities in packages that you import, but there are no
 call stacks leading to the use of these vulnerabilities. You may not
-need to take any action. There are also 2 vulnerabilities in modules
+need to take any action. There are also 4 vulnerabilities in modules
 that you require that are neither imported nor called.
 See https://pkg.go.dev/golang.org/x/vuln/cmd/govulncheck for details.
 
@@ -32,7 +32,24 @@
     Found in: net/http@go1.18
     Fixed in: net/http@go1.18.6
 
-Vulnerability #2: GO-2020-0015
+Vulnerability #2: GO-2021-0265
+    A maliciously crafted path can cause Get and other query functions to
+    consume excessive amounts of CPU and time.
+  More info: https://pkg.go.dev/vuln/GO-2021-0265
+  Module: github.com/tidwall/gjson
+    Found in: github.com/tidwall/gjson@v1.6.5
+    Fixed in: github.com/tidwall/gjson@v1.9.3
+
+Vulnerability #3: GO-2021-0054
+    Due to improper bounds checking, maliciously crafted JSON objects can cause
+    an out-of-bounds panic. If parsing user input, this may be used as a denial
+    of service vector.
+  More info: https://pkg.go.dev/vuln/GO-2021-0054
+  Module: github.com/tidwall/gjson
+    Found in: github.com/tidwall/gjson@v1.6.5
+    Fixed in: github.com/tidwall/gjson@v1.6.6
+
+Vulnerability #4: GO-2020-0015
     An attacker could provide a single byte to a UTF16 decoder instantiated with
     UseBOM or ExpectBOM to trigger an infinite loop if the String function on
     the Decoder is called, or the Decoder is passed to transform.String. If used
@@ -50,7 +67,7 @@
 #####
 # Test govulncheck runs on the subdirectory of a module
 $ govulncheck -C ${moddir}/vuln/subdir -show=traces . --> FAIL 3
-Scanning your code and P packages across M dependent module for known vulnerabilities...
+Scanning your code and P packages across M dependent modules for known vulnerabilities...
 
 Vulnerability #1: GO-2021-0113
     Due to improper index calculation, an incorrectly formatted language tag can
@@ -70,7 +87,7 @@
 
 Found 0 vulnerabilities in packages that you import, but there are no
 call stacks leading to the use of these vulnerabilities. You may not
-need to take any action. There are also 2 vulnerabilities in modules
+need to take any action. There are also 4 vulnerabilities in modules
 that you require that are neither imported nor called.
 See https://pkg.go.dev/golang.org/x/vuln/cmd/govulncheck for details.
 
@@ -83,7 +100,24 @@
     Found in: net/http@go1.18
     Fixed in: net/http@go1.18.6
 
-Vulnerability #2: GO-2020-0015
+Vulnerability #2: GO-2021-0265
+    A maliciously crafted path can cause Get and other query functions to
+    consume excessive amounts of CPU and time.
+  More info: https://pkg.go.dev/vuln/GO-2021-0265
+  Module: github.com/tidwall/gjson
+    Found in: github.com/tidwall/gjson@v1.6.5
+    Fixed in: github.com/tidwall/gjson@v1.9.3
+
+Vulnerability #3: GO-2021-0054
+    Due to improper bounds checking, maliciously crafted JSON objects can cause
+    an out-of-bounds panic. If parsing user input, this may be used as a denial
+    of service vector.
+  More info: https://pkg.go.dev/vuln/GO-2021-0054
+  Module: github.com/tidwall/gjson
+    Found in: github.com/tidwall/gjson@v1.6.5
+    Fixed in: github.com/tidwall/gjson@v1.6.6
+
+Vulnerability #4: GO-2020-0015
     An attacker could provide a single byte to a UTF16 decoder instantiated with
     UseBOM or ExpectBOM to trigger an infinite loop if the String function on
     the Decoder is called, or the Decoder is passed to transform.String. If used
diff --git a/internal/scan/source.go b/internal/scan/source.go
index 1dd32c6..c6d7eeb 100644
--- a/internal/scan/source.go
+++ b/internal/scan/source.go
@@ -11,7 +11,6 @@
 	"sort"
 
 	"golang.org/x/tools/go/packages"
-	"golang.org/x/vuln/internal"
 	"golang.org/x/vuln/internal/client"
 	"golang.org/x/vuln/internal/govulncheck"
 	"golang.org/x/vuln/internal/vulncheck"
@@ -33,18 +32,17 @@
 		Tests: cfg.test,
 		Env:   cfg.env,
 	}
-	pkgs, err := graph.LoadPackages(pkgConfig, cfg.tags, cfg.patterns)
+	mods, err := graph.LoadModules(pkgConfig)
 	if err != nil {
-		// Try to provide a meaningful and actionable error message.
-		if !fileExists(filepath.Join(dir, "go.mod")) {
-			return fmt.Errorf("govulncheck: %v", errNoGoMod)
-		}
-		if isGoVersionMismatchError(err) {
-			return fmt.Errorf("govulncheck: %v\n\n%v", errGoVersionMismatch, err)
-		}
-		return fmt.Errorf("govulncheck: loading packages: %w", err)
+		return parseLoadError(err, dir, false)
 	}
-	if err := handler.Progress(sourceProgressMessage(pkgs)); err != nil {
+	if cfg.ScanLevel.WantPackages() {
+		pkgs, err = graph.LoadPackages(pkgConfig, cfg.tags, cfg.patterns)
+		if err != nil {
+			return parseLoadError(err, dir, true)
+		}
+	}
+	if err := handler.Progress(sourceProgressMessage(pkgs, len(mods)-1)); err != nil {
 		return err
 	}
 
@@ -133,14 +131,14 @@
 // is 0, then the following message is returned
 //
 //	"No packages matching the provided pattern."
-func sourceProgressMessage(topPkgs []*packages.Package) *govulncheck.Progress {
+func sourceProgressMessage(topPkgs []*packages.Package, mods int) *govulncheck.Progress {
 	if len(topPkgs) == 0 {
 		// The package pattern is valid, but no packages are matching.
 		// Example is pkg/strace/... (see #59623).
 		return &govulncheck.Progress{Message: "No packages matching the provided pattern."}
 	}
 
-	pkgs, mods := depPkgsAndMods(topPkgs)
+	pkgs := depPkgs(topPkgs)
 
 	pkgsPhrase := fmt.Sprintf("%d package", pkgs)
 	if pkgs != 1 {
@@ -156,12 +154,10 @@
 	return &govulncheck.Progress{Message: msg}
 }
 
-// depPkgsAndMods returns the number of packages that
-// topPkgs depend on and the number of their modules.
-func depPkgsAndMods(topPkgs []*packages.Package) (int, int) {
+// depPkgs returns the number of packages that topPkgs depend on
+func depPkgs(topPkgs []*packages.Package) int {
 	tops := make(map[string]bool)
 	depPkgs := make(map[string]bool)
-	depMods := make(map[string]bool)
 
 	for _, t := range topPkgs {
 		tops[t.PkgPath] = true
@@ -185,11 +181,6 @@
 		// as a dependent package.
 		if !tops[path] {
 			depPkgs[path] = true
-			if p.Module != nil &&
-				p.Module.Path != internal.GoStdModulePath && // no module for stdlib
-				p.Module.Path != internal.UnknownModulePath { // no module for unknown
-				depMods[p.Module.Path] = true
-			}
 		}
 
 		for _, d := range p.Imports {
@@ -201,5 +192,22 @@
 		visit(t, true)
 	}
 
-	return len(depPkgs), len(depMods)
+	return len(depPkgs)
+}
+
+// parseLoadError takes a non-nil error and tries to provide a meaningful
+// and actionable error message to surface for the end user.
+func parseLoadError(err error, dir string, pkgs bool) error {
+	if !fileExists(filepath.Join(dir, "go.mod")) {
+		return fmt.Errorf("govulncheck: %v", errNoGoMod)
+	}
+	if isGoVersionMismatchError(err) {
+		return fmt.Errorf("govulncheck: %v\n\n%v", errGoVersionMismatch, err)
+	}
+
+	level := "modules"
+	if pkgs {
+		level = "packages"
+	}
+	return fmt.Errorf("govulncheck: loading %s: %w", level, err)
 }
diff --git a/internal/scan/testdata/multi-stack-modlevel.json b/internal/scan/testdata/multi-stack-modlevel.json
new file mode 100644
index 0000000..edbbab6
--- /dev/null
+++ b/internal/scan/testdata/multi-stack-modlevel.json
@@ -0,0 +1,98 @@
+{
+  "config": {
+    "protocol_version": "v0.1.0",
+    "scanner_name": "govulncheck"
+  }
+}
+{
+  "osv": {
+    "id": "GO-0000-0001",
+    "modified": "0001-01-01T00:00:00Z",
+    "published": "0001-01-01T00:00:00Z",
+    "details": "Third-party vulnerability",
+    "affected": [
+      {
+        "package": {
+          "name": "golang.org/vmod",
+          "ecosystem": ""
+        },
+        "ecosystem_specific": {
+          "imports": [
+            {
+              "goos": [
+                "amd"
+              ]
+            }
+          ]
+        }
+      }
+    ],
+    "database_specific": {
+      "url": "https://pkg.go.dev/vuln/GO-0000-0001"
+    }
+  }
+}
+{
+  "finding": {
+    "osv": "GO-0000-0001",
+    "fixed_version": "v0.1.3",
+    "trace": [
+      {
+        "module": "golang.org/vmod",
+        "version": "v0.0.1"
+      }
+    ]
+  }
+}
+{
+  "finding": {
+    "osv": "GO-0000-0001",
+    "fixed_version": "v0.1.3",
+    "trace": [
+      {
+        "module": "golang.org/vmod",
+        "version": "v0.0.1",
+        "package": "vmod",
+        "function": "VulnFoo"
+      },
+      {
+        "module": "golang.org/main",
+        "version": "v0.0.1",
+        "package": "main",
+        "function": "main"
+      }
+    ]
+  }
+}
+{
+  "osv": {
+    "id": "GO-0000-0002",
+    "modified": "0001-01-01T00:00:00Z",
+    "published": "0001-01-01T00:00:00Z",
+    "details": "Third-party vulnerability",
+    "affected": [
+      {
+        "package": {
+          "name": "golang.org/vmod",
+          "ecosystem": ""
+        },
+        "ecosystem_specific": {}
+      }
+    ],
+    "database_specific": {
+      "url": "https://pkg.go.dev/vuln/GO-0000-0002"
+    }
+  }
+}
+{
+  "finding": {
+    "osv": "GO-0000-0002",
+    "fixed_version": "v0.1.3",
+    "trace": [
+      {
+        "module": "golang.org/vmod",
+        "version": "v0.0.1"
+      }
+    ]
+  }
+}
\ No newline at end of file
diff --git a/internal/scan/testdata/multi-stack-modlevel.txt b/internal/scan/testdata/multi-stack-modlevel.txt
new file mode 100644
index 0000000..80acaee
--- /dev/null
+++ b/internal/scan/testdata/multi-stack-modlevel.txt
@@ -0,0 +1,28 @@
+Vulnerability #1: GO-0000-0001
+    Third-party vulnerability
+  More info: https://pkg.go.dev/vuln/GO-0000-0001
+  Module: golang.org/vmod
+    Found in: golang.org/vmod@v0.0.1
+    Fixed in: golang.org/vmod@v0.1.3
+    Platforms: amd
+    Example traces found:
+      #1: main.main calls vmod.VulnFoo
+
+=== Informational ===
+
+Found 0 vulnerabilities in packages that you import, but there are no
+call stacks leading to the use of these vulnerabilities. You may not
+need to take any action. There is also 1 vulnerability in modules
+that you require that is neither imported nor called.
+See https://pkg.go.dev/golang.org/x/vuln/cmd/govulncheck for details.
+
+Vulnerability #1: GO-0000-0002
+    Third-party vulnerability
+  More info: https://pkg.go.dev/vuln/GO-0000-0002
+  Module: golang.org/vmod
+    Found in: golang.org/vmod@v0.0.1
+    Fixed in: golang.org/vmod@v0.1.3
+
+Your code is affected by 1 vulnerability from 1 module.
+
+Share feedback at https://go.dev/s/govulncheck-feedback.
diff --git a/internal/vulncheck/packages.go b/internal/vulncheck/packages.go
index a829b12..94cbb12 100644
--- a/internal/vulncheck/packages.go
+++ b/internal/vulncheck/packages.go
@@ -36,13 +36,14 @@
 }
 
 func (g *PackageGraph) LoadModules(cfg *packages.Config) (mods []*packages.Module, err error) {
-	cmd := exec.Command("go", "list", "-m", "-json", "all")
+	cmd := exec.Command("go", "list", "-m", "-json", "-mod=mod", "all")
 	cmd.Env = append(cmd.Env, cfg.Env...)
 	cmd.Dir = cfg.Dir
 	out, err := cmd.Output()
 	if err != nil {
 		return nil, err
 	}
+
 	dec := json.NewDecoder(bytes.NewReader(out))
 	for dec.More() {
 		var m *packages.Module
diff --git a/internal/vulncheck/source.go b/internal/vulncheck/source.go
index 442e6f4..41b4223 100644
--- a/internal/vulncheck/source.go
+++ b/internal/vulncheck/source.go
@@ -51,7 +51,10 @@
 		}()
 	}
 
-	mods := extractModules(pkgs)
+	var mods []*packages.Module
+	for _, m := range graph.modules {
+		mods = append(mods, m)
+	}
 	mv, err := FetchVulnerabilities(ctx, client, mods)
 	if err != nil {
 		return nil, err
@@ -332,36 +335,3 @@
 		}
 	}
 }
-
-// extractModules collects modules in `pkgs` up to uniqueness of
-// module path and version.
-func extractModules(pkgs []*packages.Package) []*packages.Module {
-	modMap := map[string]*packages.Module{}
-	seen := map[*packages.Package]bool{}
-	var extract func(*packages.Package, map[string]*packages.Module)
-	extract = func(pkg *packages.Package, modMap map[string]*packages.Module) {
-		if pkg == nil || seen[pkg] {
-			return
-		}
-		if pkg.Module != nil {
-			if pkg.Module.Replace != nil {
-				modMap[pkg.Module.Replace.Path] = pkg.Module
-			} else {
-				modMap[pkg.Module.Path] = pkg.Module
-			}
-		}
-		seen[pkg] = true
-		for _, imp := range pkg.Imports {
-			extract(imp, modMap)
-		}
-	}
-	for _, pkg := range pkgs {
-		extract(pkg, modMap)
-	}
-
-	modules := []*packages.Module{}
-	for _, mod := range modMap {
-		modules = append(modules, mod)
-	}
-	return modules
-}