interna/vulncheck: call updateInitPositions

The call was likely lost somewhere during the refactoring. We now
call it. Moreover, we call it where it makes the most sense, in
CallStacks computation. This is why this CL also moves some code.

Change-Id: I6665444e81d83fcb5436509f845983c6ba35d5a2
Reviewed-on: https://go-review.googlesource.com/c/vuln/+/506998
Run-TryBot: Zvonimir Pavlinovic <zpavlinovic@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
diff --git a/internal/scan/source.go b/internal/scan/source.go
index 59cde6b..cfdc122 100644
--- a/internal/scan/source.go
+++ b/internal/scan/source.go
@@ -7,11 +7,7 @@
 import (
 	"context"
 	"fmt"
-	"go/ast"
-	"go/token"
 	"path/filepath"
-	"strconv"
-	"strings"
 
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/vuln/internal"
@@ -212,99 +208,3 @@
 
 	return len(depPkgs), len(depMods)
 }
-
-// updateInitPositions populates non-existing positions of init functions
-// and their respective calls in callStacks (see #51575).
-func updateInitPositions(callStacks map[*vulncheck.Vuln]vulncheck.CallStack) {
-	for _, cs := range callStacks {
-		for i := range cs {
-			updateInitPosition(&cs[i])
-			if i != len(cs)-1 {
-				updateInitCallPosition(&cs[i], cs[i+1])
-			}
-		}
-	}
-}
-
-// updateInitCallPosition updates the position of a call to init in a stack frame, if
-// one already does not exist:
-//
-//	P1.init -> P2.init: position of call to P2.init is the position of "import P2"
-//	statement in P1
-//
-//	P.init -> P.init#d: P.init is an implicit init. We say it calls the explicit
-//	P.init#d at the place of "package P" statement.
-func updateInitCallPosition(curr *vulncheck.StackEntry, next vulncheck.StackEntry) {
-	call := curr.Call
-	if !isInit(next.Function) || (call.Pos != nil && call.Pos.IsValid()) {
-		// Skip non-init functions and inits whose call site position is available.
-		return
-	}
-
-	var pos token.Position
-	if curr.Function.Name == "init" && curr.Function.Package == next.Function.Package {
-		// We have implicit P.init calling P.init#d. Set the call position to
-		// be at "package P" statement position.
-		pos = packageStatementPos(curr.Function.Package)
-	} else {
-		// Choose the beginning of the import statement as the position.
-		pos = importStatementPos(curr.Function.Package, next.Function.Package.PkgPath)
-	}
-
-	call.Pos = &pos
-}
-
-func importStatementPos(pkg *packages.Package, importPath string) token.Position {
-	var importSpec *ast.ImportSpec
-spec:
-	for _, f := range pkg.Syntax {
-		for _, impSpec := range f.Imports {
-			// Import spec paths have quotation marks.
-			impSpecPath, err := strconv.Unquote(impSpec.Path.Value)
-			if err != nil {
-				panic(fmt.Sprintf("import specification: package path has no quotation marks: %v", err))
-			}
-			if impSpecPath == importPath {
-				importSpec = impSpec
-				break spec
-			}
-		}
-	}
-
-	if importSpec == nil {
-		// for sanity, in case of a wild call graph imprecision
-		return token.Position{}
-	}
-
-	// Choose the beginning of the import statement as the position.
-	return pkg.Fset.Position(importSpec.Pos())
-}
-
-func packageStatementPos(pkg *packages.Package) token.Position {
-	if len(pkg.Syntax) == 0 {
-		return token.Position{}
-	}
-	// Choose beginning of the package statement as the position. Pick
-	// the first file since it is as good as any.
-	return pkg.Fset.Position(pkg.Syntax[0].Package)
-}
-
-// updateInitPosition updates the position of P.init function in a stack frame if one
-// is not available. The new position is the position of the "package P" statement.
-func updateInitPosition(se *vulncheck.StackEntry) {
-	fun := se.Function
-	if !isInit(fun) || (fun.Pos != nil && fun.Pos.IsValid()) {
-		// Skip non-init functions and inits whose position is available.
-		return
-	}
-
-	pos := packageStatementPos(fun.Package)
-	fun.Pos = &pos
-}
-
-func isInit(f *vulncheck.FuncNode) bool {
-	// A source init function, or anonymous functions used in inits, will
-	// be named "init#x" by vulncheck (more precisely, ssa), where x is a
-	// positive integer. Implicit inits are named simply "init".
-	return f.Name == "init" || strings.HasPrefix(f.Name, "init#")
-}
diff --git a/internal/scan/source_test.go b/internal/scan/source_test.go
index ec0c747..694390b 100644
--- a/internal/scan/source_test.go
+++ b/internal/scan/source_test.go
@@ -5,19 +5,10 @@
 package scan
 
 import (
-	"context"
-	"fmt"
-	"path"
-	"path/filepath"
 	"strings"
 	"testing"
 
-	"github.com/google/go-cmp/cmp"
-	"golang.org/x/tools/go/packages/packagestest"
-	"golang.org/x/vuln/internal/client"
 	"golang.org/x/vuln/internal/govulncheck"
-	"golang.org/x/vuln/internal/osv"
-	"golang.org/x/vuln/internal/vulncheck"
 )
 
 func TestSummarizeCallStack(t *testing.T) {
@@ -83,142 +74,3 @@
 	}
 	return f
 }
