Merge pull request #2568 from dheerajodha/EC-1312

improve OCI ref Parsing and Digest Resolution

Author: joejstuart

internal/rego/oci/__snapshots__/oci_test.snap

@@ -2215,3 +2215,81 @@
  ]
 }
 ---
+
+[TestOCIDescriptorManifest/tag-based_URI - 1]
+{
+ "type": "object",
+ "value": [
+  [
+   {
+    "type": "string",
+    "value": "annotations"
+   },
+   {
+    "type": "object",
+    "value": []
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "artifactType"
+   },
+   {
+    "type": "string",
+    "value": ""
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "data"
+   },
+   {
+    "type": "string",
+    "value": ""
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "digest"
+   },
+   {
+    "type": "string",
+    "value": "sha256:4e388ab32b10dc8dbc7e28144f552830adc74787c1e2c0824032078a79f227fb"
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "mediaType"
+   },
+   {
+    "type": "string",
+    "value": "application/vnd.oci.image.manifest.v1+json"
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "size"
+   },
+   {
+    "type": "number",
+    "value": 123
+   }
+  ],
+  [
+   {
+    "type": "string",
+    "value": "urls"
+   },
+   {
+    "type": "array",
+    "value": []
+   }
+  ]
+ ]
+}
+---

internal/rego/oci/oci.go

@@ -24,9 +24,9 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
-	"strings"
 
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -38,7 +38,6 @@ import (
 	"k8s.io/client-go/util/retry"
 
 	"github.com/conforma/cli/internal/fetchers/oci/files"
-	"github.com/conforma/cli/internal/image"
 	"github.com/conforma/cli/internal/utils/oci"
 )
 
