From 4b217ccbf566e6e38f2c62292ebd980d920fbcda Mon Sep 17 00:00:00 2001
From: Andrew Hsu <andrewhsu@acm.org>
Date: Thu, 21 Apr 2016 15:54:48 -0700
Subject: [PATCH] add middleware storage driver for redirect

Signed-off-by: Andrew Hsu <andrewhsu@acm.org> (github: andrewhsu)
---
 .../driver/middleware/redirect/middleware.go  | 47 ++++++++++++++
 .../middleware/redirect/middleware_test.go    | 62 +++++++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 registry/storage/driver/middleware/redirect/middleware.go
 create mode 100644 registry/storage/driver/middleware/redirect/middleware_test.go

diff --git a/registry/storage/driver/middleware/redirect/middleware.go b/registry/storage/driver/middleware/redirect/middleware.go
new file mode 100644
index 000000000..286a84abe
--- /dev/null
+++ b/registry/storage/driver/middleware/redirect/middleware.go
@@ -0,0 +1,47 @@
+package middleware
+
+import (
+	"fmt"
+	"github.com/docker/distribution/context"
+	storagedriver "github.com/docker/distribution/registry/storage/driver"
+	storagemiddleware "github.com/docker/distribution/registry/storage/driver/middleware"
+	"net/url"
+	"strings"
+)
+
+type redirectStorageMiddleware struct {
+	storagedriver.StorageDriver
+	scheme string
+	host   string
+}
+
+var _ storagedriver.StorageDriver = &redirectStorageMiddleware{}
+
+func newRedirectStorageMiddleware(sd storagedriver.StorageDriver, options map[string]interface{}) (storagedriver.StorageDriver, error) {
+	o, ok := options["baseurl"]
+	if !ok {
+		return nil, fmt.Errorf("no baseurl provided")
+	}
+	b, ok := o.(string)
+	if !ok {
+		return nil, fmt.Errorf("baseurl must be a string")
+	}
+	if !strings.Contains(b, "://") {
+		b = "https://" + b
+	}
+	u, err := url.Parse(b)
+	if err != nil {
+		return nil, fmt.Errorf("invalid baseurl: %v", err)
+	}
+
+	return &redirectStorageMiddleware{StorageDriver: sd, scheme: u.Scheme, host: u.Host}, nil
+}
+
+func (r *redirectStorageMiddleware) URLFor(ctx context.Context, path string, options map[string]interface{}) (string, error) {
+	u := &url.URL{Scheme: r.scheme, Host: r.host, Path: path}
+	return u.String(), nil
+}
+
+func init() {
+	storagemiddleware.Register("redirect", storagemiddleware.InitFunc(newRedirectStorageMiddleware))
+}
diff --git a/registry/storage/driver/middleware/redirect/middleware_test.go b/registry/storage/driver/middleware/redirect/middleware_test.go
new file mode 100644
index 000000000..31b661b63
--- /dev/null
+++ b/registry/storage/driver/middleware/redirect/middleware_test.go
@@ -0,0 +1,62 @@
+package middleware
+
+import (
+	check "gopkg.in/check.v1"
+	"testing"
+)
+
+func Test(t *testing.T) { check.TestingT(t) }
+
+type MiddlewareSuite struct{}
+
+var _ = check.Suite(&MiddlewareSuite{})
+
+func (s *MiddlewareSuite) TestNoConfig(c *check.C) {
+	options := make(map[string]interface{})
+	_, err := newRedirectStorageMiddleware(nil, options)
+	c.Assert(err, check.ErrorMatches, "no baseurl provided")
+}
+
+func (s *MiddlewareSuite) TestDefaultScheme(c *check.C) {
+	options := make(map[string]interface{})
+	options["baseurl"] = "example.com"
+	middleware, err := newRedirectStorageMiddleware(nil, options)
+	c.Assert(err, check.Equals, nil)
+
+	m, ok := middleware.(*redirectStorageMiddleware)
+	c.Assert(ok, check.Equals, true)
+	c.Assert(m.scheme, check.Equals, "https")
+	c.Assert(m.host, check.Equals, "example.com")
+}
+
+func (s *MiddlewareSuite) TestHTTPS(c *check.C) {
+	options := make(map[string]interface{})
+	options["baseurl"] = "https://example.com"
+	middleware, err := newRedirectStorageMiddleware(nil, options)
+	c.Assert(err, check.Equals, nil)
+
+	m, ok := middleware.(*redirectStorageMiddleware)
+	c.Assert(ok, check.Equals, true)
+	c.Assert(m.scheme, check.Equals, "https")
+	c.Assert(m.host, check.Equals, "example.com")
+
+	url, err := middleware.URLFor(nil, "/rick/data", nil)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(url, check.Equals, "https://example.com/rick/data")
+}
+
+func (s *MiddlewareSuite) TestHTTP(c *check.C) {
+	options := make(map[string]interface{})
+	options["baseurl"] = "http://example.com"
+	middleware, err := newRedirectStorageMiddleware(nil, options)
+	c.Assert(err, check.Equals, nil)
+
+	m, ok := middleware.(*redirectStorageMiddleware)
+	c.Assert(ok, check.Equals, true)
+	c.Assert(m.scheme, check.Equals, "http")
+	c.Assert(m.host, check.Equals, "example.com")
+
+	url, err := middleware.URLFor(nil, "morty/data", nil)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(url, check.Equals, "http://example.com/morty/data")
+}