Harden mirror-aware SOCI artifact fetch and fallback

Follow-up to c751c6d to fully address mirror behavior for SOCI artifact
discovery/fetch (index + zTOCs).

The original fix moved SOCI artifact lookups away from the image’s origin
locator, but it still effectively depended on the first resolved host and
didn’t fully preserve per-host details. In practice that meant we could still
fail in mirror-heavy setups where the first endpoint is unavailable or where
the mirror is configured with HTTP/path differences.

This change finishes that work so SOCI index and zTOC fetches behave like the
rest of our mirror-aware pull flow. We now carry the full RegistryHost context
through the fetch path, rewrite locators with host path awareness, honor mirror
scheme when constructing ORAS repositories, and retry across all configured
hosts instead of stopping at the first one.

I also tightened error handling around missing host clients and added tests that
cover the new behavior directly: locator rewriting edge cases, host scheme
handling, and fallback from one mirror endpoint to another (including the
all-hosts-fail path). The goal is to make artifact discovery/fetch resilient
and predictable in the same environments where layer pulls already are.

Also adds a simple integration test.

Signed-off-by: Milos Gajdos <milosthegajdos@gmail.com>

Author: milosgajdos

fs/artifact_fetcher.go

@@ -77,8 +77,8 @@ type orasBlobStore struct {
 	*remote.Repository
 }
 
