From 68a05a17543ad7062a80cbe5f546696437c74140 Mon Sep 17 00:00:00 2001
From: Jay Mundrawala <jay@mondoo.com>
Date: Mon, 1 Jul 2024 10:40:11 -0500
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Don't=20filter=20out=20scores=20?=
 =?UTF-8?q?of=20queries=20used=20for=20risk=20factors=20and=20not?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

If a query was used for a risk factor, its score was always being
filtered out, regardless of if it might have been used outside of risk
factors
---
 policy/executor/internal/collector.go | 38 ++++++++++++++++-----------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/policy/executor/internal/collector.go b/policy/executor/internal/collector.go
index 23e860ed..04eb534b 100644
--- a/policy/executor/internal/collector.go
+++ b/policy/executor/internal/collector.go
@@ -44,6 +44,7 @@ type BufferedCollector struct {
 	collector      *PolicyServiceCollector
 	resolvedPolicy *policy.ResolvedPolicy
 	riskMRNs       map[string][]string
+	keepQrIds      map[string]bool
 	duration       time.Duration
 	stopChan       chan struct{}
 	wg             sync.WaitGroup
@@ -55,29 +56,34 @@ func WithResolvedPolicy(resolved *policy.ResolvedPolicy) (BufferedCollectorOpt,
 	// TODO: need a more native way to integrate this part. We don't want to
 	// introduce a score type.
 	riskMRNs := map[string][]string{}
+	keepQrIds := map[string]bool{}
 	for _, rj := range resolved.CollectorJob.ReportingJobs {
-		if rj.Type != policy.ReportingJob_RISK_FACTOR {
-			continue
-		}
+		if rj.Type == policy.ReportingJob_RISK_FACTOR {
+			for k := range rj.ChildJobs {
+				cjob := resolved.CollectorJob.ReportingJobs[k]
+				if resolved.CollectorJob.RiskMrns == nil {
+					return nil, errors.New("missing query MRNs in resolved policy")
+				}
 
-		for k := range rj.ChildJobs {
-			cjob := resolved.CollectorJob.ReportingJobs[k]
-			if resolved.CollectorJob.RiskMrns == nil {
-				return nil, errors.New("missing query MRNs in resolved policy")
-			}
+				mrns := resolved.CollectorJob.RiskMrns[cjob.Uuid]
+				if mrns == nil {
+					return nil, errors.New("missing query MRNs for job uuid=" + cjob.Uuid + " checksum=" + cjob.Checksum)
+				}
 
-			mrns := resolved.CollectorJob.RiskMrns[cjob.Uuid]
-			if mrns == nil {
-				return nil, errors.New("missing query MRNs for job uuid=" + cjob.Uuid + " checksum=" + cjob.Checksum)
+				riskMRNs[cjob.QrId] = append(riskMRNs[cjob.QrId], mrns.Items...)
+			}
+		} else {
+			for k := range rj.ChildJobs {
+				cjob := resolved.CollectorJob.ReportingJobs[k]
+				keepQrIds[cjob.QrId] = true
 			}
-
-			riskMRNs[cjob.QrId] = append(riskMRNs[cjob.QrId], mrns.Items...)
 		}
 	}
 
 	return func(b *BufferedCollector) {
 		b.resolvedPolicy = resolved
 		b.riskMRNs = riskMRNs
+		b.keepQrIds = keepQrIds
 	}, nil
 }
 
@@ -132,10 +138,10 @@ func (c *BufferedCollector) run() {
 			}
 
 			for _, s := range c.scores {
-				if c.consumeRisk(s, risksIdx) {
-					continue
+				c.consumeRisk(s, risksIdx)
+				if c.keepQrIds[s.QrId] {
+					scores = append(scores, s)
 				}
-				scores = append(scores, s)
 			}
 			for k := range c.scores {
 				delete(c.scores, k)