Support Azure custom storage domains

Users of Azure Blob storage, particularly those using Azure China, Azure
Germany, or Azure Government, may need to customize the Azure
endpoint. To support this, we add a custom Azure URL opener that
supports the `domain` query parameter. This work has been submitted
upstream: https://github.com/google/go-cloud/pull/2851/files

Documentation for using Azure Blob storage:
https://gitlab.com/gitlab-org/gitlab/-/merge_requests/40693.

changelogs/unreleased/sh-support-azure-storage-domains-try2.yml

+---
+title: Support Azure custom storage domains
+merge_request: 593
+author:
+type: added

internal/config/config.go

@@ -129,10 +129,12 @@ func (c *Config) RegisterGoCloudURLOpeners() error {
 
 		pipeline := azureblob.NewPipeline(credential, azblob.PipelineOptions{})
 
-		azureURLOpener := &azureblob.URLOpener{
-			AccountName: accountName,
-			Pipeline:    pipeline,
-			Options:     azureblob.Options{Credential: credential},
+		azureURLOpener := &azureURLOpener{
+			&azureblob.URLOpener{
+				AccountName: accountName,
+				Pipeline:    pipeline,
+				Options:     azureblob.Options{Credential: credential},
+			},
 		}
 
 		c.ObjectStorageConfig.URLMux.RegisterBucket(azureblob.Scheme, azureURLOpener)

internal/config/config_test.go

@@ -8,6 +8,15 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const azureConfig = `
+[object_storage]
+provider = "AzureRM"
+
+[object_storage.azurerm]
+azure_storage_account_name = "azuretester"
+azure_storage_access_key = "deadbeef"
+`
+
 func TestLoadEmptyConfig(t *testing.T) {
 	config := ``
 
@@ -47,15 +56,7 @@ aws_secret_access_key = "gdk-minio"
 }
 
 func TestRegisterGoCloudURLOpeners(t *testing.T) {
-	config := `
-[object_storage]
-provider = "AzureRM"
-
-[object_storage.azurerm]
-azure_storage_account_name = "azuretester"
-azure_storage_access_key = "deadbeef"
-`
-	tmpFile, cfg := loadTempConfig(t, config)
+	tmpFile, cfg := loadTempConfig(t, azureConfig)
 	defer os.Remove(tmpFile.Name())
 
 	require.NotNil(t, cfg.ObjectStorageCredentials, "Expected object storage credentials")

internal/config/url_openers.go

+package config
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+
+	"gocloud.dev/blob"
+	"gocloud.dev/blob/azureblob"
+)
+
+// This code can be removed once https://github.com/google/go-cloud/pull/2851 is merged.
+
+// URLOpener opens Azure URLs like "azblob://mybucket".
+//
+// The URL host is used as the bucket name.
+//
+// The following query options are supported:
+//  - domain: The domain name used to access the Azure Blob storage (e.g. blob.core.windows.net)
+type azureURLOpener struct {
+	*azureblob.URLOpener
+}
+
+func (o *azureURLOpener) OpenBucketURL(ctx context.Context, u *url.URL) (*blob.Bucket, error) {
+	opts := new(azureblob.Options)
+	*opts = o.Options
+
+	err := setOptionsFromURLParams(u.Query(), opts)
+	if err != nil {
+		return nil, err
+	}
+	return azureblob.OpenBucket(ctx, o.Pipeline, o.AccountName, u.Host, opts)
+}
+
+func setOptionsFromURLParams(q url.Values, opts *azureblob.Options) error {
+	for param, values := range q {
+		if len(values) > 1 {
+			return fmt.Errorf("multiple values of %v not allowed", param)
+		}
+
+		value := values[0]
+		switch param {
+		case "domain":
+			opts.StorageDomain = azureblob.StorageDomain(value)
+		default:
+			return fmt.Errorf("unknown query parameter %q", param)
+		}
+	}
+
+	return nil
+}

internal/config/url_openers_test.go

+package config
+
+import (
+	"context"
+	"net/url"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gocloud.dev/blob/azureblob"
+)
+
+func TestURLOpeners(t *testing.T) {
+	tmpFile, cfg := loadTempConfig(t, azureConfig)
+	defer os.Remove(tmpFile.Name())
+
+	require.NotNil(t, cfg.ObjectStorageCredentials, "Expected object storage credentials")
+
+	err := cfg.RegisterGoCloudURLOpeners()
+	require.NoError(t, err)
+	require.NotNil(t, cfg.ObjectStorageConfig.URLMux)
+
+	tests := []struct {
+		url   string
+		valid bool
+	}{
+
+		{
+			url:   "azblob://container/object",
+			valid: true,
+		},
+		{
+			url:   "azblob://container/object?domain=core.windows.net",
+			valid: true,
+		},
+		{
+			url:   "azblob://container/object?domain=core.windows.net&domain=test",
+			valid: false,
+		},
+		{
+			url:   "azblob://container/object?param=value",
+			valid: false,
+		},
+		{
+			url:   "s3://bucket/object",
+			valid: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.url, func(t *testing.T) {
+			ctx := context.Background()
+			url, err := url.Parse(test.url)
+			require.NoError(t, err)
+
+			bucket, err := cfg.ObjectStorageConfig.URLMux.OpenBucketURL(ctx, url)
+			if bucket != nil {
+				defer bucket.Close()
+			}
+
+			if test.valid {
+				require.NotNil(t, bucket)
+				require.NoError(t, err)
+			} else {
+				require.Error(t, err)
+			}
+		})
+	}
+}
+
+func TestTestURLOpenersForParams(t *testing.T) {
+	tests := []struct {
+		name     string
+		currOpts azureblob.Options
+		query    url.Values
+		wantOpts azureblob.Options
+		wantErr  bool
+	}{
+		{
+			name: "InvalidParam",
+			query: url.Values{
+				"foo": {"bar"},
+			},
+			wantErr: true,
+		},
+		{
+			name: "StorageDomain",
+			query: url.Values{
+				"domain": {"blob.core.usgovcloudapi.net"},
+			},
+			wantOpts: azureblob.Options{StorageDomain: "blob.core.usgovcloudapi.net"},
+		},
+		{
+			name: "duplicate StorageDomain",
+			query: url.Values{
+				"domain": {"blob.core.usgovcloudapi.net", "blob.core.windows.net"},
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			o := &azureURLOpener{
+				URLOpener: &azureblob.URLOpener{
+					Options: test.currOpts,
+				},
+			}
+			err := setOptionsFromURLParams(test.query, &o.Options)
+
+			if test.wantErr {
+				assert.NotNil(t, err)
+			} else {
+				assert.Nil(t, err)
+				assert.Equal(t, test.wantOpts, o.Options)
+			}
+		})
+	}
+}