Implement registry decorator toolkit

This change provides a toolkit for intercepting registry calls, such as
`ManifestService.Get` and `LayerUpload.Finish`, with the goal of easily
supporting interesting callbacks and listeners. The package proxies
returned objects through the decorate function before creation, allowing one to
carefully choose injection points.

Use cases range from notification systems all the way to cache integration.
While such a tool isn't strictly necessary, it reduces the amount of code
required to accomplish such tasks, deferring the tricky aspects to the
decorator package.

Signed-off-by: Stephen J Day <stephen.day@docker.com>

storage/decorator/decorator.go

@@ -0,0 +1,185 @@
+package decorator
+
+import (
+	"github.com/docker/distribution/digest"
+	"github.com/docker/distribution/storage"
+)
+
+// Decorator provides an interface for intercepting object creation within a
+// registry. The single method accepts an registry storage object, such as a
+// Layer, optionally replacing it upon with an alternative object or a
+// wrapper.
+//
+// For example, if one wants to intercept the instantiation of a layer, an
+// implementation might be as follows:
+//
+// 	func (md *DecoratorImplementation) Decorate(v interface{}) interface{} {
+// 		switch v := v.(type) {
+// 		case Layer:
+// 			return wrapLayer(v)
+// 		}
+//
+//		// Make sure to return the object or nil if the decorator doesn't require
+//		// replacement.
+//		return v
+// 	}
+//
+// Such a decorator can be used to intercept calls to support implementing
+// complex features outside of the storage package.
+type Decorator interface {
+	Decorate(v interface{}) interface{}
+}
+
+// Func provides a shortcut handler for decorators that only need a
+// function. Use is similar to http.HandlerFunc.
+type Func func(v interface{}) interface{}
+
+// Decorate allows DecoratorFunc to implement the Decorator interface.
+func (df Func) Decorate(v interface{}) interface{} {
+	return df(v)
+}
+
+// DecorateRegistry the provided registry with decorator. Registries may be
+// decorated multiple times.
+func DecorateRegistry(registry storage.Registry, decorator Decorator) storage.Registry {
+	return &registryDecorator{
+		Registry:  registry,
+		decorator: decorator,
+	}
+}
+
+// registryDecorator intercepts registry object creation with a decorator.
+type registryDecorator struct {
+	storage.Registry
+	decorator Decorator
+}
+
+// Repository overrides the method of the same name on the Registry, replacing
+// the returned instance with a decorator.
+func (rd *registryDecorator) Repository(name string) storage.Repository {
+	delegate := rd.Registry.Repository(name)
+	decorated := rd.decorator.Decorate(delegate)
+	if decorated != nil {
+		repository, ok := decorated.(storage.Repository)
+
+		if ok {
+			delegate = repository
+		}
+	}
+
+	return &repositoryDecorator{
+		Repository: delegate,
+		decorator:  rd.decorator,
+	}
+}
+
+// repositoryDecorator decorates a repository, intercepting calls to Layers
+// and Manifests with injected variants.
+type repositoryDecorator struct {
+	storage.Repository
+	decorator Decorator
+}
+
+// Layers overrides the Layers method of Repository.
+func (rd *repositoryDecorator) Layers() storage.LayerService {
+	delegate := rd.Repository.Layers()
+	decorated := rd.decorator.Decorate(delegate)
+
+	if decorated != nil {
+		layers, ok := decorated.(storage.LayerService)
+
+		if ok {
+			delegate = layers
+		}
+	}
+
+	return &layerServiceDecorator{
+		LayerService: delegate,
+		decorator:    rd.decorator,
+	}
+}
+
+// Manifests overrides the Manifests method of Repository.
+func (rd *repositoryDecorator) Manifests() storage.ManifestService {
+	delegate := rd.Repository.Manifests()
+	decorated := rd.decorator.Decorate(delegate)
+
+	if decorated != nil {
+		manifests, ok := decorated.(storage.ManifestService)
+
+		if ok {
+			delegate = manifests
+		}
+	}
+
+	// NOTE(stevvooe): We do not have to intercept delegate calls to the
+	// manifest service since it doesn't produce any interfaces for which
+	// interception is supported.
+	return delegate
+}
+
+// layerServiceDecorator intercepts calls that generate Layer and LayerUpload
+// instances, replacing them with instances from the decorator.
+type layerServiceDecorator struct {
+	storage.LayerService
+	decorator Decorator
+}
+
+// Fetch overrides the Fetch method of LayerService.
+func (lsd *layerServiceDecorator) Fetch(digest digest.Digest) (storage.Layer, error) {
+	delegate, err := lsd.LayerService.Fetch(digest)
+	return decorateLayer(lsd.decorator, delegate), err
+}
+
+// Upload overrides the Upload method of LayerService.
+func (lsd *layerServiceDecorator) Upload() (storage.LayerUpload, error) {
+	delegate, err := lsd.LayerService.Upload()
+	return decorateLayerUpload(lsd.decorator, delegate), err
+}
+
+// Resume overrides the Resume method of LayerService.
+func (lsd *layerServiceDecorator) Resume(uuid string) (storage.LayerUpload, error) {
+	delegate, err := lsd.LayerService.Resume(uuid)
+	return decorateLayerUpload(lsd.decorator, delegate), err
+}
+
+// layerUploadDecorator intercepts calls that generate Layer instances,
+// replacing them with instances from the decorator.
+type layerUploadDecorator struct {
+	storage.LayerUpload
+	decorator Decorator
+}
+
+func (lud *layerUploadDecorator) Finish(dgst digest.Digest) (storage.Layer, error) {
+	delegate, err := lud.LayerUpload.Finish(dgst)
+	return decorateLayer(lud.decorator, delegate), err
+}
+
+// decorateLayer guarantees that a layer gets correctly decorated.
+func decorateLayer(decorator Decorator, delegate storage.Layer) storage.Layer {
+	decorated := decorator.Decorate(delegate)
+	if decorated != nil {
+		layer, ok := decorated.(storage.Layer)
+		if ok {
+			delegate = layer
+		}
+	}
+
+	return delegate
+}
+
+// decorateLayerUpload guarantees that an upload gets correctly decorated.
+func decorateLayerUpload(decorator Decorator, delegate storage.LayerUpload) storage.LayerUpload {
+	decorated := decorator.Decorate(delegate)
+	if decorated != nil {
+		layerUpload, ok := decorated.(storage.LayerUpload)
+		if ok {
+			delegate = layerUpload
+		}
+	}
+
+	return &layerUploadDecorator{
+		LayerUpload: delegate,
+		decorator:   decorator,
+	}
+}

