Make the blob storage in pkg/registry pluggable.

pkg/registry/blobs.go

@@ -18,13 +18,18 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"math/rand"
 	"net/http"
 	"path"
 	"strings"
 	"sync"
+
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 )
 
 // Returns whether this url should be handled by the blob handler
@@ -46,11 +51,26 @@ func isBlob(req *http.Request) bool {
 
 // blobs
 type blobs struct {
-	// Blobs are content addresses. we store them globally underneath their sha and make no distinctions per image.
-	contents map[string][]byte
 	// Each upload gets a unique id that writes occur to until finalized.
 	uploads map[string][]byte
 	lock    sync.Mutex
+
+	bh BlobHandler
+}
+
+// BlobHandler is the interface for the storage layer underneath this registry.
+type BlobHandler interface {
+	// Stat returns the size of the blob whose hash is specified,
+	// if it exists. If not, it returns (0, error).
+	Stat(repo name.Repository, h v1.Hash) (int64, error)
+
+	// Get returns true and a reader for consuming the blob specified with the hash,
+	// if it exists.  It now, it returns (nil, error).
+	Get(repo name.Repository, h v1.Hash) (io.ReadCloser, error)
+
+	// Store stores the stream of content with the given hash, or returns the error
+	// encountered doing so.
+	Store(repo name.Repository, h v1.Hash, content io.ReadCloser) error
 }
 
 func (b *blobs) handle(resp http.ResponseWriter, req *http.Request) *regError {
@@ -74,38 +94,96 @@ func (b *blobs) handle(resp http.ResponseWriter, req *http.Request) *regError {
 
 	switch req.Method {
 	case http.MethodHead:
-		b.lock.Lock()
-		defer b.lock.Unlock()
-		b, ok := b.contents[target]
-		if !ok {
+		h, err := v1.NewHash(target)
+		if err != nil {
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+		repo, err := name.NewRepository(req.URL.Host + path.Join(elem[1:len(elem)-2]...))
+		if err != nil {
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+
+		sz, err := b.bh.Stat(repo, h)
+		if errors.Is(err, ErrBlobNotFound) {
 			return &regError{
 				Status:  http.StatusNotFound,
 				Code:    "BLOB_UNKNOWN",
-				Message: "Unknown blob",
+				Message: err.Error(),
+			}
+		} else if err != nil {
+			return &regError{
+				Status:  http.StatusInternalServerError,
+				Code:    "BLOB_UNKNOWN",
+				Message: err.Error(),
 			}
 		}
 
-		resp.Header().Set("Content-Length", fmt.Sprint(len(b)))
+		resp.Header().Set("Content-Length", fmt.Sprint(sz))
 		resp.Header().Set("Docker-Content-Digest", target)
 		resp.WriteHeader(http.StatusOK)
 		return nil
 
 	case http.MethodGet:
-		b.lock.Lock()
-		defer b.lock.Unlock()
-		b, ok := b.contents[target]
-		if !ok {
+		h, err := v1.NewHash(target)
+		if err != nil {
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+		repo, err := name.NewRepository(req.URL.Host + path.Join(elem[1:len(elem)-2]...))
+		if err != nil {
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+
+		sz, err := b.bh.Stat(repo, h)
+		if errors.Is(err, ErrBlobNotFound) {
 			return &regError{
 				Status:  http.StatusNotFound,
 				Code:    "BLOB_UNKNOWN",
-				Message: "Unknown blob",
+				Message: err.Error(),
+			}
+		} else if err != nil {
+			return &regError{
+				Status:  http.StatusInternalServerError,
+				Code:    "BLOB_UNKNOWN",
+				Message: err.Error(),
+			}
+		}
+
+		b, err := b.bh.Get(repo, h)
+		if errors.Is(err, ErrBlobNotFound) {
+			return &regError{
+				Status:  http.StatusNotFound,
+				Code:    "BLOB_UNKNOWN",
+				Message: err.Error(),
+			}
+		} else if err != nil {
+			return &regError{
+				Status:  http.StatusInternalServerError,
+				Code:    "BLOB_UNKNOWN",
+				Message: err.Error(),
 			}
 		}
+		defer b.Close()
 
-		resp.Header().Set("Content-Length", fmt.Sprint(len(b)))
+		resp.Header().Set("Content-Length", fmt.Sprint(sz))
 		resp.Header().Set("Docker-Content-Digest", target)
 		resp.WriteHeader(http.StatusOK)
-		io.Copy(resp, bytes.NewReader(b))
+		io.Copy(resp, b)
 		return nil
 
 	case http.MethodPost:
@@ -131,10 +209,31 @@ func (b *blobs) handle(resp http.ResponseWriter, req *http.Request) *regError {
 					Message: "digest does not match contents",
 				}
 			}
+			h, err := v1.NewHash(d)
+			if err != nil {
+				// This is not reachable
+				return &regError{
+					Status:  http.StatusBadRequest,
+					Code:    "NAME_INVALID",
+					Message: err.Error(),
+				}
+			}
+			repo, err := name.NewRepository(req.URL.Host + path.Join(elem[1:len(elem)-2]...))
+			if err != nil {
+				return &regError{
+					Status:  http.StatusBadRequest,
+					Code:    "NAME_INVALID",
+					Message: err.Error(),
+				}
+			}
 
-			b.lock.Lock()
-			defer b.lock.Unlock()
-			b.contents[d] = l.Bytes()
+			if err := b.bh.Store(repo, h, ioutil.NopCloser(l)); err != nil {
+				return &regError{
+					Status:  http.StatusInternalServerError,
+					Code:    "BLOB_UPLOAD_INVALID",
+					Message: err.Error(),
+				}
+			}
 			resp.Header().Set("Docker-Content-Digest", d)
 			resp.WriteHeader(http.StatusCreated)
 			return nil
@@ -231,8 +330,32 @@ func (b *blobs) handle(resp http.ResponseWriter, req *http.Request) *regError {
 				Message: "digest does not match contents",
 			}
 		}
+		repo, err := name.NewRepository(req.URL.Host + path.Join(elem[1:len(elem)-3]...))
+		if err != nil {
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+		h, err := v1.NewHash(digest)
+		if err != nil {
+			// This is not reachable
+			return &regError{
+				Status:  http.StatusBadRequest,
+				Code:    "NAME_INVALID",
+				Message: err.Error(),
+			}
+		}
+
+		if err := b.bh.Store(repo, h, ioutil.NopCloser(l)); err != nil {
+			return &regError{
+				Status:  http.StatusInternalServerError,
+				Code:    "BLOB_UPLOAD_INVALID",
+				Message: err.Error(),
+			}
+		}
 
-		b.contents[d] = l.Bytes()
 		delete(b.uploads, target)
 		resp.Header().Set("Docker-Content-Digest", d)
 		resp.WriteHeader(http.StatusCreated)
@@ -246,3 +369,49 @@ func (b *blobs) handle(resp http.ResponseWriter, req *http.Request) *regError {
 		}
 	}
 }
+
+type defaultBlobStore struct {
+	m        sync.Mutex
+	contents map[v1.Hash][]byte
+}
+
+var _ BlobHandler = (*defaultBlobStore)(nil)
+
+// ErrBlobNotFound is the error returned (possibly wrapped)
+// when a given blob is not found.
+var ErrBlobNotFound = errors.New("blob not found")
+
+// Stat implements BlobHandler
+func (dbs *defaultBlobStore) Stat(repo name.Repository, h v1.Hash) (int64, error) {
+	dbs.m.Lock()
+	defer dbs.m.Unlock()
+	b, ok := dbs.contents[h]
+	if !ok {
+		return 0, fmt.Errorf("%w: %s", ErrBlobNotFound, h)
+	}
+	return int64(len(b)), nil
+}
+
+// Get implements BlobHandler
+func (dbs *defaultBlobStore) Get(repo name.Repository, h v1.Hash) (io.ReadCloser, error) {
+	dbs.m.Lock()
+	defer dbs.m.Unlock()
+	b, ok := dbs.contents[h]
+	if !ok {
+		return nil, fmt.Errorf("%w: %s", ErrBlobNotFound, h)
+	}
+	return ioutil.NopCloser(bytes.NewBuffer(b)), nil
+}
+
+// Store implements BlobHandler
+func (dbs *defaultBlobStore) Store(repo name.Repository, h v1.Hash, rc io.ReadCloser) error {
+	dbs.m.Lock()
+	defer dbs.m.Unlock()
+	defer rc.Close()
+	b, err := ioutil.ReadAll(rc)
+	if err != nil {
+		return err
+	}
+	dbs.contents[h] = b
+	return nil
+}

pkg/registry/example_test.go

@@ -24,7 +24,7 @@ import (
 func Example() {
 	s := httptest.NewServer(registry.New())
 	defer s.Close()
-	resp, _ := s.Client().Get(s.URL + "/v2/bar/blobs/sha256:...")
+	resp, _ := s.Client().Get(s.URL + "/v2/bar/blobs/sha256:2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae")
 	fmt.Println(resp.StatusCode)
 	// Output: 404
 }

pkg/registry/registry.go

@@ -27,6 +27,8 @@ import (
 	"log"
 	"net/http"
 	"os"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 )
 
 type registry struct {
@@ -77,8 +79,10 @@ func New(opts ...Option) http.Handler {
 	r := &registry{
 		log: log.New(os.Stderr, "", log.LstdFlags),
 		blobs: blobs{
-			contents: map[string][]byte{},
-			uploads:  map[string][]byte{},
+			uploads: map[string][]byte{},
+			bh: &defaultBlobStore{
+				contents: map[v1.Hash][]byte{},
+			},
 		},
 		manifests: manifests{
 			manifests: map[string]map[string]manifest{},
@@ -102,3 +106,10 @@ func Logger(l *log.Logger) Option {
 		r.manifests.log = l
 	}
 }
+
+// WithBlobHandler overrides the default BlobHandler.
+func WithBlobHandler(bh BlobHandler) Option {
+	return func(r *registry) {
+		r.blobs.bh = bh
+	}
+}