-
-// TestInits checks for correct positions of init functions
-// and their respective calls (see #51575).
-func TestInits(t *testing.T) {
-	testClient, err := client.NewInMemoryClient(
-		[]*osv.Entry{
-			{
-				ID: "A", Affected: []osv.Affected{{Module: osv.Module{Path: "golang.org/amod"}, Ranges: []osv.Range{{Type: osv.RangeTypeSemver}},
-					EcosystemSpecific: osv.EcosystemSpecific{Packages: []osv.Package{{
-						Path: "golang.org/amod/avuln", Symbols: []string{"A"}},
-					}},
-				}},
-			},
-			{
-				ID: "C", Affected: []osv.Affected{{Module: osv.Module{Path: "golang.org/cmod"}, Ranges: []osv.Range{{Type: osv.RangeTypeSemver}},
-					EcosystemSpecific: osv.EcosystemSpecific{Packages: []osv.Package{{
-						Path: "golang.org/cmod/cvuln", Symbols: []string{"C"}},
-					}},
-				}},
-			},
-		})
-	if err != nil {
-		t.Fatal(err)
-	}
-	e := packagestest.Export(t, packagestest.Modules, []packagestest.Module{
-		{
-			Name: "golang.org/entry",
-			Files: map[string]interface{}{
-				"x/x.go": `
-			package x
-
-			import (
-				_ "golang.org/amod/avuln"
-				_ "golang.org/bmod/b"
-			)
-			`,
-			},
-		},
-		{
-			Name: "golang.org/amod@v0.5.0",
-			Files: map[string]interface{}{"avuln/avuln.go": `
-			package avuln
-
-			func init() {
-				A()
-			}
-
-			func A() {}
-			`},
-		},
-		{
-			Name: "golang.org/bmod@v0.5.0",
-			Files: map[string]interface{}{"b/b.go": `
-			package b
-
-			import _ "golang.org/cmod/cvuln"
-			`},
-		},
-		{
-			Name: "golang.org/cmod@v0.5.0",
-			Files: map[string]interface{}{"cvuln/cvuln.go": `
-			package cvuln
-
-			var x int = C()
-
-			func C() int {
-				return 0
-			}
-			`},
-		},
-	})
-	defer e.Cleanup()
-
-	// Load x as entry package.
-	graph := vulncheck.NewPackageGraph("go1.18")
-	pkgs, err := graph.LoadPackages(e.Config, nil, []string{path.Join(e.Temp(), "entry/x")})
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(pkgs) != 1 {
-		t.Fatal("failed to load x test package")
-	}
-	cfg := &govulncheck.Config{ScanLevel: "symbol"}
-	result, err := vulncheck.Source(context.Background(), pkgs, cfg, testClient, graph)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	cs := vulncheck.CallStacks(result)
-	updateInitPositions(cs)
-
-	want := map[string][]string{
-		"A": {
-			// Entry init's position is the package statement.
-			// It calls avuln.init at avuln import statement.
-			"N:golang.org/entry/x.init	F:x.go:2:4	C:x.go:5:5",
-			// implicit avuln.init is calls explicit init at the avuln
-			// package statement.
-			"N:golang.org/amod/avuln.init	F:avuln.go:2:4	C:avuln.go:2:4",
-			"N:golang.org/amod/avuln.init#1	F:avuln.go:4:9	C:avuln.go:5:6",
-			"N:golang.org/amod/avuln.A	F:avuln.go:8:9	C:",
-		},
-		"C": {
-			"N:golang.org/entry/x.init	F:x.go:2:4	C:x.go:6:5",
-			"N:golang.org/bmod/b.init	F:b.go:2:4	C:b.go:4:11",
-			"N:golang.org/cmod/cvuln.init	F:cvuln.go:2:4	C:cvuln.go:4:17",
-			"N:golang.org/cmod/cvuln.C	F:cvuln.go:6:9	C:",
-		},
-	}
-	if diff := cmp.Diff(want, strStacks(cs)); diff != "" {
-		t.Errorf("modules mismatch (-want, +got):\n%s", diff)
-	}
-}
-
-// strStacks creates a string representation of a call stacks map where
-// vulnerability is represented with its ID and stack entry is a string
-// "N:<package path.function name>  F:<function position> C:< call position>"
-// File paths in positions consists of only file names.
-func strStacks(callStacks map[*vulncheck.Vuln]vulncheck.CallStack) map[string][]string {
-	m := make(map[string][]string)
-	for v, cs := range callStacks {
-		var scs []string
-		for _, se := range cs {
-			fPos := se.Function.Pos
-			fp := fmt.Sprintf("%s:%d:%d", filepath.Base(fPos.Filename), fPos.Line, fPos.Column)
-
-			var cp string
-			if se.Call != nil && se.Call.Pos.IsValid() {
-				cPos := se.Call.Pos
-				cp = fmt.Sprintf("%s:%d:%d", filepath.Base(cPos.Filename), cPos.Line, cPos.Column)
-			}
-
-			sse := fmt.Sprintf("N:%s.%s\tF:%v\tC:%v", se.Function.Package.PkgPath, se.Function.Name, fp, cp)
-			scs = append(scs, sse)
-		}
-		m[v.OSV.ID] = scs
-	}
-	return m
-}
diff --git a/internal/vulncheck/witness.go b/internal/vulncheck/witness.go
index b1e943d..7de3758 100644
--- a/internal/vulncheck/witness.go
+++ b/internal/vulncheck/witness.go
@@ -7,10 +7,14 @@
 import (
 	"container/list"
 	"fmt"
+	"go/ast"
 	"go/token"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
+
+	"golang.org/x/tools/go/packages"
 )
 
 // CallStack is a call stack starting with a client