storage/decorator/decorator_test.go

@@ -0,0 +1,138 @@
+package decorator
+
+import (
+	"io"
+	"testing"
+
+	"github.com/docker/libtrust"
+
+	"github.com/docker/distribution/digest"
+	"github.com/docker/distribution/manifest"
+	"github.com/docker/distribution/storage"
+	"github.com/docker/distribution/storagedriver/inmemory"
+	"github.com/docker/distribution/testutil"
+)
+
+func TestRegistryDecorator(t *testing.T) {
+	// Initialize the expected decorations. Call counting is a horrible way to
+	// test this but should keep this code from being atrocious.
+	expected := map[string]int{
+		"repository":      1,
+		"manifestservice": 1,
+		"layerservice":    1,
+		"layer":           4,
+		"layerupload":     4,
+	}
+	decorated := map[string]int{}
+
+	decorator := Func(func(v interface{}) interface{} {
+		switch v := v.(type) {
+		case storage.Repository:
+			t.Logf("decorate repository: %T", v)
+			decorated["repository"]++
+		case storage.ManifestService:
+			t.Logf("decorate manifestservice: %T", v)
+			decorated["manifestservice"]++
+		case storage.LayerService:
+			t.Logf("decorate layerservice: %T", v)
+			decorated["layerservice"]++
+		case storage.Layer:
+			t.Logf("decorate layer: %T", v)
+			decorated["layer"]++
+		case storage.LayerUpload:
+			t.Logf("decorate layerupload: %T", v)
+			decorated["layerupload"]++
+		default:
+			t.Fatalf("unexpected object decorated: %v", v)
+		}
+
+		return v
+	})
+
+	registry := storage.NewRegistryWithDriver(inmemory.New())
+	registry = DecorateRegistry(registry, decorator)
+
+	// Now take the registry through a number of operations
+	checkExerciseRegistry(t, registry)
+
+	for component, calls := range expected {
+		if decorated[component] != calls {
+			t.Fatalf("%v was not decorated expected number of times: %d != %d", component, decorated[component], calls)
+		}
+	}
+
+}
+
+// checkExerciseRegistry takes the registry through all of its operations,
+// carrying out generic checks.
+func checkExerciseRegistry(t *testing.T, registry storage.Registry) {
+	name := "foo/bar"
+	tag := "thetag"
+	repository := registry.Repository(name)
+	m := manifest.Manifest{
+		Versioned: manifest.Versioned{
+			SchemaVersion: 1,
+		},
+		Name: name,
+		Tag:  tag,
+	}
+
+	layers := repository.Layers()
+	for i := 0; i < 2; i++ {
+		rs, ds, err := testutil.CreateRandomTarFile()
+		if err != nil {
+			t.Fatalf("error creating test layer: %v", err)
+		}
+		dgst := digest.Digest(ds)
+		upload, err := layers.Upload()
+		if err != nil {
+			t.Fatalf("error creating layer upload: %v", err)
+		}
+
+		// Use the resumes, as well!
+		upload, err = layers.Resume(upload.UUID())
+		if err != nil {
+			t.Fatalf("error resuming layer upload: %v", err)
+		}
+
+		io.Copy(upload, rs)
+
+		if _, err := upload.Finish(dgst); err != nil {
+			t.Fatalf("unexpected error finishing upload: %v", err)
+		}
+
+		m.FSLayers = append(m.FSLayers, manifest.FSLayer{
+			BlobSum: dgst,
+		})
+
+		// Then fetch the layers
+		if _, err := layers.Fetch(dgst); err != nil {
+			t.Fatalf("error fetching layer: %v", err)
+		}
+	}
+
+	pk, err := libtrust.GenerateECP256PrivateKey()
+	if err != nil {
+		t.Fatalf("unexpected error generating key: %v", err)
+	}
+
+	sm, err := manifest.Sign(&m, pk)
+	if err != nil {
+		t.Fatalf("unexpected error signing manifest: %v", err)
+	}
+
+	manifests := repository.Manifests()
+
+	if err := manifests.Put(tag, sm); err != nil {
+		t.Fatalf("unexpected error putting the manifest: %v", err)
+	}
+
+	fetched, err := manifests.Get(tag)
+	if err != nil {
+		t.Fatalf("unexpected error fetching manifest: %v", err)
+	}
+
+	if fetched.Tag != fetched.Tag {
+		t.Fatalf("retrieved unexpected manifest: %v", err)
+	}
+}