Add proper mirror support (#74)

* Add rewrite support and fix existing tests

* Add unit tests for policy matching

* Compile regexes up front and store policies in map

* Add test for verify flow with mirror

* Rename ImageName -> ResolvedName

And only set it when necessary

* Rename Rewrite -> Replacement

but keep it as rewrite in the yaml

pkg/policy/policy.go

@@ -6,15 +6,13 @@ import (
 	"os"
 	"path"
 	"path/filepath"
-	"slices"
-	"strings"
 
 	"github.com/distribution/reference"
 	"github.com/docker/attest/pkg/config"
 	"github.com/docker/attest/pkg/oci"
 )
 
-func resolveLocalPolicy(opts *PolicyOptions, mapping *config.PolicyMapping) (*Policy, error) {
+func resolveLocalPolicy(opts *PolicyOptions, mapping *config.PolicyMapping, imageName string, matchedName string) (*Policy, error) {
 	if opts.LocalPolicyDir == "" {
 		return nil, fmt.Errorf("local policy dir not set")
 	}
@@ -35,10 +33,13 @@ func resolveLocalPolicy(opts *PolicyOptions, mapping *config.PolicyMapping) (*Po
 		InputFiles: files,
 		Mapping:    mapping,
 	}
+	if imageName != matchedName {
+		policy.ResolvedName = matchedName
+	}
 	return policy, nil
 }
 
-func resolveTufPolicy(opts *PolicyOptions, mapping *config.PolicyMapping) (*Policy, error) {
+func resolveTufPolicy(opts *PolicyOptions, mapping *config.PolicyMapping, imageName string, matchedName string) (*Policy, error) {
 	files := make([]*PolicyFile, 0, len(mapping.Files))
 	for _, f := range mapping.Files {
 		filename := f.Path
@@ -55,32 +56,68 @@ func resolveTufPolicy(opts *PolicyOptions, mapping *config.PolicyMapping) (*Poli
 		InputFiles: files,
 		Mapping:    mapping,
 	}
+	if imageName != matchedName {
+		policy.ResolvedName = matchedName
+	}
 	return policy, nil
 }
 
-func findPolicyMatch(named reference.Named, mappings *config.PolicyMappings) (*config.PolicyMapping, *config.PolicyMirror) {
-	if mappings != nil {
-		for _, mapping := range mappings.Policies {
-			if mapping.Origin.Domain == reference.Domain(named) &&
-				strings.HasPrefix(reference.Path(named), mapping.Origin.Prefix) {
-				return mapping, nil
-			}
-		}
-		// now search mirrors
-		for _, mirror := range mappings.Mirrors {
-			if (slices.Contains(mirror.Mirror.Domains, reference.Domain(named)) ||
-				slices.Contains(mirror.Mirror.Domains, "*")) &&
-				strings.HasPrefix(reference.Path(named), mirror.Mirror.Prefix) {
-				for _, mapping := range mappings.Policies {
-					if mapping.Id == mirror.PolicyId {
-						return mapping, nil
-					}
+type matchType string
+
+const (
+	matchTypePolicy        matchType = "policy"
+	matchTypeMatchNoPolicy matchType = "match_no_policy"
+	matchTypeNoMatch       matchType = "no_match"
+)
+
+type policyMatch struct {
+	matchType   matchType
+	policy      *config.PolicyMapping
+	rule        *config.PolicyRule
+	matchedName string
+}
+
+func findPolicyMatch(imageName string, mappings *config.PolicyMappings) (*policyMatch, error) {
+	if mappings == nil {
+		return &policyMatch{matchType: matchTypeNoMatch, matchedName: imageName}, nil
+	}
+	return findPolicyMatchImpl(imageName, mappings, make(map[*config.PolicyRule]bool))
+}
+
+func findPolicyMatchImpl(imageName string, mappings *config.PolicyMappings, matched map[*config.PolicyRule]bool) (*policyMatch, error) {
+	for _, rule := range mappings.Rules {
+		if rule.Pattern.MatchString(imageName) {
+			switch {
+			case rule.PolicyId == "" && rule.Replacement == "":
+				return nil, fmt.Errorf("rule %s has neither policy-id nor rewrite", rule.Pattern)
+			case rule.PolicyId != "" && rule.Replacement != "":
+				return nil, fmt.Errorf("rule %s has both policy-id and rewrite", rule.Pattern)
+			case rule.PolicyId != "":
+				policy := mappings.Policies[rule.PolicyId]
+				if policy != nil {
+					return &policyMatch{
+						matchType:   matchTypePolicy,
+						policy:      policy,
+						rule:        rule,
+						matchedName: imageName,
+					}, nil
 				}
-				return nil, mirror
+				return &policyMatch{
+					matchType:   matchTypeMatchNoPolicy,
+					rule:        rule,
+					matchedName: imageName,
+				}, nil
+			case rule.Replacement != "":
+				if matched[rule] {
+					return nil, fmt.Errorf("rewrite loop detected")
+				}
+				matched[rule] = true
+				imageName = rule.Pattern.ReplaceAllString(imageName, rule.Replacement)
+				return findPolicyMatchImpl(imageName, mappings, matched)
 			}
 		}
 	}
-	return nil, nil
+	return &policyMatch{matchType: matchTypeNoMatch, matchedName: imageName}, nil
 }
 
 func resolvePolicyById(opts *PolicyOptions) (*Policy, error) {
@@ -90,10 +127,9 @@ func resolvePolicyById(opts *PolicyOptions) (*Policy, error) {
 			return nil, fmt.Errorf("failed to load local policy mappings: %w", err)
 		}
 		if localMappings != nil {
-			for _, mapping := range localMappings.Policies {
-				if mapping.Id == opts.PolicyId {
-					return resolveLocalPolicy(opts, mapping)
-				}
+			policy := localMappings.Policies[opts.PolicyId]
+			if policy != nil {
+				return resolveLocalPolicy(opts, policy, "", "")
 			}
 		}
 
@@ -102,10 +138,9 @@ func resolvePolicyById(opts *PolicyOptions) (*Policy, error) {
 		if err != nil {
 			return nil, fmt.Errorf("failed to load tuf policy mappings by id: %w", err)
 		}
-		for _, mapping := range tufMappings.Policies {
-			if mapping.Id == opts.PolicyId {
-				return resolveTufPolicy(opts, mapping)
-			}
+		policy := tufMappings.Policies[opts.PolicyId]
+		if policy != nil {
+			return resolveTufPolicy(opts, policy, "", "")
 		}
 		return nil, fmt.Errorf("policy with id %s not found", opts.PolicyId)
 	}
@@ -124,7 +159,7 @@ func ResolvePolicy(ctx context.Context, detailsResolver oci.ImageDetailsResolver
 	if err != nil {
 		return nil, fmt.Errorf("failed to get image name: %w", err)
 	}
-	named, err := reference.ParseNormalizedNamed(imageName)
+	imageName, err = normalizeImageName(imageName)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse image name: %w", err)
 	}
@@ -132,9 +167,12 @@ func ResolvePolicy(ctx context.Context, detailsResolver oci.ImageDetailsResolver
 	if err != nil {
 		return nil, fmt.Errorf("failed to load local policy mappings: %w", err)
 	}
-	mapping, mirror := findPolicyMatch(named, localMappings)
-	if mapping != nil {
-		return resolveLocalPolicy(opts, mapping)
+	match, err := findPolicyMatch(imageName, localMappings)
+	if err != nil {
+		return nil, err
+	}
+	if match.matchType == matchTypePolicy {
+		return resolveLocalPolicy(opts, match.policy, imageName, match.matchedName)
 	}
 	// must check tuf
 	tufMappings, err := config.LoadTufMappings(opts.TufClient, opts.LocalTargetsDir)
@@ -143,20 +181,31 @@ func ResolvePolicy(ctx context.Context, detailsResolver oci.ImageDetailsResolver
 	}
 
 	// it's a mirror of a tuf policy
-	if mirror != nil {
+	if match.matchType == matchTypeMatchNoPolicy {
 		for _, mapping := range tufMappings.Policies {
-			if mapping.Id == mirror.PolicyId {
-				return resolveTufPolicy(opts, mapping)
+			if mapping.Id == match.rule.PolicyId {
+				return resolveTufPolicy(opts, mapping, imageName, match.matchedName)
 			}
 		}
 	}
 
 	// try to resolve a tuf policy directly
-	mapping, _ = findPolicyMatch(named, tufMappings)
-	if mapping == nil {
-		return nil, nil
+	match, err = findPolicyMatch(imageName, tufMappings)
+	if err != nil {
+		return nil, err
 	}
-	return resolveTufPolicy(opts, mapping)
+	if match.matchType == matchTypePolicy {
+		return resolveTufPolicy(opts, match.policy, imageName, match.matchedName)
+	}
+	return nil, nil
+}
+
+func normalizeImageName(imageName string) (string, error) {
+	named, err := reference.ParseNormalizedNamed(imageName)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse image name: %w", err)
+	}
+	return named.Name(), nil
 }
 
 func CreateImageDetailsResolver(imageSource *oci.ImageSpec) (oci.ImageDetailsResolver, error) {

pkg/policy/policy_match_test.go

@@ -0,0 +1,121 @@
+package policy
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/docker/attest/pkg/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFindPolicyMatch(t *testing.T) {
+	testCases := []struct {
+		name       string
+		imageName  string
+		mappingDir string
+
+		expectError       bool
+		expectedMatchType matchType
+		expectedPolicyID  string
+		expectedImageName string
+	}{
+		{
+			name:       "alpine",
+			mappingDir: "doi",
+			imageName:  "docker.io/library/alpine",
+
+			expectedMatchType: matchTypePolicy,
+			expectedPolicyID:  "docker-official-images",
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "no match",
+			mappingDir: "doi",
+			imageName:  "docker.io/something/else",
+
+			expectedMatchType: matchTypeNoMatch,
+			expectedImageName: "docker.io/something/else",
+		},
+		{
+			name:       "match, no policy",
+			mappingDir: "local",
+			imageName:  "docker.io/library/alpine",
+
+			expectedMatchType: matchTypeMatchNoPolicy,
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "simple rewrite",
+			mappingDir: "simple-rewrite",
+			imageName:  "mycoolmirror.org/library/alpine",
+
+			expectedMatchType: matchTypePolicy,
+			expectedPolicyID:  "docker-official-images",
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "rewrite no match",
+			mappingDir: "rewrite-to-no-match",
+			imageName:  "mycoolmirror.org/library/alpine",
+
+			expectedMatchType: matchTypeNoMatch,
+			expectedImageName: "badredirect.org/alpine",
+		},
+		{
+			name:       "rewrite to match, no policy",
+			mappingDir: "rewrite-to-local",
+			imageName:  "mycoolmirror.org/library/alpine",
+
+			expectedMatchType: matchTypeMatchNoPolicy,
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "multiple rewrites",
+			mappingDir: "rewrite-multiple",
+			imageName:  "myevencoolermirror.org/library/alpine",
+
+			expectedMatchType: matchTypePolicy,
+			expectedPolicyID:  "docker-official-images",
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "invalid rewrites",
+			mappingDir: "rewrite-invalid",
+			imageName:  "mycoolmirror.org/library/alpine",
+
+			expectError:       true,
+			expectedMatchType: matchTypePolicy,
+			expectedPolicyID:  "docker-official-images",
+			expectedImageName: "docker.io/library/alpine",
+		},
+		{
+			name:       "rewrite loop",
+			mappingDir: "rewrite-loop",
+			imageName:  "yin/alpine",
+
+			expectError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mappings, err := config.LoadLocalMappings(filepath.Join("testdata", "mappings", tc.mappingDir))
+			require.NoError(t, err)
+			match, err := findPolicyMatch(tc.imageName, mappings)
+			if tc.expectError {
+				require.Error(t, err)
+				// TODO: check error matches expected error message
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectedMatchType, match.matchType)
+			if match.matchType == matchTypePolicy {
+				if assert.NotNil(t, match.policy) {
+					assert.Equal(t, tc.expectedPolicyID, match.policy.Id)
+				}
+			}
+			assert.Equal(t, tc.expectedImageName, match.matchedName)
+		})
+	}
+}

pkg/policy/policy_test.go

@@ -84,6 +84,7 @@ func TestRegoEvaluator_Evaluate(t *testing.T) {
 			src, err := oci.ParseImageSpec(imageName, oci.WithPlatform(platform.String()))
 			require.NoError(t, err)
 			resolver, err := policy.CreateImageDetailsResolver(src)
+			require.NoError(t, err)
 			policy, err := policy.ResolvePolicy(ctx, resolver, tc.policy)
 			if tc.errorStr != "" {
 				require.Error(t, err)
@@ -91,6 +92,7 @@ func TestRegoEvaluator_Evaluate(t *testing.T) {
 				return
 			}
 			require.NoErrorf(t, err, "failed to resolve policy")
+			require.NotNil(t, policy, "policy should not be nil")
 			result, err := re.Evaluate(ctx, tc.resolver, policy, input)
 			require.NoErrorf(t, err, "Evaluate failed")
 
@@ -107,8 +109,10 @@ func TestRegoEvaluator_Evaluate(t *testing.T) {
 func TestLoadingMappings(t *testing.T) {
 	policyMappings, err := config.LoadLocalMappings(filepath.Join("testdata", "mock-tuf-allow"))
 	require.NoError(t, err)
-	assert.Equal(t, len(policyMappings.Mirrors), 1)
-	for _, mirror := range policyMappings.Mirrors {
-		assert.Equal(t, "docker-official-images", mirror.PolicyId)
+	assert.Equal(t, len(policyMappings.Rules), 3)
+	for _, mirror := range policyMappings.Rules {
+		if mirror.PolicyId != "" {
+			assert.Equal(t, "docker-official-images", mirror.PolicyId)
+		}
 	}
 }

pkg/policy/testdata/mappings/doi/mapping.yaml

@@ -0,0 +1,10 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images

pkg/policy/testdata/mappings/local/mapping.yaml

@@ -0,0 +1,10 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: local-policy
+    description: Local Policy
+    files:
+      - path: local-policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images # note this policy does not exist in this file

pkg/policy/testdata/mappings/rewrite-invalid/mapping.yaml

@@ -0,0 +1,13 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: "^mycoolmirror[.]org/library/(.*)$"
+    rewrite: "docker.io/library/$1"
+    policy-id: docker-official-images # invalid to specify both rewrite and policy-id

pkg/policy/testdata/mappings/rewrite-loop/mapping.yaml

@@ -0,0 +1,14 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: "^yin/(.*)$"
+    rewrite: "yang/$1"
+  - pattern: "^yang/(.*)$"
+    rewrite: "yin/$1"

pkg/policy/testdata/mappings/rewrite-multiple/mapping.yaml

@@ -0,0 +1,14 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: "^mycoolmirror[.]org/library/(.*)$"
+    rewrite: "docker.io/library/$1"
+  - pattern: "^myevencoolermirror[.]org/library/(.*)$"
+    rewrite: "mycoolmirror.org/library/$1"

pkg/policy/testdata/mappings/rewrite-to-local/mapping.yaml

@@ -0,0 +1,12 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: local-policy
+    description: Local Policy
+    files:
+      - path: local-policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images # note this policy does not exist in this file
+  - pattern: "^mycoolmirror[.]org/library/(.*)$"
+    rewrite: "docker.io/library/$1"

pkg/policy/testdata/mappings/rewrite-to-no-match/mapping.yaml

@@ -0,0 +1,12 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: "^mycoolmirror[.]org/library/(.*)$"
+    rewrite: "badredirect.org/$1" # no matching rule for this rewrite

pkg/policy/testdata/mappings/simple-rewrite/mapping.yaml

@@ -0,0 +1,12 @@
+version: v1
+kind: policy-mapping
+policies:
+  - id: docker-official-images
+    description: Docker Official Images
+    files:
+      - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: "^mycoolmirror[.]org/library/(.*)$"
+    rewrite: "docker.io/library/$1"

pkg/policy/testdata/mock-tuf-allow-canonical/mapping.yaml

@@ -2,15 +2,16 @@
 version: v1
 kind: policy-mapping
 policies:
-  - origin:
-      domain: docker.io
-      prefix: library/
-    id: docker-official-images
+  - id: docker-official-images
     description: Docker Official Images
+    attestations:
+      repo: "localhost:5001/library-refs"
     files:
       - path: doi/policy.rego
-mirrors:
-  - policy-id: docker-official-images
-    mirror:
-      domains: [localhost:5001, registry.local:5000]
-      prefix: ""
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: ^localhost:5001/(.*)$
+    rewrite: docker.io/library/$1
+  - pattern: ^registry[.]local:5000/(.*)$
+    rewrite: docker.io/library/$1

pkg/policy/testdata/mock-tuf-allow/mapping.yaml

@@ -2,17 +2,16 @@
 version: v1
 kind: policy-mapping
 policies:
-  - origin:
-      domain: docker.io
-      prefix: library/
-    id: docker-official-images
+  - id: docker-official-images
     description: Docker Official Images
     attestations:
       repo: "localhost:5001/library-refs"
     files:
       - path: doi/policy.rego
-mirrors:
-  - policy-id: docker-official-images
-    mirror:
-      domains: [localhost:5001, registry.local:5000]
-      prefix: ""
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images
+  - pattern: ^localhost:5001/(.*)$
+    rewrite: docker.io/library/$1
+  - pattern: ^registry[.]local:5000/(.*)$
+    rewrite: docker.io/library/$1

pkg/policy/testdata/mock-tuf-deny/mapping.yaml

@@ -2,10 +2,10 @@
 version: v1
 kind: policy-mapping
 policies:
-  - origin:
-      domain: docker.io
-      prefix: library/
-    id: docker-official-images
+  - id: docker-official-images
     description: Docker Official Images
     files:
       - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images

pkg/policy/testdata/mock-tuf-verify-sig/mapping.yaml

@@ -2,10 +2,10 @@
 version: v1
 kind: policy-mapping
 policies:
-  - origin:
-      domain: docker.io
-      prefix: library/
-    id: docker-official-images
+  - id: docker-official-images
     description: Docker Official Images
     files:
       - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images

pkg/policy/testdata/mock-tuf-wrong-key/mapping.yaml

@@ -2,10 +2,10 @@
 version: v1
 kind: policy-mapping
 policies:
-  - origin:
-      domain: docker.io
-      prefix: library/
-    id: docker-official-images
+  - id: docker-official-images
     description: Docker Official Images
     files:
       - path: doi/policy.rego
+rules:
+  - pattern: "^docker[.]io/library/(.*)$"
+    policy-id: docker-official-images

pkg/policy/types.go

@@ -36,9 +36,10 @@ type PolicyOptions struct {
 }
 
 type Policy struct {
-	InputFiles []*PolicyFile
-	Query      string
-	Mapping    *config.PolicyMapping
+	InputFiles   []*PolicyFile
+	Query        string
+	Mapping      *config.PolicyMapping
+	ResolvedName string
 }
 
 type PolicyInput struct {