From 558ace139143057e0f8e3dbbcc5b695dda33462a Mon Sep 17 00:00:00 2001
From: Andrey Smirnov <andrey.smirnov@siderolabs.com>
Date: Tue, 31 Oct 2023 22:59:42 +0400
Subject: [PATCH] feat: implement 'rewrite' storage middleware
This allows to rewrite 'URLFor' of the storage driver to use a specific
host/trim the base path.
It is different from the 'redirect' middleware, as it still calls the
storage driver URLFor.
For example, with Azure storage provider, this allows to transform the
SAS Azure Blob Storage URL into the URL compatible with Azure Front
Door.
Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>
---
cmd/registry/main.go | 1 +
.../storage-drivers/middleware/_index.md | 15 ++++
.../storage-drivers/middleware/rewrite.md | 32 +++++++
.../driver/middleware/rewrite/middleware.go | 86 +++++++++++++++++++
.../middleware/rewrite/middleware_test.go | 75 ++++++++++++++++
5 files changed, 209 insertions(+)
create mode 100644 docs/content/storage-drivers/middleware/_index.md
create mode 100644 docs/content/storage-drivers/middleware/rewrite.md
create mode 100644 registry/storage/driver/middleware/rewrite/middleware.go
create mode 100644 registry/storage/driver/middleware/rewrite/middleware_test.go
diff --git a/cmd/registry/main.go b/cmd/registry/main.go
index de160f301..104eb0ecf 100644
--- a/cmd/registry/main.go
+++ b/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"
)
diff --git a/docs/content/storage-drivers/middleware/_index.md b/docs/content/storage-drivers/middleware/_index.md
new file mode 100644
index 000000000..979fe92c6
--- /dev/null
+++ b/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.
diff --git a/docs/content/storage-drivers/middleware/rewrite.md b/docs/content/storage-drivers/middleware/rewrite.md
new file mode 100644
index 000000000..cba2a05d7
--- /dev/null
+++ b/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
+```
diff --git a/registry/storage/driver/middleware/rewrite/middleware.go b/registry/storage/driver/middleware/rewrite/middleware.go
new file mode 100644
index 000000000..7baef518d
--- /dev/null
+++ b/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
+}
diff --git a/registry/storage/driver/middleware/rewrite/middleware_test.go b/registry/storage/driver/middleware/rewrite/middleware_test.go
new file mode 100644
index 000000000..651866cbc
--- /dev/null
+++ b/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)
+}