Repository: spegel-org/spegel
Pull request #538 (2024-07-10 13:56:11)

Replace mock OCI client with in memory client
This change implements an in memory OCI client and adds it to the e2e tests. I have also refactored some of the client interface code to work better with multiple implementations of the OCI client. This will hopefully make adding more implementations easier in the future while also enabling more tests using the memory client.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/spegel-org/spegel/pull/538?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org) Report
Attention: Patch coverage is `78.94737%` with `32 lines` in your changes missing coverage. Please review.

| [Files](https://app.codecov.io/gh/spegel-org/spegel/pull/538?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org) | Coverage Δ | |
|---|---|---|
| [pkg/state/state.go](https://app.codecov.io/gh/spegel-org/spegel/pull/538?src=pr&el=tree&filepath=pkg%2Fstate%2Fstate.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org#diff-cGtnL3N0YXRlL3N0YXRlLmdv) | `58.02% <100.00%> (ø)` | |
| [pkg/oci/containerd.go](https://app.codecov.io/gh/spegel-org/spegel/pull/538?src=pr&el=tree&filepath=pkg%2Foci%2Fcontainerd.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org#diff-cGtnL29jaS9jb250YWluZXJkLmdv) | `58.38% <81.81%> (+7.63%)` | :arrow_up: |
| [pkg/routing/memory.go](https://app.codecov.io/gh/spegel-org/spegel/pull/538?src=pr&el=tree&filepath=pkg%2Frouting%2Fmemory.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org#diff-cGtnL3JvdXRpbmcvbWVtb3J5Lmdv) | `95.12% <75.00%> (-4.88%)` | :arrow_down: |
| [pkg/oci/memory.go](https://app.codecov.io/gh/spegel-org/spegel/pull/538?src=pr&el=tree&filepath=pkg%2Foci%2Fmemory.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org#diff-cGtnL29jaS9tZW1vcnkuZ28=) | `83.05% <83.05%> (ø)` | |
| [pkg/oci/oci.go](https://app.codecov.io/gh/spegel-org/spegel/pull/538?src=pr&el=tree&filepath=pkg%2Foci%2Foci.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spegel-org#diff-cGtnL29jaS9vY2kuZ28=) | `74.64% <74.64%> (ø)` | |


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 482fb8f..0da9de5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added 
 
+- [#538](https://github.com/spegel-org/spegel/pull/538) Replace mock OCI client with in memory client.
+
 ### Changed
 
 - [#518](https://github.com/spegel-org/spegel/pull/518) Extend tests for image.
diff --git a/go.mod b/go.mod
index cf3e136..f13bb6f 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/alexflint/go-arg v1.5.1
 	github.com/containerd/containerd v1.7.19
 	github.com/containerd/containerd/api v1.7.19
+	github.com/containerd/errdefs v0.1.0
 	github.com/containerd/typeurl/v2 v2.1.1
 	github.com/go-logr/logr v1.4.2
 	github.com/ipfs/go-cid v0.4.1
@@ -43,7 +44,6 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/containerd/cgroups v1.1.0 // indirect
 	github.com/containerd/continuity v0.4.2 // indirect
-	github.com/containerd/errdefs v0.1.0 // indirect
 	github.com/containerd/fifo v1.1.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/containerd/platforms v0.2.1 // indirect
diff --git a/pkg/oci/containerd.go b/pkg/oci/containerd.go
index 776a6be..fb32a95 100644
--- a/pkg/oci/containerd.go
+++ b/pkg/oci/containerd.go
@@ -15,7 +15,7 @@ import (
 	"github.com/containerd/containerd"
 	eventtypes "github.com/containerd/containerd/api/events"
 	"github.com/containerd/containerd/content"
-	"github.com/containerd/containerd/images"
+	"github.com/containerd/errdefs"
 	"github.com/containerd/typeurl/v2"
 	"github.com/go-logr/logr"
 	"github.com/opencontainers/go-digest"
@@ -201,67 +201,6 @@ func (c *Containerd) ListImages(ctx context.Context) ([]Image, error) {
 	return imgs, nil
 }
 
-func (c *Containerd) AllIdentifiers(ctx context.Context, img Image) ([]string, error) {
-	client, err := c.Client()
-	if err != nil {
-		return nil, err
-	}
-	cImg, err := client.ImageService().Get(ctx, img.Name)
-	if err != nil {
-		return nil, err
-	}
-	keys := []string{}
-	err = images.Walk(ctx, images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
-		keys = append(keys, desc.Digest.String())
-		switch desc.MediaType {
-		case images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
-			var idx ocispec.Index
-			b, err := content.ReadBlob(ctx, client.ContentStore(), desc)
-			if err != nil {
-				return nil, fmt.Errorf("failed to read blob for manifest list: %w", err)
-			}
-			if err := json.Unmarshal(b, &idx); err != nil {
-				return nil, err
-			}
-			var descs []ocispec.Descriptor
-			for _, m := range idx.Manifests {
-				// Skip index layers that do not exist locally
-				if _, err := client.ContentStore().Info(ctx, m.Digest); err != nil {
-					continue
-				}
-				descs = append(descs, m)
-			}
-			if len(descs) == 0 {
-				return nil, fmt.Errorf("could not find any platforms with local content in manifest list: %v", desc.Digest)
-			}
-			return descs, nil
-		case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
-			var manifest ocispec.Manifest
-			b, err := content.ReadBlob(ctx, client.ContentStore(), desc)
-			if err != nil {
-				return nil, fmt.Errorf("failed to read blob for manifest: %w", err)
-			}
-			if err := json.Unmarshal(b, &manifest); err != nil {
-				return nil, err
-			}
-			keys = append(keys, manifest.Config.Digest.String())
-			for _, layer := range manifest.Layers {
-				keys = append(keys, layer.Digest.String())
-			}
-			return nil, nil
-		default:
-			return nil, fmt.Errorf("unexpected media type %v for digest: %v", desc.MediaType, desc.Digest)
-		}
-	}), cImg.Target)
-	if err != nil {
-		return nil, fmt.Errorf("failed to walk image manifests: %w", err)
-	}
-	if len(keys) == 0 {
-		return nil, errors.New("no image digests found")
-	}
-	return keys, nil
-}
-
 func (c *Containerd) Resolve(ctx context.Context, ref string) (digest.Digest, error) {
 	client, err := c.Client()
 	if err != nil {
@@ -280,6 +219,9 @@ func (c *Containerd) Size(ctx context.Context, dgst digest.Digest) (int64, error
 		return 0, err
 	}
 	info, err := client.ContentStore().Info(ctx, dgst)
+	if errors.Is(err, errdefs.ErrNotFound) {
+		return 0, errors.Join(ErrNotFound, err)
+	}
 	if err != nil {
 		return 0, err
 	}
@@ -292,27 +234,15 @@ func (c *Containerd) GetManifest(ctx context.Context, dgst digest.Digest) ([]byt
 		return nil, "", err
 	}
 	b, err := content.ReadBlob(ctx, client.ContentStore(), ocispec.Descriptor{Digest: dgst})
-	if err != nil {
-		return nil, "", err
-	}
-	var ud UnknownDocument
-	if err := json.Unmarshal(b, &ud); err != nil {
-		return nil, "", err
+	if errors.Is(err, errdefs.ErrNotFound) {
+		return nil, "", errors.Join(ErrNotFound, err)
 	}
-	if ud.MediaType != "" {
-		return b, ud.MediaType, nil
-	}
-	var ic ocispec.Image
-	if err := json.Unmarshal(b, &ic); err != nil {
+	if err != nil {
 		return nil, "", err
 	}
-	if isImageConfig(ic) {
-		return b, ocispec.MediaTypeImageConfig, nil
-	}
-	// Media type is not a required field. We need a fallback method if the field is not set.
-	mt, err := c.lookupMediaType(ctx, dgst)
+	mt, err := DetermineMediaType(b)
 	if err != nil {
-		return nil, "", fmt.Errorf("could not get media type for %s: %w", dgst.String(), err)
+		return nil, "", err
 	}
 	return b, mt, nil
 }
@@ -321,6 +251,9 @@ func (c *Containerd) GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCl
 	if c.contentPath != "" {
 		path := filepath.Join(c.contentPath, "blobs", dgst.Algorithm().String(), dgst.Encoded())
 		file, err := os.Open(path)
+		if errors.Is(err, os.ErrNotExist) {
+			return nil, errors.Join(ErrNotFound, err)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -331,6 +264,9 @@ func (c *Containerd) GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCl
 		return nil, err
 	}
 	ra, err := client.ContentStore().ReaderAt(ctx, ocispec.Descriptor{Digest: dgst})
+	if errors.Is(err, errdefs.ErrNotFound) {
+		return nil, errors.Join(ErrNotFound, err)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -343,85 +279,6 @@ func (c *Containerd) GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCl
 	}, nil
 }
 
-// lookupMediaType will resolve the media type for a digest without looking at the content.
-// Only use this as a fallback method as it is a lot slower than reading it from the file.
-// TODO: A cache would be helpful to speed up lookups for the same digets.
-func (c *Containerd) lookupMediaType(ctx context.Context, dgst digest.Digest) (string, error) {
-	logr.FromContextOrDiscard(ctx).Info("using Containerd fallback method to determine media type", "digest", dgst.String())
-	client, err := c.Client()
-	if err != nil {
-		return "", err
-	}
-	images, err := client.ImageService().List(ctx, fmt.Sprintf("target.digest==%s", dgst.String()))
-	if err != nil {
-		return "", err
-	}
-	if len(images) > 0 {
-		return images[0].Target.MediaType, nil
-	}
-	info, err := client.ContentStore().Info(ctx, dgst)
-	if err != nil {
-		return "", err
-	}
-	filter, err := getContentFilter(info.Labels)
-	if err != nil {
-		return "", err
-	}
-	var parentDgst digest.Digest
-	err = client.ContentStore().Walk(ctx, func(info content.Info) error {
-		for _, v := range info.Labels {
-			if v != dgst.String() {
-				continue
-			}
-			parentDgst = info.Digest
-			return filepath.SkipAll
-		}
-		return nil
-	}, filter)
-	if err != nil && !errors.Is(err, filepath.SkipAll) {
-		return "", err
-	}
-	if parentDgst == "" {
-		return "", errors.New("could not find parent")
-	}
-	b, err := content.ReadBlob(ctx, client.ContentStore(), ocispec.Descriptor{Digest: parentDgst})
-	if err != nil {
-		return "", err
-	}
-	var idx ocispec.Index
-	if err := json.Unmarshal(b, &idx); err != nil {
-		return "", err
-	}
-	for _, desc := range idx.Manifests {
-		if desc.Digest == dgst {
-			return desc.MediaType, nil
-		}
-	}
-	var manifest ocispec.Manifest
-	if err := json.Unmarshal(b, &manifest); err != nil {
-		return "", err
-	}
-	if manifest.Config.Digest == dgst {
-		return manifest.Config.MediaType, nil
-	}
-	for _, layer := range manifest.Layers {
-		if layer.Digest == dgst {
-			return layer.MediaType, nil
-		}
-	}
-	return "", errors.New("could not find reference in parent")
-}
-
-func getContentFilter(labels map[string]string) (string, error) {
-	for k, v := range labels {
-		if !strings.HasPrefix(k, "containerd.io/distribution.source") {
-			continue
-		}
-		return fmt.Sprintf(`labels.%q==%s`, k, v), nil
-	}
-	return "", errors.New("could not find distribution label to create content filter")
-}
-
 func getEventImage(e typeurl.Any) (string, EventType, error) {
 	if e == nil {
 		return "", "", errors.New("any cannot be nil")
@@ -452,19 +309,6 @@ func createFilters(registries []url.URL) (string, string) {
 	return listFilter, eventFilter
 }
 
-func isImageConfig(ic ocispec.Image) bool {
-	if ic.Architecture == "" {
-		return false
-	}
-	if ic.OS == "" {
-		return false
-	}
-	if ic.RootFS.Type == "" {
-		return false
-	}
-	return true
-}
-
 type hostFile struct {
 	HostConfigs map[string]hostConfig `toml:"host"`
 	Server      string                `toml:"server"`
@@ -479,7 +323,7 @@ type hostConfig struct {
 	Capabilities []string               `toml:"capabilities"`
 }
 
-// Refer to containerd registry configuration documentation for mor information about required configuration.
+// Refer to containerd registry configuration documentation for more information about required configuration.
 // https://github.com/containerd/containerd/blob/main/docs/cri/config.md#registry-configuration
 // https://github.com/containerd/containerd/blob/main/docs/hosts.md#registry-configuration---examples
 func AddMirrorConfiguration(ctx context.Context, fs afero.Fs, configPath string, registryURLs, mirrorURLs []url.URL, resolveTags, appendToBackup bool) error {
diff --git a/pkg/oci/memory.go b/pkg/oci/memory.go
new file mode 100644
index 0000000..7bfd4dc
--- /dev/null
+++ b/pkg/oci/memory.go
@@ -0,0 +1,116 @@
+package oci
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+
+	"github.com/opencontainers/go-digest"
+)
+
+var _ Client = &Memory{}
+
+type Memory struct {
+	blobs  map[digest.Digest][]byte
+	tags   map[string]digest.Digest
+	images []Image
+	mx     sync.RWMutex
+}
+
+func NewMemory() *Memory {
+	return &Memory{
+		images: []Image{},
+		tags:   map[string]digest.Digest{},
+		blobs:  map[digest.Digest][]byte{},
+	}
+}
+
+func (m *Memory) Name() string {
+	return "memory"
+}
+
+func (m *Memory) Verify(ctx context.Context) error {
+	return nil
+}
+
+func (m *Memory) Subscribe(ctx context.Context) (<-chan ImageEvent, <-chan error, error) {
+	return nil, nil, nil
+}
+
+func (m *Memory) ListImages(ctx context.Context) ([]Image, error) {
+	m.mx.RLock()
+	defer m.mx.RUnlock()
+
+	return m.images, nil
+}
+
+func (m *Memory) Resolve(ctx context.Context, ref string) (digest.Digest, error) {
+	m.mx.RLock()
+	defer m.mx.RUnlock()
+
+	dgst, ok := m.tags[ref]
+	if !ok {
+		return "", fmt.Errorf("could not resolve tag %s to a digest", ref)
+	}
+	return dgst, nil
+}
+
+func (m *Memory) Size(ctx context.Context, dgst digest.Digest) (int64, error) {
+	m.mx.RLock()
+	defer m.mx.RUnlock()
+
+	b, ok := m.blobs[dgst]
+	if !ok {
+		return 0, errors.Join(ErrNotFound, fmt.Errorf("size information for digest %s not found", dgst))
+	}
+	return int64(len(b)), nil
+}
+
+func (m *Memory) GetManifest(ctx context.Context, dgst digest.Digest) ([]byte, string, error) {
+	m.mx.RLock()
+	defer m.mx.RUnlock()
+
+	b, ok := m.blobs[dgst]
+	if !ok {
+		return nil, "", errors.Join(ErrNotFound, fmt.Errorf("manifest with digest %s not found", dgst))
+	}
+	mt, err := DetermineMediaType(b)
+	if err != nil {
+		return nil, "", err
+	}
+	return b, mt, nil
+}
+
+func (m *Memory) GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCloser, error) {
+	m.mx.RLock()
+	defer m.mx.RUnlock()
+
+	b, ok := m.blobs[dgst]
+	if !ok {
+		return nil, errors.Join(ErrNotFound, fmt.Errorf("blob with digest %s not found", dgst))
+	}
+	rc := io.NopCloser(bytes.NewBuffer(b))
+	return rc, nil
+}
+
+func (m *Memory) AddImage(img Image) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	m.images = append(m.images, img)
+	tagName, ok := img.TagName()
+	if !ok {
+		return
+	}
+	m.tags[tagName] = img.Digest
+}
+
+func (m *Memory) AddBlob(b []byte, dgst digest.Digest) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+
+	m.blobs[dgst] = b
+}
diff --git a/pkg/oci/mock.go b/pkg/oci/mock.go
deleted file mode 100644
index 424239a..0000000
--- a/pkg/oci/mock.go
+++ /dev/null
@@ -1,56 +0,0 @@
-package oci
-
-import (
-	"context"
-	"io"
-
-	"github.com/opencontainers/go-digest"
-)
-
-var _ Client = &MockClient{}
-
-type MockClient struct {
-	images []Image
-}
-
-func NewMockClient(images []Image) *MockClient {
-	return &MockClient{
-		images: images,
-	}
-}
-
-func (m *MockClient) Name() string {
-	return "mock"
-}
-
-func (m *MockClient) Verify(ctx context.Context) error {
-	return nil
-}
-
-func (m *MockClient) Subscribe(ctx context.Context) (<-chan ImageEvent, <-chan error, error) {
-	return nil, nil, nil
-}
-
-func (m *MockClient) ListImages(ctx context.Context) ([]Image, error) {
-	return m.images, nil
-}
-
-func (m *MockClient) AllIdentifiers(ctx context.Context, img Image) ([]string, error) {
-	return []string{img.Digest.String()}, nil
-}
-
-func (m *MockClient) Resolve(ctx context.Context, ref string) (digest.Digest, error) {
-	return "", nil
-}
-
-func (m *MockClient) Size(ctx context.Context, dgst digest.Digest) (int64, error) {
-	return 0, nil
-}
-
-func (m *MockClient) GetManifest(ctx context.Context, dgst digest.Digest) ([]byte, string, error) {
-	return nil, "", nil
-}
-
-func (m *MockClient) GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCloser, error) {
-	return nil, nil
-}
diff --git a/pkg/oci/oci.go b/pkg/oci/oci.go
index 13f86ed..ab702e1 100644
--- a/pkg/oci/oci.go
+++ b/pkg/oci/oci.go
@@ -2,23 +2,151 @@ package oci
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
 	"io"
 
+	"github.com/containerd/containerd/images"
 	"github.com/opencontainers/go-digest"
+	"github.com/opencontainers/image-spec/specs-go"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
-type UnknownDocument struct {
-	MediaType string `json:"mediaType,omitempty"`
-}
+var (
+	ErrNotFound = errors.New("content not found")
+)
 
 type Client interface {
+	// Name returns the name of the Client implementation.
 	Name() string
+
+	// Verify checks that all expected configuration is set.
 	Verify(ctx context.Context) error
+
+	// Subscribe will notify for any image events ocuring in the store backend.
 	Subscribe(ctx context.Context) (<-chan ImageEvent, <-chan error, error)
+
+	// ListImages returns a list of all local images.
 	ListImages(ctx context.Context) ([]Image, error)
-	AllIdentifiers(ctx context.Context, img Image) ([]string, error)
+
+	// Resolve returns the digest for the tagged image name reference.
+	// The ref is expected to be in the format `registry/name:tag`.
 	Resolve(ctx context.Context, ref string) (digest.Digest, error)
+
+	// Size returns the content byte size for the given digest.
+	// Will return ErrNotFound if the digest cannot be found.
 	Size(ctx context.Context, dgst digest.Digest) (int64, error)
+
+	// GetManifest returns the manifest content for the given digest.
+	// Will return ErrNotFound if the digest cannot be found.
 	GetManifest(ctx context.Context, dgst digest.Digest) ([]byte, string, error)
+
+	// GetBlob returns a stream of the blob content for the given digest.
+	// Will return ErrNotFound if the digest cannot be found.
 	GetBlob(ctx context.Context, dgst digest.Digest) (io.ReadCloser, error)
 }
+
+type UnknownDocument struct {
+	MediaType string `json:"mediaType"`
+	specs.Versioned
+}
+
+func DetermineMediaType(b []byte) (string, error) {
+	var ud UnknownDocument
+	if err := json.Unmarshal(b, &ud); err != nil {
+		return "", err
+	}
+	if ud.SchemaVersion == 2 && ud.MediaType != "" {
+		return ud.MediaType, nil
+	}
+	data := map[string]json.RawMessage{}
+	if err := json.Unmarshal(b, &data); err != nil {
+		return "", err
+	}
+	_, architectureOk := data["architecture"]
+	_, osOk := data["os"]
+	_, rootfsOk := data["rootfs"]
+	if architectureOk && osOk && rootfsOk {
+		return ocispec.MediaTypeImageConfig, nil
+	}
+	_, manifestsOk := data["manifests"]
+	if ud.SchemaVersion == 2 && manifestsOk {
+		return ocispec.MediaTypeImageIndex, nil
+	}
+	_, configOk := data["config"]
+	if ud.SchemaVersion == 2 && configOk {
+		return ocispec.MediaTypeImageManifest, nil
+	}
+	return "", errors.New("not able to determine media type")
+}
+
+func WalkImage(ctx context.Context, client Client, img Image) ([]string, error) {
+	keys := []string{}
+	err := walk(ctx, []digest.Digest{img.Digest}, func(dgst digest.Digest) ([]digest.Digest, error) {
+		b, mt, err := client.GetManifest(ctx, dgst)
+		if err != nil {
+			return nil, err
+		}
+		keys = append(keys, dgst.String())
+		switch mt {
+		case images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
+			var idx ocispec.Index
+			if err := json.Unmarshal(b, &idx); err != nil {
+				return nil, err
+			}
+			manifestDgsts := []digest.Digest{}
+			for _, m := range idx.Manifests {
+				_, err := client.Size(ctx, m.Digest)
+				if errors.Is(err, ErrNotFound) {
+					continue
+				}
+				if err != nil {
+					return nil, err
+				}
+				manifestDgsts = append(manifestDgsts, m.Digest)
+			}
+			if len(manifestDgsts) == 0 {
+				return nil, fmt.Errorf("could not find any platforms with local content in manifest %s", dgst)
+			}
+			return manifestDgsts, nil
+		case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
+			var manifest ocispec.Manifest
+			err := json.Unmarshal(b, &manifest)
+			if err != nil {
+				return nil, err
+			}
+			keys = append(keys, manifest.Config.Digest.String())
+			for _, layer := range manifest.Layers {
+				keys = append(keys, layer.Digest.String())
+			}
+			return nil, nil
+		default:
+			return nil, fmt.Errorf("unexpected media type %s for digest %s", mt, dgst)
+		}
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to walk image manifests: %w", err)
+	}
+	if len(keys) == 0 {
+		return nil, errors.New("no image digests found")
+	}
+	return keys, nil
+}
+
+func walk(ctx context.Context, dgsts []digest.Digest, handler func(dgst digest.Digest) ([]digest.Digest, error)) error {
+	for _, dgst := range dgsts {
+		children, err := handler(dgst)
+		if err != nil {
+			return err
+		}
+		if len(children) == 0 {
+			continue
+		}
+		err = walk(ctx, children, handler)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/pkg/routing/memory.go b/pkg/routing/memory.go
index b761e03..dc1186b 100644
--- a/pkg/routing/memory.go
+++ b/pkg/routing/memory.go
@@ -22,14 +22,16 @@ func NewMemoryRouter(resolver map[string][]netip.AddrPort, self netip.AddrPort)
 func (m *MemoryRouter) Ready(ctx context.Context) (bool, error) {
 	m.mx.RLock()
 	defer m.mx.RUnlock()
+
 	return len(m.resolver) > 0, nil
 }
 
 func (m *MemoryRouter) Resolve(ctx context.Context, key string, allowSelf bool, count int) (<-chan netip.AddrPort, error) {
-	peerCh := make(chan netip.AddrPort, count)
 	m.mx.RLock()
 	peers, ok := m.resolver[key]
 	m.mx.RUnlock()
+
+	peerCh := make(chan netip.AddrPort, count)
 	// If no peers exist close the channel to stop any consumer.
 	if !ok {
 		close(peerCh)
@@ -54,16 +56,24 @@ func (m *MemoryRouter) Advertise(ctx context.Context, keys []string) error {
 func (m *MemoryRouter) Add(key string, ap netip.AddrPort) {
 	m.mx.Lock()
 	defer m.mx.Unlock()
-	if v, ok := m.resolver[key]; ok {
-		m.resolver[key] = append(v, ap)
+
+	v, ok := m.resolver[key]
+	if !ok {
+		m.resolver[key] = []netip.AddrPort{ap}
 		return
 	}
-	m.resolver[key] = []netip.AddrPort{ap}
+	for _, h := range v {
+		if h == ap {
+			return
+		}
+	}
+	m.resolver[key] = append(v, ap)
 }
 
 func (m *MemoryRouter) Lookup(key string) ([]netip.AddrPort, bool) {
 	m.mx.RLock()
 	defer m.mx.RUnlock()
+
 	v, ok := m.resolver[key]
 	return v, ok
 }
diff --git a/pkg/state/state.go b/pkg/state/state.go
index 5df6abb..77e8ce0 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -94,8 +94,8 @@ func all(ctx context.Context, ociClient oci.Client, router routing.Router, resol
 func update(ctx context.Context, ociClient oci.Client, router routing.Router, event oci.ImageEvent, skipDigests, resolveLatestTag bool) (int, error) {
 	keys := []string{}
 	if !(!resolveLatestTag && event.Image.IsLatestTag()) {
-		if tagRef, ok := event.Image.TagName(); ok {
-			keys = append(keys, tagRef)
+		if tagName, ok := event.Image.TagName(); ok {
+			keys = append(keys, tagName)
 		}
 	}
 	if event.Type == oci.DeleteEvent {
@@ -107,7 +107,7 @@ func update(ctx context.Context, ociClient oci.Client, router routing.Router, ev
 		return 0, nil
 	}
 	if !skipDigests {
-		dgsts, err := ociClient.AllIdentifiers(ctx, event.Image)
+		dgsts, err := oci.WalkImage(ctx, ociClient, event.Image)
 		if err != nil {
 			return 0, fmt.Errorf("could not get digests for image %s: %w", event.Image.String(), err)
 		}


Test patch:
diff --git a/pkg/oci/containerd_test.go b/pkg/oci/containerd_test.go
index 7a6fcd0..e1b6097 100644
--- a/pkg/oci/containerd_test.go
+++ b/pkg/oci/containerd_test.go
@@ -9,7 +9,6 @@ import (
 
 	eventtypes "github.com/containerd/containerd/api/events"
 	"github.com/containerd/typeurl/v2"
-	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/require"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1"
@@ -193,44 +192,6 @@ func TestGetEventImage(t *testing.T) {
 	}
 }
 
-func TestIsImageConfig(t *testing.T) {
-	t.Parallel()
-
-	tests := []struct {
-		name     string
-		cImg     ocispec.Image
-		expected bool
-	}{
-		{
-			name: "image with required fields set",
-			cImg: ocispec.Image{
-				Platform: ocispec.Platform{
-					Architecture: "dummy",
-					OS:           "dummy",
-				},
-				RootFS: ocispec.RootFS{
-					Type: "dummy",
-				},
-			},
-			expected: true,
-		},
-		{
-			name:     "image with other fields setk",
-			cImg:     ocispec.Image{},
-			expected: false,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			t.Parallel()
-
-			ok := isImageConfig(tt.cImg)
-			require.Equal(t, tt.expected, ok)
-		})
-
-	}
-}
-
 func TestMirrorConfiguration(t *testing.T) {
 	t.Parallel()
 
diff --git a/pkg/oci/oci_test.go b/pkg/oci/oci_test.go
index 305c0ad..a700bf3 100644
--- a/pkg/oci/oci_test.go
+++ b/pkg/oci/oci_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"path"
+	"path/filepath"
 	"testing"
 
 	"github.com/containerd/containerd"
@@ -84,10 +85,33 @@ func TestOCIClient(t *testing.T) {
 		client:      containerdClient,
 	}
 
-	for _, ociClient := range []Client{remoteContainerd, localContainerd} {
+	memoryClient := NewMemory()
+	for _, img := range imgs {
+		dgst, err := digest.Parse(img["digest"])
+		require.NoError(t, err)
+		img, err := Parse(img["name"], dgst)
+		require.NoError(t, err)
+		memoryClient.AddImage(img)
+	}
+	for k, v := range blobs {
+		memoryClient.AddBlob(v, k)
+	}
+
+	for _, ociClient := range []Client{remoteContainerd, localContainerd, memoryClient} {
 		t.Run(ociClient.Name(), func(t *testing.T) {
 			t.Parallel()
 
+			b, mt, err := ociClient.GetManifest(ctx, digest.FromString("foo"))
+			require.Empty(t, b)
+			require.Empty(t, mt)
+			require.ErrorIs(t, err, ErrNotFound)
+			rc, err := ociClient.GetBlob(ctx, digest.FromString("foo"))
+			require.Empty(t, rc)
+			require.ErrorIs(t, err, ErrNotFound)
+			size, err := ociClient.Size(ctx, digest.FromString("foo"))
+			require.Empty(t, size)
+			require.ErrorIs(t, err, ErrNotFound)
+
 			imgs, err := ociClient.ListImages(ctx)
 			require.NoError(t, err)
 			require.Len(t, imgs, 5)
@@ -103,8 +127,8 @@ func TestOCIClient(t *testing.T) {
 				Name:   noPlatformName,
 				Digest: dgst,
 			}
-			_, err = ociClient.AllIdentifiers(ctx, img)
-			require.EqualError(t, err, "failed to walk image manifests: could not find any platforms with local content in manifest list: sha256:addc990c58744bdf96364fe89bd4aab38b1e824d51c688edb36c75247cd45fa9")
+			_, err = WalkImage(ctx, ociClient, img)
+			require.EqualError(t, err, "failed to walk image manifests: could not find any platforms with local content in manifest sha256:addc990c58744bdf96364fe89bd4aab38b1e824d51c688edb36c75247cd45fa9")
 
 			contentTests := []struct {
 				mediaType string
@@ -259,7 +283,7 @@ func TestOCIClient(t *testing.T) {
 
 					img, err := Parse(tt.imageName, digest.Digest(tt.imageDigest))
 					require.NoError(t, err)
-					keys, err := ociClient.AllIdentifiers(ctx, img)
+					keys, err := WalkImage(ctx, ociClient, img)
 					require.NoError(t, err)
 					require.Equal(t, tt.expectedKeys, keys)
 				})
@@ -267,3 +291,54 @@ func TestOCIClient(t *testing.T) {
 		})
 	}
 }
+
+func TestDetermineMediaType(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name              string
+		dgst              digest.Digest
+		expectedMediaType string
+	}{
+		{
+			name:              "image config",
+			dgst:              digest.Digest("sha256:68b8a989a3e08ddbdb3a0077d35c0d0e59c9ecf23d0634584def8bdbb7d6824f"),
+			expectedMediaType: ocispec.MediaTypeImageConfig,
+		},
+		{
+			name:              "image index",
+			dgst:              digest.Digest("sha256:9430beb291fa7b96997711fc486bc46133c719631aefdbeebe58dd3489217bfe"),
+			expectedMediaType: ocispec.MediaTypeImageIndex,
+		},
+		{
+			name:              "image index without media type",
+			dgst:              digest.Digest("sha256:d8df04365d06181f037251de953aca85cc16457581a8fc168f4957c978e1008b"),
+			expectedMediaType: ocispec.MediaTypeImageIndex,
+		},
+		{
+			name:              "image manifest",
+			dgst:              digest.Digest("sha256:dce623533c59af554b85f859e91fc1cbb7f574e873c82f36b9ea05a09feb0b53"),
+			expectedMediaType: ocispec.MediaTypeImageManifest,
+		},
+		{
+			name:              "image manifest without media type",
+			dgst:              digest.Digest("sha256:b6d6089ca6c395fd563c2084f5dd7bc56a2f5e6a81413558c5be0083287a77e9"),
+			expectedMediaType: ocispec.MediaTypeImageManifest,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			b, err := os.ReadFile(filepath.Join("testdata", "blobs", tt.dgst.Algorithm().String(), tt.dgst.Encoded()))
+			require.NoError(t, err)
+			mt, err := DetermineMediaType(b)
+			require.NoError(t, err)
+			require.Equal(t, tt.expectedMediaType, mt)
+		})
+	}
+
+	mt, err := DetermineMediaType([]byte("{}"))
+	require.EqualError(t, err, "not able to determine media type")
+	require.Empty(t, mt)
+}
diff --git a/pkg/oci/testdata/blobs/sha256/b6d6089ca6c395fd563c2084f5dd7bc56a2f5e6a81413558c5be0083287a77e9 b/pkg/oci/testdata/blobs/sha256/b6d6089ca6c395fd563c2084f5dd7bc56a2f5e6a81413558c5be0083287a77e9
new file mode 100644
index 0000000..715af06
--- /dev/null
+++ b/pkg/oci/testdata/blobs/sha256/b6d6089ca6c395fd563c2084f5dd7bc56a2f5e6a81413558c5be0083287a77e9
@@ -0,0 +1,15 @@
+{
+  "schemaVersion": 2,
+  "config": {
+    "mediaType": "application/vnd.oci.image.config.v1+json",
+    "digest": "sha256:68b8a989a3e08ddbdb3a0077d35c0d0e59c9ecf23d0634584def8bdbb7d6824f",
+    "size": 529
+  },
+  "layers": [
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:3caa2469de2a23cbcc209dd0b9d01cd78ff9a0f88741655991d36baede5b0996",
+      "size": 118
+    }
+  ]
+}
diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go
index 62771ee..835ad8f 100644
--- a/pkg/state/state_test.go
+++ b/pkg/state/state_test.go
@@ -2,18 +2,60 @@ package state
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/json"
+	"math/rand/v2"
 	"net/netip"
+	"strconv"
 	"testing"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
+	"github.com/go-logr/logr"
+	tlog "github.com/go-logr/logr/testing"
+	"github.com/opencontainers/go-digest"
+	"github.com/opencontainers/image-spec/specs-go"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/require"
 
 	"github.com/spegel-org/spegel/pkg/oci"
 	"github.com/spegel-org/spegel/pkg/routing"
 )
 
-func TestBasic(t *testing.T) {
+func TestTrack(t *testing.T) {
 	t.Parallel()
+	ociClient := oci.NewMemory()
+
+	imgRefs := []string{
+		"docker.io/library/ubuntu:latest",
+		"ghcr.io/spegel-org/spegel:v0.0.9",
+		"docker.io/library/alpine",
+	}
+	imgs := []oci.Image{}
+	for _, imageStr := range imgRefs {
+		manifest := ocispec.Manifest{
+			Versioned: specs.Versioned{
+				SchemaVersion: 2,
+			},
+			MediaType: ocispec.MediaTypeImageManifest,
+			Annotations: map[string]string{
+				"random": strconv.Itoa(rand.Int()),
+			},
+		}
+		b, err := json.Marshal(&manifest)
+		require.NoError(t, err)
+		hash := sha256.New()
+		_, err = hash.Write(b)
+		require.NoError(t, err)
+		dgst := digest.NewDigest(digest.SHA256, hash)
+		ociClient.AddBlob(b, dgst)
+		img, err := oci.Parse(imageStr, dgst)
+		require.NoError(t, err)
+		ociClient.AddImage(img)
+
+		imgs = append(imgs, img)
+	}
 
 	tests := []struct {
 		name             string
@@ -28,33 +70,20 @@ func TestBasic(t *testing.T) {
 			resolveLatestTag: false,
 		},
 	}
-
-	imgRefs := []string{
-		"docker.io/library/ubuntu:latest@sha256:b060fffe8e1561c9c3e6dea6db487b900100fc26830b9ea2ec966c151ab4c020",
-		"ghcr.io/spegel-org/spegel:v0.0.9@sha256:fa32bd3bcd49a45a62cfc1b0fed6a0b63bf8af95db5bad7ec22865aee0a4b795",
-		"docker.io/library/alpine@sha256:25fad2a32ad1f6f510e528448ae1ec69a28ef81916a004d3629874104f8a7f70",
-	}
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			imgs := []oci.Image{}
-			for _, imageStr := range imgRefs {
-				img, err := oci.Parse(imageStr, "")
-				require.NoError(t, err)
-				imgs = append(imgs, img)
-			}
-			ociClient := oci.NewMockClient(imgs)
-			router := routing.NewMemoryRouter(map[string][]netip.AddrPort{}, netip.MustParseAddrPort("127.0.0.1:5000"))
+			log := tlog.NewTestLogger(t)
+			ctx := logr.NewContext(context.Background(), log)
+			ctx, cancel := context.WithCancel(ctx)
 
-			ctx, cancel := context.WithCancel(context.TODO())
-			go func() {
-				time.Sleep(2 * time.Second)
-				cancel()
-			}()
-			err := Track(ctx, ociClient, router, tt.resolveLatestTag)
-			require.NoError(t, err)
+			router := routing.NewMemoryRouter(map[string][]netip.AddrPort{}, netip.MustParseAddrPort("127.0.0.1:5000"))
+			g, gCtx := errgroup.WithContext(ctx)
+			g.Go(func() error {
+				return Track(gCtx, ociClient, router, tt.resolveLatestTag)
+			})
+			time.Sleep(100 * time.Millisecond)
 
 			for _, img := range imgs {
 				peers, ok := router.Lookup(img.Digest.String())
@@ -72,6 +101,10 @@ func TestBasic(t *testing.T) {
 				require.True(t, ok)
 				require.Len(t, peers, 1)
 			}
+
+			cancel()
+			err := g.Wait()
+			require.NoError(t, err)
 		})
 	}
 }