feat: implement 'rewrite' storage middleware (#4146)

cmd/registry/main.go

@@ -14,6 +14,7 @@ import (
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/inmemory"
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/middleware/cloudfront"
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/middleware/redirect"
+	_ "github.com/distribution/distribution/v3/registry/storage/driver/middleware/rewrite"
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/s3-aws"
 )
 

docs/content/storage-drivers/middleware/_index.md

@@ -0,0 +1,15 @@
+---
+description: Explains how to use storage middleware
+keywords: registry, on-prem, images, tags, repository, distribution, storage drivers, advanced
+title: Storage middleware
+---
+
+This document describes the registry storage middleware.
+
+## Provided middleware
+
+This storage driver package comes bundled with several middleware options:
+
+- cloudfront
+- redirect
+- [rewrite](rewrite): Partially rewrites the URL returned by the storage driver.

docs/content/storage-drivers/middleware/rewrite.md

@@ -0,0 +1,32 @@
+---
+description: Explains how to use the rewrite storage middleware
+keywords: registry, service, driver, images, storage, middleware, rewrite
+title: Rewrite middleware
+---
+
+A storage middleware which allows to rewrite the URL returned by the storage driver.
+
+For example, it can be used to rewrite the Blob Storage URL returned by the Azure Blob Storage driver to use Azure CDN.
+
+## Parameters
+
+* `scheme`: (optional): Rewrite the returned URL scheme (if set).
+* `host`: (optional): Rewrite the returned URL host (if set).
+* `trimpathprefix` (optional): Trim the prefix from the returned URL path (if set).
+
+## Example configuration
+
+```yaml
+storage:
+  azure:
+    accountname: "ACCOUNT_NAME"
+    accountkey: "******"
+    container: container-name
+middleware:
+  storage:
+    - name: rewrite
+      options:
+        scheme: https
+        host: example-cdn-endpoint.azurefd.net
+        trimpathprefix: /container-name
+```

registry/storage/driver/middleware/rewrite/middleware.go

@@ -0,0 +1,86 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+
+	storagedriver "github.com/distribution/distribution/v3/registry/storage/driver"
+	storagemiddleware "github.com/distribution/distribution/v3/registry/storage/driver/middleware"
+	"github.com/sirupsen/logrus"
+)
+
+func init() {
+	if err := storagemiddleware.Register("rewrite", newRewriteStorageMiddleware); err != nil {
+		logrus.Errorf("tailed to register redirect storage middleware: %v", err)
+	}
+}
+
+type rewriteStorageMiddleware struct {
+	storagedriver.StorageDriver
+	overrideScheme string
+	overrideHost   string
+	trimPathPrefix string
+}
+
+var _ storagedriver.StorageDriver = &rewriteStorageMiddleware{}
+
+func getStringOption(key string, options map[string]interface{}) (string, error) {
+	o, ok := options[key]
+	if !ok {
+		return "", nil
+	}
+	s, ok := o.(string)
+	if !ok {
+		return "", fmt.Errorf("%s must be a string", key)
+	}
+	return s, nil
+}
+
+func newRewriteStorageMiddleware(ctx context.Context, sd storagedriver.StorageDriver, options map[string]interface{}) (storagedriver.StorageDriver, error) {
+	var err error
+
+	r := &rewriteStorageMiddleware{StorageDriver: sd}
+
+	if r.overrideScheme, err = getStringOption("scheme", options); err != nil {
+		return nil, err
+	}
+
+	if r.overrideHost, err = getStringOption("host", options); err != nil {
+		return nil, err
+	}
+
+	if r.trimPathPrefix, err = getStringOption("trimpathprefix", options); err != nil {
+		return nil, err
+	}
+
+	return r, nil
+}
+
+func (r *rewriteStorageMiddleware) RedirectURL(req *http.Request, path string) (string, error) {
+	storagePath, err := r.StorageDriver.RedirectURL(req, path)
+	if err != nil {
+		return "", err
+	}
+
+	u, err := url.Parse(storagePath)
+	if err != nil {
+		return "", err
+	}
+
+	if r.overrideScheme != "" {
+		u.Scheme = r.overrideScheme
+	}
+
+	if r.overrideHost != "" {
+		u.Host = r.overrideHost
+	}
+
+	if r.trimPathPrefix != "" {
+		u.Path = strings.TrimPrefix(u.Path, r.trimPathPrefix)
+	}
+
+	return u.String(), nil
+}

registry/storage/driver/middleware/rewrite/middleware_test.go

@@ -0,0 +1,75 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/distribution/distribution/v3/registry/storage/driver/base"
+	"github.com/stretchr/testify/require"
+)
+
+type mockSD struct {
+	base.Base
+}
+
+func (*mockSD) RedirectURL(_ *http.Request, urlPath string) (string, error) {
+	return "http://some.host/some/path/file", nil
+}
+
+func TestNoConfig(t *testing.T) {
+	options := make(map[string]interface{})
+	middleware, err := newRewriteStorageMiddleware(context.Background(), &mockSD{}, options)
+	require.NoError(t, err)
+
+	_, ok := middleware.(*rewriteStorageMiddleware)
+	require.True(t, ok)
+
+	url, err := middleware.RedirectURL(nil, "")
+	require.NoError(t, err)
+	require.Equal(t, "http://some.host/some/path/file", url)
+}
+
+func TestWrongType(t *testing.T) {
+	options := map[string]interface{}{
+		"scheme": 1,
+	}
+	_, err := newRewriteStorageMiddleware(context.TODO(), nil, options)
+	require.ErrorContains(t, err, "scheme must be a string")
+}
+
+func TestRewriteHostsScheme(t *testing.T) {
+	options := map[string]interface{}{
+		"scheme": "https",
+		"host":   "example.com",
+	}
+
+	middleware, err := newRewriteStorageMiddleware(context.TODO(), &mockSD{}, options)
+	require.NoError(t, err)
+
+	m, ok := middleware.(*rewriteStorageMiddleware)
+	require.True(t, ok)
+	require.Equal(t, "https", m.overrideScheme)
+	require.Equal(t, "example.com", m.overrideHost)
+
+	url, err := middleware.RedirectURL(nil, "")
+	require.NoError(t, err)
+	require.Equal(t, "https://example.com/some/path/file", url)
+}
+
+func TestTrimPrefix(t *testing.T) {
+	options := map[string]interface{}{
+		"trimpathprefix": "/some/path",
+	}
+
+	middleware, err := newRewriteStorageMiddleware(context.TODO(), &mockSD{}, options)
+	require.NoError(t, err)
+
+	m, ok := middleware.(*rewriteStorageMiddleware)
+	require.True(t, ok)
+	require.Equal(t, "/some/path", m.trimPathPrefix)
+
+	url, err := middleware.RedirectURL(nil, "")
+	require.NoError(t, err)
+	require.Equal(t, "http://some.host/file", url)
+}