From 8d5b7c523d1785855cf79be0965f0c40c33552e5 Mon Sep 17 00:00:00 2001
From: Dominik Richter <dominik.richter@gmail.com>
Date: Thu, 4 Apr 2024 02:33:39 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix=20risk=20factors=20missing?=
 =?UTF-8?q?=20in=20inmemory=20w/=20upstream=20(#1224)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When running against upstream policies, we don't get risk factors from
the policy bundle, since we generally don't need the bundle to run the
scan. However, we do require basic risk factor info to help score
everything before we send data up.

In this change we pull risk info from the resolved policy and inject it
into the inmemory datastore before the scan is started (unless the risk
factor exists for any reason, like prior policy bundles). It has enough
information for scoring.

For any later reporting step we may still inject all risk factor
metadata into the inmemory store before printing the output.

Signed-off-by: Dominik Richter <dominik.richter@gmail.com>
---
 internal/datalakes/inmemory/policyhub.go      |  8 ++++++++
 internal/datalakes/inmemory/policyresolver.go | 13 +++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/internal/datalakes/inmemory/policyhub.go b/internal/datalakes/inmemory/policyhub.go
index 5778db0f..216d0446 100644
--- a/internal/datalakes/inmemory/policyhub.go
+++ b/internal/datalakes/inmemory/policyhub.go
@@ -316,6 +316,14 @@ func (db *Db) invalidateFrameworkAndBundleAncestors(ctx context.Context, wrap *w
 	return nil
 }
 
+func (db *Db) getRiskFactor(ctx context.Context, mrn string) (*policy.RiskFactor, error) {
+	found, ok := db.cache.Get(dbIDRiskFactor + mrn)
+	if !ok {
+		return nil, errors.New("risk factor " + mrn + " not found")
+	}
+	return found.(*policy.RiskFactor), nil
+}
+
 func (db *Db) SetRiskFactor(ctx context.Context, riskFactor *policy.RiskFactor) error {
 	db.cache.Set(dbIDRiskFactor+riskFactor.Mrn, riskFactor, 1)
 	return nil
diff --git a/internal/datalakes/inmemory/policyresolver.go b/internal/datalakes/inmemory/policyresolver.go
index e4b4e044..f8e6c44b 100644
--- a/internal/datalakes/inmemory/policyresolver.go
+++ b/internal/datalakes/inmemory/policyresolver.go
@@ -615,6 +615,19 @@ func (db *Db) SetResolvedPolicy(ctx context.Context, mrn string, resolvedPolicy
 		return errors.New("failed to save resolved policy '" + mrn + "'")
 	}
 
+	// We need the risk factors for initial reporting, but don't require all
+	// their metadata. The risk factors in the resolved policy provides everything
+	// we need for scoring. If we fetch a full risk factor we can replace it.
+	for mrn, rf := range resolvedPolicy.CollectorJob.RiskFactors {
+		rf.Mrn = mrn
+		if _, err := db.getRiskFactor(ctx, rf.Mrn); err == nil {
+			continue
+		}
+		if err := db.SetRiskFactor(ctx, rf); err != nil {
+			return err
+		}
+	}
+
 	if cached {
 		x, ok := db.cache.Get(dbIDPolicy + mrn)
 		if !ok {