@@ -57,6 +61,8 @@
 		}()
 	}
 	wg.Wait()
+
+	updateInitPositions(stackPerVuln)
 	return stackPerVuln
 }
 
@@ -288,3 +294,99 @@
 	// should happen only for inits
 	return f1.String() < f2.String()
 }
+
+// updateInitPositions populates non-existing positions of init functions
+// and their respective calls in callStacks (see #51575).
+func updateInitPositions(callStacks map[*Vuln]CallStack) {
+	for _, cs := range callStacks {
+		for i := range cs {
+			updateInitPosition(&cs[i])
+			if i != len(cs)-1 {
+				updateInitCallPosition(&cs[i], cs[i+1])
+			}
+		}
+	}
+}
+
+// updateInitCallPosition updates the position of a call to init in a stack frame, if
+// one already does not exist:
+//
+//	P1.init -> P2.init: position of call to P2.init is the position of "import P2"
+//	statement in P1
+//
+//	P.init -> P.init#d: P.init is an implicit init. We say it calls the explicit
+//	P.init#d at the place of "package P" statement.
+func updateInitCallPosition(curr *StackEntry, next StackEntry) {
+	call := curr.Call
+	if !isInit(next.Function) || (call.Pos != nil && call.Pos.IsValid()) {
+		// Skip non-init functions and inits whose call site position is available.
+		return
+	}
+
+	var pos token.Position
+	if curr.Function.Name == "init" && curr.Function.Package == next.Function.Package {
+		// We have implicit P.init calling P.init#d. Set the call position to
+		// be at "package P" statement position.
+		pos = packageStatementPos(curr.Function.Package)
+	} else {
+		// Choose the beginning of the import statement as the position.
+		pos = importStatementPos(curr.Function.Package, next.Function.Package.PkgPath)
+	}
+
+	call.Pos = &pos
+}
+
+func importStatementPos(pkg *packages.Package, importPath string) token.Position {
+	var importSpec *ast.ImportSpec
+spec:
+	for _, f := range pkg.Syntax {
+		for _, impSpec := range f.Imports {
+			// Import spec paths have quotation marks.
+			impSpecPath, err := strconv.Unquote(impSpec.Path.Value)
+			if err != nil {
+				panic(fmt.Sprintf("import specification: package path has no quotation marks: %v", err))
+			}
+			if impSpecPath == importPath {
+				importSpec = impSpec
+				break spec
+			}
+		}
+	}
+
+	if importSpec == nil {
+		// for sanity, in case of a wild call graph imprecision
+		return token.Position{}
+	}
+
+	// Choose the beginning of the import statement as the position.
+	return pkg.Fset.Position(importSpec.Pos())
+}
+
+func packageStatementPos(pkg *packages.Package) token.Position {
+	if len(pkg.Syntax) == 0 {
+		return token.Position{}
+	}
+	// Choose beginning of the package statement as the position. Pick
+	// the first file since it is as good as any.
+	return pkg.Fset.Position(pkg.Syntax[0].Package)
+}
+
+// updateInitPosition updates the position of P.init function in a stack frame if one
+// is not available. The new position is the position of the "package P" statement.
+func updateInitPosition(se *StackEntry) {
+	fun := se.Function
+	if !isInit(fun) || (fun.Pos != nil && fun.Pos.IsValid()) {
+		// Skip non-init functions and inits whose position is available.
+		return
+	}
+
+	pos := packageStatementPos(fun.Package)
+	fun.Pos = &pos
+}
+
+func isInit(f *FuncNode) bool {
+	// A source init function, or anonymous functions used in inits, will
+	// be named "init#x" by vulncheck (more precisely, ssa), where x is a
+	// positive integer. Implicit inits are named simply "init".
+	return f.Name == "init" || strings.HasPrefix(f.Name, "init#")
+}
diff --git a/internal/vulncheck/witness_test.go b/internal/vulncheck/witness_test.go
index 6dbe634..619541a 100644
--- a/internal/vulncheck/witness_test.go
+++ b/internal/vulncheck/witness_test.go
@@ -5,11 +5,19 @@
 package vulncheck
 
 import (
+	"context"
+	"fmt"
+	"path"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"golang.org/x/tools/go/packages"
+	"golang.org/x/tools/go/packages/packagestest"
+	"golang.org/x/vuln/internal/client"
+	"golang.org/x/vuln/internal/govulncheck"
 	"golang.org/x/vuln/internal/osv"
 )
 
