Merge pull request #2204 from dgageot/better-oci-refs

Better oci refs

Author: dgageot

pkg/config/sources.go

@@ -110,70 +110,82 @@ func (a ociSource) ParentDir() string {
 	return ""
 }
 
-// Read loads an agent configuration from an OCI artifact
+// Read loads an agent configuration from an OCI artifact.
 //
-// The OCI registry remains the source of truth
-// The local content store is used as a cache and fallback only
-// A forced re-pull is triggered exclusively when store corruption is detected
+// The OCI registry remains the source of truth.
+// The local content store is used as a cache and fallback only.
+// A forced re-pull is triggered exclusively when store corruption is detected.
 func (a ociSource) Read(ctx context.Context) ([]byte, error) {
 	store, err := content.NewStore()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create content store: %w", err)
 	}
 
-	tryLoad := func() ([]byte, error) {
-		af, err := store.GetArtifact(a.reference)
-		if err != nil {
-			return nil, err
+	// Normalize the reference so that equivalent forms (e.g.
+	// "agentcatalog/review-pr" and "index.docker.io/agentcatalog/review-pr:latest")
+	// resolve to the same store key that remote.Pull uses.
+	storeKey, err := remote.NormalizeReference(a.reference)
+	if err != nil {
+		return nil, fmt.Errorf("normalizing OCI reference %s: %w", a.reference, err)
+	}
+
+	// For digest references, the content is immutable. If we already have
+	// the artifact locally, serve it directly without any network call.
+	if remote.IsDigestReference(a.reference) {
+		if data, loadErr := loadArtifact(store, storeKey); loadErr == nil {
+			slog.Debug("Serving digest-pinned OCI artifact from cache", "ref", a.reference)
+			return data, nil
 		}
-		return []byte(af), nil
 	}
 
-	// Check if we have any local metadata (same as before)
-	_, metaErr := store.GetArtifactMetadata(a.reference)
-	hasLocal := metaErr == nil
+	// Check whether we have a local copy to fall back on.
+	hasLocal := hasLocalArtifact(store, storeKey)
 
-	// Always try normal pull first (preserves pull-interval behavior)
+	// Pull from registry (checks remote digest, skips download if unchanged).
 	if _, pullErr := remote.Pull(ctx, a.reference, false); pullErr != nil {
 		if !hasLocal {
 			return nil, fmt.Errorf("failed to pull OCI image %s: %w", a.reference, pullErr)
 		}
-
-		slog.Debug(
-			"Failed to check for OCI reference updates, using cached version",
-			"ref", a.reference,
-			"error", pullErr,
-		)
+		slog.Debug("Failed to check for OCI reference updates, using cached version",
+			"ref", a.reference, "error", pullErr)
 	}
 
-	// Try loading from store
-	data, err := tryLoad()
+	// Try loading from store.
+	data, err := loadArtifact(store, storeKey)
 	if err == nil {
 		return data, nil
 	}
 
-	// If loading failed due to corruption, force re-pull
-	if errors.Is(err, content.ErrStoreCorrupted) {
-		slog.Warn(
-			"Local OCI store corrupted, forcing re-pull",
-			"ref", a.reference,
-		)
+	// If corrupted, force re-pull and try once more.
+	if !errors.Is(err, content.ErrStoreCorrupted) {
+		return nil, fmt.Errorf("failed to load agent from OCI source %s: %w", a.reference, err)
+	}
 
-		if _, pullErr := remote.Pull(ctx, a.reference, true); pullErr != nil {
-			return nil, fmt.Errorf("failed to force re-pull OCI image %s: %w", a.reference, pullErr)
-		}
+	slog.Warn("Local OCI store corrupted, forcing re-pull", "ref", a.reference)
+	if _, pullErr := remote.Pull(ctx, a.reference, true); pullErr != nil {
+		return nil, fmt.Errorf("failed to force re-pull OCI image %s: %w", a.reference, pullErr)
+	}
 
-		data, err = tryLoad()
-		if err == nil {
-			return data, nil
-		}
+	data, err = loadArtifact(store, storeKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load agent from OCI source %s: %w", a.reference, err)
+	}
+	return data, nil
+}
+
+// loadArtifact reads the agent YAML from the content store.
+func loadArtifact(store *content.Store, storeKey string) ([]byte, error) {
+	af, err := store.GetArtifact(storeKey)
+	if err != nil {
+		return nil, err
 	}
+	return []byte(af), nil
+}
 
-	return nil, fmt.Errorf(
-		"failed to load agent from OCI source %s: %w",
-		a.reference,
-		err,
-	)
+// hasLocalArtifact reports whether the content store has metadata for the given key.
+func hasLocalArtifact(store *content.Store, storeKey string) bool {
+	_, err := store.GetArtifactMetadata(storeKey)
+	return err == nil
 }
 
 // urlSource is used to load an agent configuration from an HTTP/HTTPS URL.

pkg/config/sources_test.go

@@ -8,12 +8,59 @@ import (
 	"sync/atomic"
 	"testing"
 
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/static"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/docker/docker-agent/pkg/content"
 	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/remote"
 )
 
+func TestOCISource_DigestReference_ServesFromCache(t *testing.T) {
+	t.Parallel()
+
+	// Create a temporary content store and store a test artifact.
+	storeDir := t.TempDir()
+	store, err := content.NewStore(content.WithBaseDir(storeDir))
+	require.NoError(t, err)
+
+	testData := []byte("version: v1\nname: test-agent")
+	layer := static.NewLayer(testData, "application/yaml")
+	img, err := mutate.AppendLayers(empty.Image, layer)
+	require.NoError(t, err)
+	img = mutate.Annotations(img, map[string]string{
+		"io.docker.agent.version": "test",
+	}).(v1.Image)
+
+	ref := "test-digest-cache/agent:latest"
+	digest, err := store.StoreArtifact(img, ref)
+	require.NoError(t, err)
+
+	// Build a digest reference using the stored digest.
+	digestRef := "test-digest-cache/agent@" + digest
+
+	// Read via ociSource. Since the reference is pinned by digest and is
+	// present in the local store, this must succeed without any network call.
+	// We override the default store directory via an env-based approach;
+	// instead, we directly exercise the cache-hit logic by verifying the
+	// store lookup works with the normalized key.
+	storeKey, err := remote.NormalizeReference(digestRef)
+	require.NoError(t, err)
+
+	// Verify the store can resolve the digest key directly.
+	data, err := store.GetArtifact(storeKey)
+	require.NoError(t, err)
+	assert.Equal(t, string(testData), data)
+
+	// Also verify that IsDigestReference correctly identifies this.
+	assert.True(t, remote.IsDigestReference(digestRef))
+	assert.False(t, remote.IsDigestReference(ref))
+}
+
 func TestURLSource_Read(t *testing.T) {
 	t.Parallel()
 
@@ -204,11 +251,11 @@ func TestURLSource_Read_FallsBackToCacheOnHTTPError(t *testing.T) {
 func TestURLSource_Read_UpdatesCacheWhenContentChanges(t *testing.T) {
 	// Not parallel - uses shared cache directory
 
-	var content atomic.Value
-	content.Store("initial content update")
+	var serverContent atomic.Value
+	serverContent.Store("initial content update")
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		currentContent := content.Load().(string)
+		currentContent := serverContent.Load().(string)
 		etag := `"etag-` + currentContent + `"`
 
 		if r.Header.Get("If-None-Match") == etag {
@@ -239,7 +286,7 @@ func TestURLSource_Read_UpdatesCacheWhenContentChanges(t *testing.T) {
 	assert.Equal(t, "initial content update", string(data))
 
 	// Change content
-	content.Store("updated content update")
+	serverContent.Store("updated content update")
 
 	// Second read should get new content
 	data, err = source.Read(t.Context())

pkg/content/store.go

@@ -267,12 +267,18 @@ func (s *Store) DeleteArtifact(identifier string) error {
 // resolveIdentifier resolves a user-provided identifier (digest or reference)
 // into a concrete content digest stored in the local artifact store.
 func (s *Store) resolveIdentifier(identifier string) (string, error) {
-	// If the identifier is already a digest, we can return it directly.
-	// This bypasses the refs lookup entirely.
+	// If the identifier is already a bare digest, return it directly.
 	if strings.HasPrefix(identifier, "sha256:") {
 		return identifier, nil
 	}
 
+	// If the identifier is a digest reference (e.g. "repo@sha256:abc..."),
+	// extract and return the digest portion directly. Digest references
+	// are content-addressable, so the digest alone identifies the artifact.
+	if i := strings.LastIndex(identifier, "@sha256:"); i >= 0 {
+		return identifier[i+1:], nil
+	}
+
 	// If no tag is provided, default to ":latest".
 	// This mirrors standard OCI reference semantics.
 	if !strings.Contains(identifier, ":") {

pkg/content/store_test.go

@@ -109,3 +109,34 @@ func TestStoreResolution(t *testing.T) {
 		assert.NotNil(t, img)
 	}
 }
+
+func TestStoreResolution_DigestReference(t *testing.T) {
+	store, err := NewStore(WithBaseDir(t.TempDir()))
+	require.NoError(t, err)
+
+	testData := []byte("Digest resolution test")
+	layer := static.NewLayer(testData, types.OCIUncompressedLayer)
+	img := empty.Image
+	img, err = mutate.AppendLayers(img, layer)
+	require.NoError(t, err)
+
+	tagRef := "myrepo/agent:v1"
+	digest, err := store.StoreArtifact(img, tagRef)
+	require.NoError(t, err)
+
+	// Bare digest should resolve.
+	retrievedImg, err := store.GetArtifactImage(digest)
+	require.NoError(t, err)
+	assert.NotNil(t, retrievedImg)
+
+	// Digest reference (repo@sha256:...) should also resolve.
+	digestRef := "myrepo/agent@" + digest
+	retrievedImg, err = store.GetArtifactImage(digestRef)
+	require.NoError(t, err)
+	assert.NotNil(t, retrievedImg)
+
+	// Metadata lookup via digest reference should work too.
+	meta, err := store.GetArtifactMetadata(digestRef)
+	require.NoError(t, err)
+	assert.Equal(t, digest, meta.Digest)
+}

pkg/remote/pull.go

@@ -10,6 +10,29 @@ import (
 	"github.com/docker/docker-agent/pkg/content"
 )
 
+// NormalizeReference parses an OCI reference and returns the normalized
+// store key that Pull uses to store artifacts. This ensures that equivalent
+// references (e.g. "agentcatalog/review-pr" and
+// "index.docker.io/agentcatalog/review-pr:latest") map to the same key.
+func NormalizeReference(registryRef string) (string, error) {
+	ref, err := name.ParseReference(registryRef)
+	if err != nil {
+		return "", fmt.Errorf("parsing registry reference %s: %w", registryRef, err)
+	}
+	return ref.Context().RepositoryStr() + separator(ref) + ref.Identifier(), nil
+}
+
+// IsDigestReference reports whether the given reference pins a specific
+// image digest (e.g. "repo@sha256:abc...").
+func IsDigestReference(registryRef string) bool {
+	ref, err := name.ParseReference(registryRef)
+	if err != nil {
+		return false
+	}
+	_, ok := ref.(name.Digest)
+	return ok
+}
+
 // Pull pulls an artifact from a registry and stores it in the content store
 func Pull(ctx context.Context, registryRef string, force bool, opts ...crane.Option) (string, error) {
 	opts = append(opts, crane.WithContext(ctx))

pkg/remote/pull_test.go

@@ -72,6 +72,81 @@ func TestPullIntegration(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestNormalizeReference(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		ref      string
+		expected string
+	}{
+		{
+			name:     "short reference gets normalized",
+			ref:      "agentcatalog/review-pr",
+			expected: "agentcatalog/review-pr:latest",
+		},
+		{
+			name:     "fully qualified reference gets normalized to same key",
+			ref:      "index.docker.io/agentcatalog/review-pr:latest",
+			expected: "agentcatalog/review-pr:latest",
+		},
+		{
+			name:     "tagged reference preserves tag",
+			ref:      "agentcatalog/review-pr:v1",
+			expected: "agentcatalog/review-pr:v1",
+		},
+		{
+			name:     "digest reference preserves digest",
+			ref:      "agentcatalog/review-pr@sha256:0000000000000000000000000000000000000000000000000000000000000000",
+			expected: "agentcatalog/review-pr@sha256:0000000000000000000000000000000000000000000000000000000000000000",
+		},
+		{
+			name:     "non-docker-hub registry",
+			ref:      "ghcr.io/myorg/agent:v2",
+			expected: "myorg/agent:v2",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			result, err := NormalizeReference(tt.ref)
+			require.NoError(t, err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestNormalizeReference_InvalidReference(t *testing.T) {
+	t.Parallel()
+
+	_, err := NormalizeReference(":::invalid")
+	require.Error(t, err)
+}
+
+func TestIsDigestReference(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		ref      string
+		expected bool
+	}{
+		{"tag reference", "agentcatalog/review-pr:latest", false},
+		{"implicit tag", "agentcatalog/review-pr", false},
+		{"digest reference", "agentcatalog/review-pr@sha256:0000000000000000000000000000000000000000000000000000000000000000", true},
+		{"fully qualified digest", "index.docker.io/agentcatalog/review-pr@sha256:0000000000000000000000000000000000000000000000000000000000000000", true},
+		{"invalid reference", ":::invalid", false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			assert.Equal(t, tt.expected, IsDigestReference(tt.ref))
+		})
+	}
+}
+
 func TestSeparator(t *testing.T) {
 	t.Parallel()