Context instead of global state for download cache

Makes use of Context to store the download cache, this will help running
benchmarks in parallel as they perform within the same process and share
the same global state.

cmd/fetch/fetch_policy.go

@@ -115,8 +115,10 @@ func fetchPolicyCmd() *cobra.Command {
 				sources = append(sources, &source.PolicyUrl{Url: url, Kind: source.DataKind})
 			}
 
+			ctx := cmd.Context()
 			for _, s := range sources {
-				_, err := s.GetPolicy(cmd.Context(), destDir, true)
+				c, _, err := s.GetPolicy(ctx, destDir, true)
+				ctx = c
 				if err != nil {
 					return err
 				}

cmd/inspect/inspect_policy.go

@@ -121,7 +121,8 @@ func inspectPolicyCmd() *cobra.Command {
 				s := &source.PolicyUrl{Url: url, Kind: source.PolicyKind}
 
 				// Download
-				policyDir, err := s.GetPolicy(ctx, destDir, false)
+				c, policyDir, err := s.GetPolicy(ctx, destDir, false)
+				ctx = c
 				if err != nil {
 					return err
 				}

cmd/inspect/inspect_policy_data.go

@@ -91,7 +91,8 @@ func inspectPolicyDataCmd() *cobra.Command {
 				s := &source.PolicyUrl{Url: url, Kind: source.PolicyKind}
 
 				// Download
-				policyDir, err := s.GetPolicy(ctx, destDir, false)
+				c, policyDir, err := s.GetPolicy(ctx, destDir, false)
+				ctx = c
 				if err != nil {
 					return err
 				}

cmd/validate/common_test.go

@@ -41,10 +41,10 @@ type mockEvaluator struct {
 	mock.Mock
 }
 
-func (e *mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) ([]evaluator.Outcome, evaluator.Data, error) {
+func (e *mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) (context.Context, []evaluator.Outcome, evaluator.Data, error) {
 	args := e.Called(ctx, target.Inputs)
 
-	return args.Get(0).([]evaluator.Outcome), args.Get(1).(evaluator.Data), args.Error(2)
+	return ctx, args.Get(0).([]evaluator.Outcome), args.Get(1).(evaluator.Data), args.Error(2)
 }
 
 func (e *mockEvaluator) Destroy() {

cmd/validate/image.go

@@ -209,7 +209,8 @@ func validateImageCmd(validate imageValidationFunc) *cobra.Command {
 				data.spec = s
 			}
 
-			policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			c, policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			ctx = c
 			if err != nil {
 				allErrors = errors.Join(allErrors, err)
 				return
@@ -232,9 +233,10 @@ func validateImageCmd(validate imageValidationFunc) *cobra.Command {
 
 			// We're not currently using the policyCache returned from PreProcessPolicy, but we could
 			// use it to cache the policy for future use.
-			if p, _, err := policy.PreProcessPolicy(ctx, policyOptions); err != nil {
+			if c, p, _, err := policy.PreProcessPolicy(ctx, policyOptions); err != nil {
 				allErrors = errors.Join(allErrors, err)
 			} else {
+				ctx = c
 				// inject extra variables into rule data per source
 				if len(data.extraRuleData) > 0 {
 					policySpec := p.Spec()
@@ -262,7 +264,8 @@ func validateImageCmd(validate imageValidationFunc) *cobra.Command {
 							if len(parts) < 2 {
 								log.Errorf("Incorrect syntax for --extra-rule-data")
 							}
-							extraRuleDataPolicyConfig, err := validate_utils.GetPolicyConfig(ctx, parts[1])
+							c, extraRuleDataPolicyConfig, err := validate_utils.GetPolicyConfig(ctx, parts[1])
+							ctx = c
 							if err != nil {
 								log.Errorf("Unable to load data from extraRuleData: %s", err.Error())
 							}

cmd/validate/image_integration_test.go

@@ -71,7 +71,7 @@ func TestEvaluatorLifecycle(t *testing.T) {
 		evaluators[i].On("Destroy").NotBefore(expectations...)
 	}
 
-	newConftestEvaluator = func(_ context.Context, s []source.PolicySource, _ evaluator.ConfigProvider, _ v1alpha1.Source) (evaluator.Evaluator, error) {
+	newConftestEvaluator = func(ctx context.Context, s []source.PolicySource, _ evaluator.ConfigProvider, _ v1alpha1.Source) (evaluator.Evaluator, error) {
 		// We are splitting this url to get to the index of the evaluator.
 		idx, err := strconv.Atoi(strings.Split(strings.Split(s[0].PolicyUrl(), "@")[0], "::")[1])
 		require.NoError(t, err)
@@ -84,7 +84,7 @@ func TestEvaluatorLifecycle(t *testing.T) {
 
 	validate := func(_ context.Context, component app.SnapshotComponent, _ *app.SnapshotSpec, _ policy.Policy, evaluators []evaluator.Evaluator, _ bool) (*output.Output, error) {
 		for _, e := range evaluators {
-			_, _, err := e.Evaluate(ctx, evaluator.EvaluationTarget{Inputs: []string{}})
+			_, _, _, err := e.Evaluate(ctx, evaluator.EvaluationTarget{Inputs: []string{}})
 			require.NoError(t, err)
 		}
 

cmd/validate/input.go

@@ -88,7 +88,8 @@ func validateInputCmd(validate InputValidationFunc) *cobra.Command {
 		PreRunE: func(cmd *cobra.Command, args []string) (allErrors error) {
 			ctx := cmd.Context()
 
-			policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			c, policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			cmd.SetContext(c)
 			if err != nil {
 				allErrors = errors.Join(allErrors, err)
 				return

cmd/validate/policy.go

@@ -53,7 +53,8 @@ func ValidatePolicyCmd(validate policyValidationFunc) *cobra.Command {
 		PreRunE: func(cmd *cobra.Command, args []string) (allErrors error) {
 			ctx := cmd.Context()
 
-			policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			c, policyConfiguration, err := validate_utils.GetPolicyConfig(ctx, data.policyConfiguration)
+			cmd.SetContext(c)
 			if err != nil {
 				allErrors = errors.Join(allErrors, err)
 				return

internal/evaluator/conftest_evaluator.go

@@ -364,7 +364,7 @@ func (r *policyRules) collect(a *ast.AnnotationsRef) error {
 	return nil
 }
 
-func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget) ([]Outcome, Data, error) {
+func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget) (context.Context, []Outcome, Data, error) {
 	var results []Outcome
 
 	if trace.IsEnabled() {
@@ -379,11 +379,11 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 	rules := policyRules{}
 	// Download all sources
 	for _, s := range c.policySources {
-		dir, err := s.GetPolicy(ctx, c.workDir, false)
+		ctx, dir, err := s.GetPolicy(ctx, c.workDir, false)
 		if err != nil {
 			log.Debugf("Unable to download source from %s!", s.PolicyUrl())
 			// TODO do we want to download other policies instead of erroring out?
-			return nil, nil, err
+			return ctx, nil, nil, err
 		}
 		annotations := []*ast.AnnotationsRef{}
 		fs := utils.FS(ctx)
@@ -396,7 +396,7 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 					// Let's try to give some more robust messaging to the user.
 					policyURL, err := url.Parse(s.PolicyUrl())
 					if err != nil {
-						return nil, nil, errMsg
+						return ctx, nil, nil, errMsg
 					}
 					// Do we have a prefix at the end of the URL path?
 					// If not, this means we aren't trying to access a specific file.
@@ -410,7 +410,7 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 						}
 					}
 				}
-				return nil, nil, errMsg
+				return ctx, nil, nil, errMsg
 			}
 		}
 
@@ -419,7 +419,7 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 				continue
 			}
 			if err := rules.collect(a); err != nil {
-				return nil, nil, err
+				return ctx, nil, nil, err
 			}
 		}
 	}
@@ -453,7 +453,7 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 	runResults, data, err := r.Run(ctx, target.Inputs)
 	if err != nil {
 		// TODO do we want to evaluate further policies instead of erroring out?
-		return nil, nil, err
+		return ctx, nil, nil, err
 	}
 
 	effectiveTime := c.policy.EffectiveTime()
@@ -535,10 +535,10 @@ func (c conftestEvaluator) Evaluate(ctx context.Context, target EvaluationTarget
 	// ran due to input error, etc.
 	if totalRules == 0 {
 		log.Error("no successes, warnings, or failures, check input")
-		return nil, nil, fmt.Errorf("no successes, warnings, or failures, check input")
+		return ctx, nil, nil, fmt.Errorf("no successes, warnings, or failures, check input")
 	}
 
-	return results, data, nil
+	return ctx, results, data, nil
 }
 
 func toRules(results []output.Result) []Result {

internal/evaluator/conftest_evaluator_test.go

@@ -67,8 +67,8 @@ func withTestRunner(ctx context.Context, clnt testRunner) context.Context {
 
 type testPolicySource struct{}
 
-func (t testPolicySource) GetPolicy(ctx context.Context, dest string, showMsg bool) (string, error) {
-	return "/policy", nil
+func (t testPolicySource) GetPolicy(ctx context.Context, dest string, showMsg bool) (context.Context, string, error) {
+	return ctx, "/policy", nil
 }
 
 func (t testPolicySource) PolicyUrl() string {
@@ -302,7 +302,7 @@ func TestConftestEvaluatorEvaluateSeverity(t *testing.T) {
 	}, pol, ecc.Source{})
 
 	assert.NoError(t, err)
-	actualResults, data, err := evaluator.Evaluate(ctx, inputs)
+	_, actualResults, data, err := evaluator.Evaluate(ctx, inputs)
 	assert.NoError(t, err)
 	assert.Equal(t, expectedResults, actualResults)
 	assert.Equal(t, expectedData, data)
@@ -413,7 +413,7 @@ func TestConftestEvaluatorEvaluateNoSuccessWarningsOrFailures(t *testing.T) {
 			}, p, ecc.Source{Config: tt.sourceConfig})
 
 			assert.NoError(t, err)
-			actualResults, data, err := evaluator.Evaluate(ctx, inputs)
+			_, actualResults, data, err := evaluator.Evaluate(ctx, inputs)
 			assert.ErrorContains(t, err, "no successes, warnings, or failures, check input")
 			assert.Nil(t, actualResults)
 			assert.Nil(t, data)
@@ -1305,7 +1305,7 @@ func TestConftestEvaluatorIncludeExclude(t *testing.T) {
 			}, p, ecc.Source{})
 
 			assert.NoError(t, err)
-			got, data, err := evaluator.Evaluate(ctx, inputs)
+			_, got, data, err := evaluator.Evaluate(ctx, inputs)
 			assert.NoError(t, err)
 			assert.Equal(t, tt.want, got)
 			assert.Equal(t, Data(nil), data)
@@ -1825,7 +1825,7 @@ func TestConftestEvaluatorEvaluate(t *testing.T) {
 	}, config, ecc.Source{})
 	require.NoError(t, err)
 
-	results, data, err := evaluator.Evaluate(ctx, EvaluationTarget{Inputs: []string{path.Join(dir, "inputs")}})
+	_, results, data, err := evaluator.Evaluate(ctx, EvaluationTarget{Inputs: []string{path.Join(dir, "inputs")}})
 	require.NoError(t, err)
 
 	// sort the slice by code for test stability
@@ -1888,7 +1888,7 @@ func TestUnconformingRule(t *testing.T) {
 	}, p, ecc.Source{})
 	require.NoError(t, err)
 
-	_, _, err = evaluator.Evaluate(ctx, EvaluationTarget{Inputs: []string{path.Join(dir, "inputs")}})
+	_, _, _, err = evaluator.Evaluate(ctx, EvaluationTarget{Inputs: []string{path.Join(dir, "inputs")}})
 	assert.EqualError(t, err, `the rule "deny = true { true }" returns an unsupported value, at no_msg.rego:3`)
 }
 

internal/evaluator/evaluator.go

@@ -26,7 +26,7 @@ type EvaluationTarget struct {
 }
 
 type Evaluator interface {
-	Evaluate(ctx context.Context, target EvaluationTarget) ([]Outcome, Data, error)
+	Evaluate(ctx context.Context, target EvaluationTarget) (context.Context, []Outcome, Data, error)
 
 	// Destroy performs any cleanup needed
 	Destroy()

internal/image/validate.go

@@ -122,7 +122,8 @@ func ValidateImage(ctx context.Context, comp app.SnapshotComponent, snap *app.Sn
 		} else {
 			target.Target = digest
 		}
-		results, data, err := e.Evaluate(ctx, target)
+		c, results, data, err := e.Evaluate(ctx, target)
+		ctx = c
 		log.Debug("\n\nRunning conftest policy check\n\n")
 
 		if err != nil {

internal/image/validate_test.go

@@ -289,10 +289,10 @@ type mockEvaluator struct {
 	mock.Mock
 }
 
-func (e *mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) ([]evaluator.Outcome, evaluator.Data, error) {
+func (e *mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) (context.Context, []evaluator.Outcome, evaluator.Data, error) {
 	args := e.Called(ctx, target.Inputs)
 
-	return args.Get(0).([]evaluator.Outcome), args.Get(1).(evaluator.Data), args.Error(2)
+	return ctx, args.Get(0).([]evaluator.Outcome), args.Get(1).(evaluator.Data), args.Error(2)
 }
 
 func (e *mockEvaluator) Destroy() {

internal/input/validate.go

@@ -55,7 +55,8 @@ func ValidateInput(ctx context.Context, fpath string, policy policy.Policy, deta
 
 	var allResults []evaluator.Outcome
 	for _, e := range p.Evaluators {
-		results, _, err := e.Evaluate(ctx, evaluator.EvaluationTarget{Inputs: inputFiles})
+		c, results, _, err := e.Evaluate(ctx, evaluator.EvaluationTarget{Inputs: inputFiles})
+		ctx = c
 		if err != nil {
 			return nil, fmt.Errorf("evaluating policy: %w", err)
 		}

internal/input/validate_test.go

@@ -40,8 +40,8 @@ type (
 	badMockEvaluator struct{}
 )
 
-func (e mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) ([]evaluator.Outcome, evaluator.Data, error) {
-	return []evaluator.Outcome{}, nil, nil
+func (e mockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) (context.Context, []evaluator.Outcome, evaluator.Data, error) {
+	return ctx, []evaluator.Outcome{}, nil, nil
 }
 
 func (e mockEvaluator) Destroy() {
@@ -51,8 +51,8 @@ func (e mockEvaluator) CapabilitiesPath() string {
 	return ""
 }
 
-func (b badMockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) ([]evaluator.Outcome, evaluator.Data, error) {
-	return nil, nil, errors.New("Evaluator error")
+func (b badMockEvaluator) Evaluate(ctx context.Context, target evaluator.EvaluationTarget) (context.Context, []evaluator.Outcome, evaluator.Data, error) {
+	return ctx, nil, nil, errors.New("Evaluator error")
 }
 
 func (e badMockEvaluator) Destroy() {

internal/policy/policy.go

@@ -567,17 +567,17 @@ func validatePolicyConfig(policyConfig string) error {
 
 // PreProcessPolicy fetches policy sources and returns a policy object with
 // pinned SHA/image digest URL where applicable, along with a policy cache object.
-func PreProcessPolicy(ctx context.Context, policyOptions Options) (Policy, *cache.PolicyCache, error) {
+func PreProcessPolicy(ctx context.Context, policyOptions Options) (context.Context, Policy, *cache.PolicyCache, error) {
 	var policyCache *cache.PolicyCache
 	pinnedPolicyUrls := map[string][]string{}
 	policyCache, err := cache.NewPolicyCache(ctx)
 	if err != nil {
-		return nil, nil, err
+		return ctx, nil, nil, err
 	}
 
 	p, err := NewPolicy(ctx, policyOptions)
 	if err != nil {
-		return nil, nil, err
+		return ctx, nil, nil, err
 	}
 
 	sources := p.Spec().Sources
@@ -589,18 +589,18 @@ func PreProcessPolicy(ctx context.Context, policyOptions Options) (Policy, *cach
 		dir, err := utils.CreateWorkDir(fs)
 		if err != nil {
 			log.Debug("Failed to create work dir!")
-			return nil, nil, err
+			return ctx, nil, nil, err
 		}
 
 		for _, policySource := range policySources {
 			if strings.HasPrefix(policySource.PolicyUrl(), "data:") {
 				continue
 			}
 
-			destDir, err := policySource.GetPolicy(ctx, dir, false)
+			ctx, destDir, err := policySource.GetPolicy(ctx, dir, false)
 			if err != nil {
 				log.Debugf("Unable to download source from %s!", policySource.PolicyUrl())
-				return nil, nil, err
+				return ctx, nil, nil, err
 			}
 			log.Debugf("Downloaded policy source from %s to %s\n", policySource.PolicyUrl(), destDir)
 
@@ -626,7 +626,7 @@ func PreProcessPolicy(ctx context.Context, policyOptions Options) (Policy, *cach
 		}
 	}
 
-	return p, policyCache, err
+	return ctx, p, policyCache, err
 }
 
 func urls(s []source.PolicySource, kind source.PolicyType) []string {

internal/policy/source/git_config.go

@@ -54,25 +54,25 @@ func SourceIsHttp(src string) bool {
 	return err == nil && strings.HasPrefix(normalizedUrl, "http")
 }
 
-func GoGetterDownload(ctx context.Context, tmpDir, src string) (string, error) {
+func GoGetterDownload(ctx context.Context, tmpDir, src string) (context.Context, string, error) {
 	// Download the config from a url
 	c := PolicyUrl{
 		Url:  src,
 		Kind: ConfigKind,
 	}
-	configDir, err := c.GetPolicy(ctx, tmpDir, false)
+	ctx, configDir, err := c.GetPolicy(ctx, tmpDir, false)
 	if err != nil {
 		log.Debugf("Failed to download policy config from %s", c.Url)
-		return "", err
+		return ctx, "", err
 	}
 	log.Debugf("Downloaded policy config from %s to %s", c.Url, configDir)
 
 	// Look for a suitable file to use for the config
 	configFile, err := choosePolicyFile(ctx, configDir)
 	if err != nil {
 		// A more useful error message:
-		return "", fmt.Errorf("no suitable config file found at %s", c.Url)
+		return ctx, "", fmt.Errorf("no suitable config file found at %s", c.Url)
 	}
 	log.Debugf("Chose file %s to use for the policy config", configFile)
-	return configFile, nil
+	return ctx, configFile, nil
 }