Merge pull request #90 from stevvooe/registry-decorator

Implement registry decorator toolkit

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)
+	}
+}