-func newRemoteBlobStore(refspec reference.Spec, client *http.Client) (*orasBlobStore, error) {
-	repo, err := newRemoteStore(refspec, client)
+func newRemoteBlobStoreFromHost(refspec reference.Spec, host docker.RegistryHost) (*orasBlobStore, error) {
+	repo, err := newRemoteStoreFromHost(refspec, host)
 	if err != nil {
 		return nil, fmt.Errorf("cannot create remote store: %w", err)
 	}
@@ -230,16 +230,38 @@ func (c *clientWrapper) RoundTrip(req *http.Request) (*http.Response, error) {
 	return c.Client.Do(req)
 }
 
-func withLocatorHost(refspec reference.Spec, host string) (reference.Spec, error) {
+func withLocatorHost(refspec reference.Spec, host docker.RegistryHost) (reference.Spec, error) {
+	if host.Host == "" || strings.Contains(host.Host, "/") {
+		return reference.Spec{}, fmt.Errorf("invalid registry host %q", host.Host)
+	}
+
 	repoPath := strings.TrimPrefix(refspec.Locator, refspec.Hostname())
 	repoPath = strings.TrimPrefix(repoPath, "/")
 	if repoPath == "" {
 		return reference.Spec{}, fmt.Errorf("invalid image locator %q", refspec.Locator)
 	}
-	refspec.Locator = path.Join(host, repoPath)
+
+	repoPrefix, err := repositoryPrefixFromHostPath(host.Path)
+	if err != nil {
+		return reference.Spec{}, err
+	}
+
+	refspec.Locator = path.Join(host.Host, repoPrefix, repoPath)
 	return refspec, nil
 }
 
+func repositoryPrefixFromHostPath(hostPath string) (string, error) {
+	clean := path.Clean(hostPath)
+	switch {
+	case clean == ".", clean == "/", clean == "/v2":
+		return "", nil
+	case strings.HasPrefix(clean, "/v2/"):
+		return strings.TrimPrefix(clean, "/v2/"), nil
+	default:
+		return "", fmt.Errorf("unsupported registry host path %q; expected /v2 or /v2/<prefix>", hostPath)
+	}
+}
+
 func newRemoteStore(refspec reference.Spec, client *http.Client) (*remote.Repository, error) {
 	repo, err := remote.NewRepository(refspec.Locator)
 	if err != nil {
@@ -254,6 +276,26 @@ func newRemoteStore(refspec reference.Spec, client *http.Client) (*remote.Reposi
 	return repo, nil
 }
 
+func newRemoteStoreFromHost(refspec reference.Spec, host docker.RegistryHost) (*remote.Repository, error) {
+	repo, err := newRemoteStore(refspec, host.Client)
+	if err != nil {
+		return nil, err
+	}
+
+	switch strings.ToLower(host.Scheme) {
+	case "http":
+		repo.PlainHTTP = true
+	case "", "https":
+		// Keep the default behavior from newRemoteStore:
+		// - localhost uses plain HTTP
+		// - all other hosts use HTTPS
+	default:
+		return nil, fmt.Errorf("unsupported registry scheme %q for host %q", host.Scheme, host.Host)
+	}
+
+	return repo, nil
+}
+
 // Constructs a new artifact fetcher
 // Takes in the image reference, the local store and the resolver
 func newArtifactFetcher(refspec reference.Spec, localStore store.BasicStore, remoteStore resolverStorage) (*artifactFetcher, error) {

fs/artifact_fetcher_test.go

@@ -23,11 +23,13 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"testing"
 
 	"github.com/awslabs/soci-snapshotter/soci"
 	"github.com/awslabs/soci-snapshotter/soci/store"
 	"github.com/containerd/containerd/reference"
+	"github.com/containerd/containerd/remotes/docker"
 	"github.com/google/go-cmp/cmp"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -306,16 +308,129 @@ func TestNewRemoteStore(t *testing.T) {
 }
 
 func TestWithLocatorHost(t *testing.T) {
+	testCases := []struct {
+		name            string
+		host            docker.RegistryHost
+		expectedLocator string
+		expectedError   string
+	}{
+		{
+			name: "rewrites locator with mirror host",
+			host: docker.RegistryHost{
+				Host: "example-mirror.local",
+				Path: "/v2",
+			},
+			expectedLocator: "example-mirror.local/library/nginx",
+		},
+		{
+			name: "rewrites locator with mirror path prefix",
+			host: docker.RegistryHost{
+				Host: "example-mirror.local",
+				Path: "/v2/team-a",
+			},
+			expectedLocator: "example-mirror.local/team-a/library/nginx",
+		},
+		{
+			name: "invalid mirror path",
+			host: docker.RegistryHost{
+				Host: "example-mirror.local",
+				Path: "/custom/path",
+			},
+			expectedError: "unsupported registry host path",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			refspec, err := reference.Parse("docker.io/library/nginx:latest")
+			if err != nil {
+				t.Fatalf("unexpected failure parsing reference: %v", err)
+			}
+			rewritten, err := withLocatorHost(refspec, tc.host)
+			if tc.expectedError != "" {
+				if err == nil {
+					t.Fatalf("expected error containing %q, got nil", tc.expectedError)
+				}
+				if got := err.Error(); !strings.Contains(got, tc.expectedError) {
+					t.Fatalf("unexpected error, expected substring %q, got %q", tc.expectedError, got)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected failure rewriting reference: %v", err)
+			}
+			if rewritten.Locator != tc.expectedLocator {
+				t.Fatalf("unexpected locator, expected %q, got %q", tc.expectedLocator, rewritten.Locator)
+			}
+		})
+	}
+}
+
+func TestNewRemoteStoreFromHost(t *testing.T) {
+	client := &http.Client{}
 	refspec, err := reference.Parse("docker.io/library/nginx:latest")
 	if err != nil {
 		t.Fatalf("unexpected failure parsing reference: %v", err)
 	}
-	rewritten, err := withLocatorHost(refspec, "example-mirror.local")
-	if err != nil {
-		t.Fatalf("unexpected failure rewriting reference: %v", err)
+
+	testCases := []struct {
+		name              string
+		host              docker.RegistryHost
+		shouldBePlainHTTP bool
+		expectedError     string
+	}{
+		{
+			name: "http mirror uses plain http",
+			host: docker.RegistryHost{
+				Host:   "example-mirror.local",
+				Scheme: "http",
+				Client: client,
+			},
+			shouldBePlainHTTP: true,
+		},
+		{
+			name: "https mirror does not use plain http",
+			host: docker.RegistryHost{
+				Host:   "example-mirror.local",
+				Scheme: "https",
+				Client: client,
+			},
+			shouldBePlainHTTP: false,
+		},
+		{
+			name: "invalid scheme fails",
+			host: docker.RegistryHost{
+				Host:   "example-mirror.local",
+				Scheme: "ftp",
+				Client: client,
+			},
+			expectedError: "unsupported registry scheme",
+		},
 	}
-	if rewritten.Locator != "example-mirror.local/library/nginx" {
-		t.Fatalf("unexpected locator, expected %q, got %q", "example-mirror.local/library/nginx", rewritten.Locator)
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			rewritten, err := withLocatorHost(refspec, docker.RegistryHost{Host: tc.host.Host, Path: "/v2"})
+			if err != nil {
+				t.Fatalf("unexpected failure rewriting reference: %v", err)
+			}
+			store, err := newRemoteStoreFromHost(rewritten, tc.host)
+			if tc.expectedError != "" {
+				if err == nil {
+					t.Fatalf("expected error containing %q, got nil", tc.expectedError)
+				}
+				if got := err.Error(); !strings.Contains(got, tc.expectedError) {
+					t.Fatalf("unexpected error, expected substring %q, got %q", tc.expectedError, got)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error, got %v", err)
+			}
+			if store.PlainHTTP != tc.shouldBePlainHTTP {
+				t.Fatalf("unexpected plain http setting, expected %v, got %v", tc.shouldBePlainHTTP, store.PlainHTTP)
+			}
+		})
 	}
 }
 

fs/fetch_soci_index_test.go

@@ -0,0 +1,166 @@
+/*
+   Copyright The Soci Snapshotter Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package fs
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strconv"
+	"sync/atomic"
+	"testing"
+
+	"github.com/awslabs/soci-snapshotter/snapshot"
+	"github.com/awslabs/soci-snapshotter/soci"
+	"github.com/awslabs/soci-snapshotter/soci/store"
+	"github.com/containerd/containerd/remotes/docker"
+	"github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"oras.land/oras-go/v2/content/memory"
+)
+
+type inMemoryTestStore struct {
+	*memory.Store
+}
+
+func newInMemoryTestStore() *inMemoryTestStore {
+	return &inMemoryTestStore{
+		Store: memory.New(),
+	}
+}
+
+func (s *inMemoryTestStore) Label(context.Context, ocispec.Descriptor, string, string) error {
+	return nil
+}
+
+func (s *inMemoryTestStore) Delete(context.Context, digest.Digest) error {
+	return nil
+}
+
+func (s *inMemoryTestStore) BatchOpen(ctx context.Context) (context.Context, store.CleanupFunc, error) {
+	return ctx, store.NopCleanup, nil
+}
+
+func registryHostFromServer(ts *httptest.Server) docker.RegistryHost {
+	u, _ := url.Parse(ts.URL)
+	return docker.RegistryHost{
+		Host:   u.Host,
+		Scheme: u.Scheme,
+		Path:   "/v2",
+		Client: ts.Client(),
+	}
+}
+
+func TestFetchSociIndexFallsBackToNextHost(t *testing.T) {
+	indexBytes, err := soci.MarshalIndex(soci.NewIndex(soci.V2, nil, nil, nil))
+	if err != nil {
+		t.Fatalf("failed to marshal index: %v", err)
+	}
+	indexDigest := digest.FromBytes(indexBytes)
+	manifestDigest := digest.FromString("manifest")
+	indexBlobPath := "/v2/library/nginx/blobs/" + indexDigest.String()
+	indexManifestPath := "/v2/library/nginx/manifests/" + indexDigest.String()
+
+	var firstHostRequests int32
+	firstHost := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&firstHostRequests, 1)
+		http.Error(w, "temporary upstream failure", http.StatusBadGateway)
+	}))
+	defer firstHost.Close()
+
+	var secondHostRequests int32
+	secondHost := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&secondHostRequests, 1)
+		if r.URL.Path != indexBlobPath && r.URL.Path != indexManifestPath {
+			http.NotFound(w, r)
+			return
+		}
+		w.Header().Set("Content-Type", "application/octet-stream")
+		w.Header().Set("Content-Length", strconv.Itoa(len(indexBytes)))
+		switch r.Method {
+		case http.MethodHead:
+			w.WriteHeader(http.StatusOK)
+		case http.MethodGet:
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write(indexBytes)
+		default:
+			w.WriteHeader(http.StatusMethodNotAllowed)
+		}
+	}))
+	defer secondHost.Close()
+
+	fs := &filesystem{
+		contentStore: newInMemoryTestStore(),
+	}
+	index, err := fs.fetchSociIndex(
+		context.Background(),
+		"docker.io/library/nginx:latest",
+		indexDigest.String(),
+		manifestDigest.String(),
+		[]docker.RegistryHost{
+			registryHostFromServer(firstHost),
+			registryHostFromServer(secondHost),
+		},
+	)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if index == nil {
+		t.Fatal("expected non-nil index")
+	}
+	if got := atomic.LoadInt32(&firstHostRequests); got == 0 {
+		t.Fatal("expected first registry host to be attempted")
+	}
+	if got := atomic.LoadInt32(&secondHostRequests); got == 0 {
+		t.Fatal("expected second registry host to be attempted")
+	}
+}
+
+func TestFetchSociIndexReturnsNoIndexWhenAllHostsFail(t *testing.T) {
+	indexBytes, err := soci.MarshalIndex(soci.NewIndex(soci.V2, nil, nil, nil))
+	if err != nil {
+		t.Fatalf("failed to marshal index: %v", err)
+	}
+	indexDigest := digest.FromBytes(indexBytes)
+	manifestDigest := digest.FromString("manifest")
+
+	failingHost := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "temporary upstream failure", http.StatusBadGateway)
+	}))
+	defer failingHost.Close()
+
+	fs := &filesystem{
+		contentStore: newInMemoryTestStore(),
+	}
+	_, err = fs.fetchSociIndex(
+		context.Background(),
+		"docker.io/library/nginx:latest",
+		indexDigest.String(),
+		manifestDigest.String(),
+		[]docker.RegistryHost{
+			registryHostFromServer(failingHost),
+		},
+	)
+	if err == nil {
+		t.Fatal("expected error but got nil")
+	}
+	if !errors.Is(err, snapshot.ErrNoIndex) {
+		t.Fatalf("expected error to match snapshot.ErrNoIndex, got: %v", err)
+	}
+}