@@ -98,3 +106,139 @@
 		t.Errorf("want %v; got %v", want, got)
 	}
 }
+
+// TestInits checks for correct positions of init functions
+// and their respective calls (see #51575).
+func TestInits(t *testing.T) {
+	testClient, err := client.NewInMemoryClient(
+		[]*osv.Entry{
+			{
+				ID: "A", Affected: []osv.Affected{{Module: osv.Module{Path: "golang.org/amod"}, Ranges: []osv.Range{{Type: osv.RangeTypeSemver}},
+					EcosystemSpecific: osv.EcosystemSpecific{Packages: []osv.Package{{
+						Path: "golang.org/amod/avuln", Symbols: []string{"A"}},
+					}},
+				}},
+			},
+			{
+				ID: "C", Affected: []osv.Affected{{Module: osv.Module{Path: "golang.org/cmod"}, Ranges: []osv.Range{{Type: osv.RangeTypeSemver}},
+					EcosystemSpecific: osv.EcosystemSpecific{Packages: []osv.Package{{
+						Path: "golang.org/cmod/cvuln", Symbols: []string{"C"}},
+					}},
+				}},
+			},
+		})
+	if err != nil {
+		t.Fatal(err)
+	}
+	e := packagestest.Export(t, packagestest.Modules, []packagestest.Module{
+		{
+			Name: "golang.org/entry",
+			Files: map[string]interface{}{
+				"x/x.go": `
+			package x
+
+			import (
+				_ "golang.org/amod/avuln"
+				_ "golang.org/bmod/b"
+			)
+			`,
+			},
+		},
+		{
+			Name: "golang.org/amod@v0.5.0",
+			Files: map[string]interface{}{"avuln/avuln.go": `
+			package avuln
+
+			func init() {
+				A()
+			}
+
+			func A() {}
+			`},
+		},
+		{
+			Name: "golang.org/bmod@v0.5.0",
+			Files: map[string]interface{}{"b/b.go": `
+			package b
+
+			import _ "golang.org/cmod/cvuln"
+			`},
+		},
+		{
+			Name: "golang.org/cmod@v0.5.0",
+			Files: map[string]interface{}{"cvuln/cvuln.go": `
+			package cvuln
+
+			var x int = C()
+
+			func C() int {
+				return 0
+			}
+			`},
+		},
+	})
+	defer e.Cleanup()
+
+	// Load x as entry package.
+	graph := NewPackageGraph("go1.18")
+	pkgs, err := graph.LoadPackages(e.Config, nil, []string{path.Join(e.Temp(), "entry/x")})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(pkgs) != 1 {
+		t.Fatal("failed to load x test package")
+	}
+	cfg := &govulncheck.Config{ScanLevel: "symbol"}
+	result, err := Source(context.Background(), pkgs, cfg, testClient, graph)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cs := CallStacks(result)
+	want := map[string][]string{
+		"A": {
+			// Entry init's position is the package statement.
+			// It calls avuln.init at avuln import statement.
+			"N:golang.org/entry/x.init	F:x.go:2:4	C:x.go:5:5",
+			// implicit avuln.init is calls explicit init at the avuln
+			// package statement.
+			"N:golang.org/amod/avuln.init	F:avuln.go:2:4	C:avuln.go:2:4",
+			"N:golang.org/amod/avuln.init#1	F:avuln.go:4:9	C:avuln.go:5:6",
+			"N:golang.org/amod/avuln.A	F:avuln.go:8:9	C:",
+		},
+		"C": {
+			"N:golang.org/entry/x.init	F:x.go:2:4	C:x.go:6:5",
+			"N:golang.org/bmod/b.init	F:b.go:2:4	C:b.go:4:11",
+			"N:golang.org/cmod/cvuln.init	F:cvuln.go:2:4	C:cvuln.go:4:17",
+			"N:golang.org/cmod/cvuln.C	F:cvuln.go:6:9	C:",
+		},
+	}
+	if diff := cmp.Diff(want, fullStacksToString(cs)); diff != "" {
+		t.Errorf("modules mismatch (-want, +got):\n%s", diff)
+	}
+}
+
+// fullStacksToString is like stacksToString but the stack stringification
+// is a slice of strings, each containing detailed information on each on
+// the corresponding frame.
+func fullStacksToString(callStacks map[*Vuln]CallStack) map[string][]string {
+	m := make(map[string][]string)
+	for v, cs := range callStacks {
+		var scs []string
+		for _, se := range cs {
+			fPos := se.Function.Pos
+			fp := fmt.Sprintf("%s:%d:%d", filepath.Base(fPos.Filename), fPos.Line, fPos.Column)
+
+			var cp string
+			if se.Call != nil && se.Call.Pos.IsValid() {
+				cPos := se.Call.Pos
+				cp = fmt.Sprintf("%s:%d:%d", filepath.Base(cPos.Filename), cPos.Line, cPos.Column)
+			}
+
+			sse := fmt.Sprintf("N:%s.%s\tF:%v\tC:%v", se.Function.Package.PkgPath, se.Function.Name, fp, cp)
+			scs = append(scs, sse)
+		}
+		m[v.OSV.ID] = scs
+	}
+	return m
+}