Guided Remediation: Make VulnerabilityClient for OSV queries (#773)

Implementing
#766 (comment)
- Created `VulnerabilityClient` interface for OSV queries & to store
cache
- Renamed `ResolutionClient` to `DependencyClient`
- Made new `ResolutionClient` struct, that's just both
`DependencyClient` and `VulnerabilityClient` together

Author: michaelkedar

internal/remediation/relax.go

@@ -8,10 +8,11 @@ import (
 	"deps.dev/util/resolve"
 	"github.com/google/osv-scanner/internal/remediation/relaxer"
 	"github.com/google/osv-scanner/internal/resolution"
+	"github.com/google/osv-scanner/internal/resolution/client"
 )
 
 // ComputeRelaxPatches attempts to resolve each vulnerability found in result independently, returning the list of unique possible patches
-func ComputeRelaxPatches(ctx context.Context, cl resolve.Client, result *resolution.ResolutionResult, opts RemediationOptions) ([]resolution.ResolutionDiff, error) {
+func ComputeRelaxPatches(ctx context.Context, cl client.ResolutionClient, result *resolution.ResolutionResult, opts RemediationOptions) ([]resolution.ResolutionDiff, error) {
 	// Filter the original result just in case it hasn't been already
 	result.FilterVulns(opts.MatchVuln)
 
@@ -74,7 +75,7 @@ var errRelaxRemediateImpossible = errors.New("cannot fix vulns by relaxing")
 
 func tryRelaxRemediate(
 	ctx context.Context,
-	cl resolve.Client,
+	cl client.ResolutionClient,
 	orig *resolution.ResolutionResult,
 	vulnIDs []string,
 	opts RemediationOptions,

internal/resolution/client/resolution_client.go renamed to internal/resolution/client/client.go

@@ -4,9 +4,15 @@ import (
 	"context"
 
 	"deps.dev/util/resolve"
+	"github.com/google/osv-scanner/pkg/models"
 )
 
-type ResolutionClient interface {
+type ResolutionClient struct {
+	DependencyClient
+	VulnerabilityClient
+}
+
+type DependencyClient interface {
 	resolve.Client
 	// WriteCache writes a manifest-specific resolution cache.
 	WriteCache(filepath string) error
@@ -15,3 +21,9 @@ type ResolutionClient interface {
 	// PreFetch loads cache, then makes and caches likely queries needed for resolving a package with a list of requirements
 	PreFetch(ctx context.Context, requirements []resolve.RequirementVersion, manifestPath string)
 }
+
+type VulnerabilityClient interface {
+	// FindVulns finds the vulnerabilities affecting each of Nodes in the graph.
+	// The returned Vulnerabilities[i] corresponds to the vulnerabilities in g.Nodes[i].
+	FindVulns(g *resolve.Graph) ([]models.Vulnerabilities, error)
+}

internal/resolution/client/osv_client.go

@@ -0,0 +1,100 @@
+package client
+
+import (
+	"sync"
+
+	"deps.dev/util/resolve"
+	"github.com/google/osv-scanner/internal/resolution/util"
+	"github.com/google/osv-scanner/internal/utility/vulns"
+	"github.com/google/osv-scanner/pkg/lockfile"
+	"github.com/google/osv-scanner/pkg/models"
+	"github.com/google/osv-scanner/pkg/osv"
+	"golang.org/x/exp/maps"
+)
+
+type OSVClient struct {
+	// vulnCache caches all vulnerabilities affecting any versions of particular packages.
+	// We cache call vulns & manually check affected, rather than querying the affected versions directly
+	// since remediation needs to query for OSV vulnerabilities multiple times for the same packages.
+	vulnCache sync.Map // map[resolve.PackageKey][]models.Vulnerability
+	// TODO: This tends to get the full info of a lot of vulns that never show up in the dependency graphs.
+	// Worst case is something like PyPI:tensorflow, which has >600 vulns across all versions, but a specific version may be affected by 0.
+}
+
+func NewOSVClient() *OSVClient {
+	return &OSVClient{}
+}
+
+func (c *OSVClient) FindVulns(g *resolve.Graph) ([]models.Vulnerabilities, error) {
+	// Determine which packages we don't already have cached
+	toQuery := make(map[resolve.PackageKey]struct{})
+	for _, node := range g.Nodes[1:] { // skipping the root node
+		pk := node.Version.PackageKey
+		if _, ok := c.vulnCache.Load(pk); !ok {
+			toQuery[pk] = struct{}{}
+		}
+	}
+
+	// Query OSV for the missing records
+	if len(toQuery) > 0 {
+		pks := maps.Keys(toQuery)
+		var batchRequest osv.BatchedQuery
+		batchRequest.Queries = make([]*osv.Query, len(pks))
+		for i, pk := range pks {
+			batchRequest.Queries[i] = &osv.Query{
+				Package: osv.Package{
+					Name:      pk.Name,
+					Ecosystem: string(util.OSVEcosystem[pk.System]),
+				},
+				// Omitting the Version from the query gets all vulns affecting any version of the package
+				// (I'm not actually sure if this behaviour is explicitly documented anywhere)
+			}
+		}
+		batchResponse, err := osv.MakeRequest(batchRequest)
+		if err != nil {
+			return nil, err
+		}
+		hydrated, err := osv.Hydrate(batchResponse)
+		if err != nil {
+			return nil, err
+		}
+		// fill in the cache with the responses
+		for i, pk := range pks {
+			c.vulnCache.Store(pk, hydrated.Results[i].Vulns)
+		}
+	}
+
+	// Compute the actual affected vulnerabilities for each node
+	nodeVulns := make([]models.Vulnerabilities, len(g.Nodes))
+	// For convenience, include the root node as an empty slice in the results
+	for i, n := range g.Nodes {
+		if i == 0 {
+			continue
+		}
+		pkgVulnsAny, ok := c.vulnCache.Load(n.Version.PackageKey)
+		if !ok {
+			// This should be impossible
+			panic("vulnerability caching failed")
+		}
+		pkgVulns, ok := pkgVulnsAny.([]models.Vulnerability)
+		if !ok {
+			panic("vulnerability caching failed")
+		}
+
+		var affectedVulns []models.Vulnerability
+		pkgDetails := lockfile.PackageDetails{
+			Name:      n.Version.Name,
+			Version:   n.Version.Version,
+			Ecosystem: lockfile.Ecosystem(util.OSVEcosystem[n.Version.System]),
+			CompareAs: lockfile.Ecosystem(util.OSVEcosystem[n.Version.System]),
+		}
+		for _, vuln := range pkgVulns {
+			if vulns.IsAffected(vuln, pkgDetails) {
+				affectedVulns = append(affectedVulns, vuln)
+			}
+		}
+		nodeVulns[i] = affectedVulns
+	}
+
+	return nodeVulns, nil
+}

internal/resolution/client/override_client.go

@@ -7,20 +7,19 @@ import (
 	"deps.dev/util/resolve"
 )
 
-// OvverideClient wraps a resolve.Client, allowing for custom packages & versions to be added
+// OverrideClient wraps a DependencyClient, allowing for custom packages & versions to be added
 type OverrideClient struct {
-	c resolve.Client
-
+	DependencyClient
 	// Can't quite reuse resolve.LocalClient because it automatically creates dependencies
 	pkgVers map[resolve.PackageKey][]resolve.Version            // versions of a package
 	verDeps map[resolve.VersionKey][]resolve.RequirementVersion // dependencies of a version
 }
 
-func NewOverrideClient(c resolve.Client) *OverrideClient {
+func NewOverrideClient(c DependencyClient) *OverrideClient {
 	return &OverrideClient{
-		c:       c,
-		pkgVers: make(map[resolve.PackageKey][]resolve.Version),
-		verDeps: make(map[resolve.VersionKey][]resolve.RequirementVersion),
+		DependencyClient: c,
+		pkgVers:          make(map[resolve.PackageKey][]resolve.Version),
+		verDeps:          make(map[resolve.VersionKey][]resolve.RequirementVersion),
 	}
 }
 
@@ -46,29 +45,29 @@ func (c *OverrideClient) Version(ctx context.Context, vk resolve.VersionKey) (re
 		}
 	}
 
-	return c.c.Version(ctx, vk)
+	return c.DependencyClient.Version(ctx, vk)
 }
 
 func (c *OverrideClient) Versions(ctx context.Context, pk resolve.PackageKey) ([]resolve.Version, error) {
 	if vers, ok := c.pkgVers[pk]; ok {
 		return vers, nil
 	}
 
-	return c.c.Versions(ctx, pk)
+	return c.DependencyClient.Versions(ctx, pk)
 }
 
 func (c *OverrideClient) Requirements(ctx context.Context, vk resolve.VersionKey) ([]resolve.RequirementVersion, error) {
 	if deps, ok := c.verDeps[vk]; ok {
 		return deps, nil
 	}
 
-	return c.c.Requirements(ctx, vk)
+	return c.DependencyClient.Requirements(ctx, vk)
 }
 
 func (c *OverrideClient) MatchingVersions(ctx context.Context, vk resolve.VersionKey) ([]resolve.Version, error) {
 	if vs, ok := c.pkgVers[vk.PackageKey]; ok {
 		return resolve.MatchRequirement(vk, vs), nil
 	}
 
-	return c.c.MatchingVersions(ctx, vk)
+	return c.DependencyClient.MatchingVersions(ctx, vk)
 }

internal/resolution/dependency_chain.go

@@ -6,6 +6,7 @@ import (
 
 	"deps.dev/util/resolve"
 	"github.com/google/osv-scanner/internal/resolution/manifest"
+	"github.com/google/osv-scanner/internal/resolution/util"
 	vulnUtil "github.com/google/osv-scanner/internal/utility/vulns"
 	"github.com/google/osv-scanner/pkg/lockfile"
 	"github.com/google/osv-scanner/pkg/models"
@@ -28,7 +29,7 @@ func (dc DependencyChain) EndDependency() (resolve.VersionKey, string) {
 
 func ChainIsDev(dc DependencyChain, m manifest.Manifest) bool {
 	direct, _ := dc.DirectDependency()
-	ecosystem, ok := OSVEcosystem[direct.System]
+	ecosystem, ok := util.OSVEcosystem[direct.System]
 	if !ok {
 		return false
 	}
@@ -112,8 +113,8 @@ func chainConstrains(ctx context.Context, cl resolve.Client, chain DependencyCha
 	pkg := lockfile.PackageDetails{
 		Name:      bestVk.Name,
 		Version:   bestVk.Version,
-		Ecosystem: lockfile.Ecosystem(OSVEcosystem[vk.System]),
-		CompareAs: lockfile.Ecosystem(OSVEcosystem[vk.System]),
+		Ecosystem: lockfile.Ecosystem(util.OSVEcosystem[vk.System]),
+		CompareAs: lockfile.Ecosystem(util.OSVEcosystem[vk.System]),
 	}
 
 	return vulnUtil.IsAffected(*vuln, pkg)

internal/resolution/resolve.go

@@ -40,14 +40,15 @@ func getResolver(sys resolve.System, cl resolve.Client) (resolve.Resolver, error
 	}
 }
 
-func Resolve(ctx context.Context, cl resolve.Client, m manifest.Manifest) (*ResolutionResult, error) {
-	c := client.NewOverrideClient(cl)
+func Resolve(ctx context.Context, cl client.ResolutionClient, m manifest.Manifest) (*ResolutionResult, error) {
+	c := client.NewOverrideClient(cl.DependencyClient)
 	c.AddVersion(m.Root, m.Requirements)
 	for _, loc := range m.LocalManifests {
 		c.AddVersion(loc.Root, loc.Requirements)
 		// TODO: may need to do this recursively
 	}
-	r, err := getResolver(m.System(), c)
+	cl.DependencyClient = c
+	r, err := getResolver(m.System(), cl.DependencyClient)
 	if err != nil {
 		return nil, err
 	}
@@ -66,7 +67,7 @@ func Resolve(ctx context.Context, cl resolve.Client, m manifest.Manifest) (*Reso
 		Graph:    graph,
 	}
 
-	if err := result.computeVulns(ctx, c); err != nil {
+	if err := result.computeVulns(ctx, cl); err != nil {
 		return nil, err
 	}
 
@@ -76,9 +77,56 @@ func Resolve(ctx context.Context, cl resolve.Client, m manifest.Manifest) (*Reso
 	return result, nil
 }
 
-var OSVEcosystem = map[resolve.System]models.Ecosystem{
-	resolve.NPM:   models.EcosystemNPM,
-	resolve.Maven: models.EcosystemMaven,
+// computeVulns scans for vulnerabilities in a resolved graph and populates res.Vulns
+func (res *ResolutionResult) computeVulns(ctx context.Context, cl client.ResolutionClient) error {
+	nodeVulns, err := cl.FindVulns(res.Graph)
+	if err != nil {
+		return err
+	}
+	// Find all dependency paths to the vulnerable dependencies
+	var vulnerableNodes []resolve.NodeID
+	vulnInfo := make(map[string]models.Vulnerability)
+	for i, vulns := range nodeVulns {
+		if len(vulns) > 0 {
+			vulnerableNodes = append(vulnerableNodes, resolve.NodeID(i))
+		}
+		for _, vuln := range vulns {
+			vulnInfo[vuln.ID] = vuln
+		}
+	}
+
+	nodeChains := computeChains(res.Graph, vulnerableNodes)
+	vulnChains := make(map[string][]DependencyChain)
+	for i, idx := range vulnerableNodes {
+		for _, vuln := range nodeVulns[idx] {
+			vulnChains[vuln.ID] = append(vulnChains[vuln.ID], nodeChains[i]...)
+		}
+	}
+
+	// construct the ResolutionVulns
+	// TODO: This constructs a single ResolutionVuln per vulnerability ID.
+	// The scan action treats vulns with the same ID but affecting different versions of a package as distinct.
+	// TODO: Combine aliased IDs
+	for id, vuln := range vulnInfo {
+		rv := ResolutionVuln{Vulnerability: vuln, DevOnly: true}
+		for _, chain := range vulnChains[id] {
+			if chainConstrains(ctx, cl, chain, &rv.Vulnerability) {
+				rv.ProblemChains = append(rv.ProblemChains, chain)
+			} else {
+				rv.NonProblemChains = append(rv.NonProblemChains, chain)
+			}
+			rv.DevOnly = rv.DevOnly && ChainIsDev(chain, res.Manifest)
+		}
+		if len(rv.ProblemChains) == 0 {
+			// There has to be at least one problem chain for the vulnerability to appear.
+			// If our heuristic couldn't determine any, treat them all as problematic.
+			rv.ProblemChains = rv.NonProblemChains
+			rv.NonProblemChains = nil
+		}
+		res.Vulns = append(res.Vulns, rv)
+	}
+
+	return nil
 }
 
 // FilterVulns populates Vulns with the UnfilteredVulns that satisfy matchFn