@@ -388,22 +387,13 @@ func ociDescriptor(bctx rego.BuiltinContext, a *ast.Term) (*ast.Term, error) {
 
 	client := oci.NewClient(bctx.Context)
 
-	uri, err := resolveIfNeeded(client, string(uriValue))
+	uri, ref, err := resolveIfNeeded(client, string(uriValue))
 	if err != nil {
 		logger.WithField("action", "resolveIfNeeded").Error(err)
 		return nil, nil
 	}
 	logger = logger.WithField("ref", uri)
 
-	ref, err := name.NewDigest(uri)
-	if err != nil {
-		logger.WithFields(log.Fields{
-			"action": "new digest",
-			"error":  err,
-		}).Error("failed to create new digest")
-		return nil, nil
-	}
-
 	descriptor, err := client.Head(ref)
 	if err != nil {
 		logger.WithFields(log.Fields{
@@ -430,22 +420,13 @@ func ociImageManifest(bctx rego.BuiltinContext, a *ast.Term) (*ast.Term, error)
 
 	client := oci.NewClient(bctx.Context)
 
-	uri, err := resolveIfNeeded(client, string(uriValue))
+	uri, ref, err := resolveIfNeeded(client, string(uriValue))
 	if err != nil {
 		logger.WithField("action", "resolveIfNeeded").Error(err)
 		return nil, nil
 	}
 	logger = logger.WithField("ref", uri)
 
-	ref, err := name.NewDigest(uri)
-	if err != nil {
-		logger.WithFields(log.Fields{
-			"action": "new digest",
-			"error":  err,
-		}).Error("failed to create new digest")
-		return nil, nil
-	}
-
 	var image v1.Image
 	err = retry.OnError(retry.DefaultRetry, func(_ error) bool { return true }, func() error {
 		image, err = client.Image(ref)
@@ -575,22 +556,13 @@ func ociImageIndex(bctx rego.BuiltinContext, a *ast.Term) (*ast.Term, error) {
 
 	client := oci.NewClient(bctx.Context)
 
-	uri, err := resolveIfNeeded(client, string(uriValue))
+	uri, ref, err := resolveIfNeeded(client, string(uriValue))
 	if err != nil {
 		logger.WithField("action", "resolveIfNeeded").Error(err)
 		return nil, nil
 	}
 	logger = logger.WithField("ref", uri)
 
-	ref, err := name.NewDigest(uri)
-	if err != nil {
-		logger.WithFields(log.Fields{
-			"action": "new digest",
-			"error":  err,
-		}).Error("failed to create new digest")
-		return nil, nil
-	}
-
 	imageIndex, err := client.Index(ref)
 	if err != nil {
 		logger.WithFields(log.Fields{
@@ -686,23 +658,43 @@ func newAnnotationsTerm(annotations map[string]string) *ast.Term {
 	return ast.ObjectTerm(annotationTerms...)
 }
 
-func resolveIfNeeded(client oci.Client, uri string) (string, error) {
-	if !strings.Contains(uri, "@") {
-		original := uri
-		ref, err := image.NewImageReference(uri)
-		if err != nil {
-			return "", fmt.Errorf("unable to parse reference: %w", err)
-		}
+func resolveIfNeeded(client oci.Client, uri string) (string, name.Reference, error) {
+	ref, err := parseReference(uri)
+	if err != nil {
+		return "", nil, fmt.Errorf("unable to parse reference: %w", err)
+	}
 
-		digest, err := client.ResolveDigest(ref.Ref())
-		if err != nil {
-			return "", fmt.Errorf("unable to resolve digest: %w", err)
-		}
-		uri = fmt.Sprintf("%s@%s", uri, digest)
+	// If it's already a digest reference, return as is
+	if _, ok := ref.(name.Digest); ok {
+		return uri, ref, nil
+	}
+
+	// For tag references, resolve to digest
+	digest, err := client.ResolveDigest(ref)
+	if err != nil {
+		return "", nil, fmt.Errorf("unable to resolve digest: %w", err)
+	}
 
-		log.Debugf("resolved image reference %q to %q", original, uri)
+	resolved := fmt.Sprintf("%s@%s", uri, digest)
+	log.Debugf("resolved image reference %q to %q", uri, resolved)
+	return resolved, ref, nil
+}
+
+func parseReference(uri string) (name.Reference, error) {
+	// Try to parse as digest first, if that fails with ErrBadName, try as tag
+	ref, err := name.NewDigest(uri)
+	if err != nil {
+		if errors.Is(err, &name.ErrBadName{}) {
+			tag, err := name.NewTag(uri)
+			if err != nil {
+				return nil, fmt.Errorf("invalid reference format: %w", err)
+			}
+			return tag, nil
+		} else {
+			return nil, fmt.Errorf("invalid digest format: %w", err)
+		}
 	}
-	return uri, nil
+	return ref, nil
 }
 
 func init() {

internal/rego/oci/oci_test.go

@@ -184,12 +184,42 @@ func TestOCIDescriptorManifest(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:           "tag-based URI with error",
+			ref:            ast.StringTerm("registry.local/spam:latest"),
+			resolvedDigest: "sha256:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b",
+			headErr:        errors.New("tag error"),
+			wantErr:        true,
+		},
 		{
 			name:       "resolve error",
 			ref:        ast.StringTerm("registry.local/spam:latest"),
 			resolveErr: errors.New("kaboom!"),
 			wantErr:    true,
 		},
+		{
+			name:           "unsupported digest algorithm",
+			ref:            ast.StringTerm("registry.local/spam@sha512:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b"),
+			resolvedDigest: "sha256:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b",
+			wantErr:        true,
+		},
+		{
+			name:           "malformed digest with extra @",
+			ref:            ast.StringTerm("registry.local/spam@@sha256:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b"),
+			resolvedDigest: "sha256:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b",
+			wantErr:        true,
+		},
+		{
+			name:           "invalid tag after digest fallback",
+			ref:            ast.StringTerm("registry.local/spam:!nv@lid"),
+			resolvedDigest: "sha256:01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b",
+			wantErr:        true,
+		},
+		{
+			name:    "invalid digest format",
+			ref:     ast.StringTerm("registry.local/spam@sha256:invalid-digest-format"),
+			wantErr: true,
+		},
 	}
 
 	for _, c := range cases {
@@ -464,6 +494,7 @@ func TestOCIImageManifest(t *testing.T) {
 		})
 	}
 }
+
 func TestOCIImageFiles(t *testing.T) {
 
 	image, err := crane.Image(map[string][]byte{
@@ -747,3 +778,115 @@ func TestFunctionsRegistered(t *testing.T) {
 		})
 	}
 }
+
+func TestParseReference(t *testing.T) {
+	cases := []struct {
+		name    string
+		uri     string
+		wantErr bool
+	}{
+		{
+			name: "valid digest",
+			uri:  "registry.local/spam@sha256:4bbf56a3a9231f752d3b9c174637975f0f83ed2b15e65799837c571e4ef3374b",
+		},
+		{
+			name: "valid tag",
+			uri:  "registry.local/spam:latest",
+		},
+		{
+			name:    "invalid digest format",
+			uri:     "registry.local/spam@sha256:invalid",
+			wantErr: true,
+		},
+		{
+			name:    "invalid tag format",
+			uri:     "registry.local/spam:!nv@lid",
+			wantErr: true,
+		},
+		{
+			name:    "trailing @",
+			uri:     "registry.local/spam@",
+			wantErr: true,
+		},
+		{
+			name:    "multiple @",
+			uri:     "registry.local/spam@@sha256:abc123",
+			wantErr: true,
+		},
+		{
+			name:    "empty string",
+			uri:     "",
+			wantErr: true,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			ref, err := parseReference(c.uri)
+			if c.wantErr {
+				require.Error(t, err)
+				require.Nil(t, ref)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, ref)
+			}
+		})
+	}
+}
+
+func TestResolveIfNeeded(t *testing.T) {
+	cases := []struct {
+		name           string
+		uri            string
+		resolvedDigest string
+		resolveErr     error
+		wantErr        bool
+	}{
+		{
+			name: "digest reference unchanged",
+			uri:  "registry.local/spam@sha256:4bbf56a3a9231f752d3b9c174637975f0f83ed2b15e65799837c571e4ef3374b",
+		},
+		{
+			name:           "tag reference resolved",
+			uri:            "registry.local/spam:latest",
+			resolvedDigest: "sha256:4bbf56a3a9231f752d3b9c174637975f0f83ed2b15e65799837c571e4ef3374b",
+		},
+		{
+			name:       "resolve error",
+			uri:        "registry.local/spam:latest",
+			resolveErr: errors.New("resolve error"),
+			wantErr:    true,
+		},
+		{
+			name:    "invalid reference",
+			uri:     "registry.local/spam@",
+			wantErr: true,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			client := fake.FakeClient{}
+			if c.resolveErr != nil {
+				client.On("ResolveDigest", mock.Anything).Return("", c.resolveErr)
+			} else if c.resolvedDigest != "" {
+				client.On("ResolveDigest", mock.Anything).Return(c.resolvedDigest, nil)
+			}
+
+			uri, ref, err := resolveIfNeeded(&client, c.uri)
+			if c.wantErr {
+				require.Error(t, err)
+				require.Empty(t, uri)
+				require.Nil(t, ref)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, ref)
+				if c.resolvedDigest != "" {
+					require.Contains(t, uri, c.resolvedDigest)
+				} else {
+					require.Equal(t, c.uri, uri)
+				}
+			}
+		})
+	}
+}