From e5fc12042b4a348ee660dee0e1385c2a96c3db86 Mon Sep 17 00:00:00 2001
From: jonaustin09 <jonaustin460@gmail.com>
Date: Thu, 4 Jan 2024 13:16:05 -0500
Subject: [PATCH] feat: Added sas token authentication for azure backend

---
 backend/azure/azure.go | 39 +++++++++++++++++++++++++++++++++++----
 cmd/versitygw/azure.go | 11 +++++++++--
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/backend/azure/azure.go b/backend/azure/azure.go
index fe563044..f32831ec 100644
--- a/backend/azure/azure.go
+++ b/backend/azure/azure.go
@@ -49,11 +49,22 @@ type Azure struct {
 	client     *azblob.Client
 	creds      *azblob.SharedKeyCredential
 	serviceURL string
+	sasToken   string
 }
 
 var _ backend.Backend = &Azure{}
 
-func New(accountName, accountKey, serviceURL string) (*Azure, error) {
+func New(accountName, accountKey, serviceURL, sasToken string) (*Azure, error) {
+	if sasToken != "" && (accountName != "" || accountKey != "") {
+		return nil, fmt.Errorf("choose one of the authentication methods: with sas token or with access key/name")
+	}
+	if sasToken != "" {
+		client, err := azblob.NewClientWithNoCredential(serviceURL+"?"+sasToken, nil)
+		if err != nil {
+			return nil, fmt.Errorf("init client: %w", err)
+		}
+		return &Azure{client: client, serviceURL: serviceURL, sasToken: sasToken}, nil
+	}
 	cred, err := azblob.NewSharedKeyCredential(accountName, accountKey)
 	if err != nil {
 		return nil, fmt.Errorf("init credentials: %w", err)
@@ -645,16 +656,36 @@ func (az *Azure) GetBucketAcl(ctx context.Context, input *s3.GetBucketAclInput)
 	return []byte(*aclPtr), nil
 }
 
+func (az *Azure) getContainerURL(container string) string {
+	return fmt.Sprintf("%v/%v", az.serviceURL, container)
+}
+
+func (az *Azure) getBlobURL(container, blob string) string {
+	return az.getContainerURL(container) + "/" + blob
+}
+
 func (az *Azure) getBlobClient(container, blb string) (*blob.Client, error) {
-	return blob.NewClientWithSharedKeyCredential(fmt.Sprintf("%v/%v/%v", az.serviceURL, container, blb), az.creds, nil)
+	blobURL := az.getBlobURL(container, blb)
+	if az.sasToken != "" {
+		return blob.NewClientWithNoCredential(blobURL+"?"+az.sasToken, nil)
+	}
+	return blob.NewClientWithSharedKeyCredential(blobURL, az.creds, nil)
 }
 
 func (az *Azure) getContainerClient(ctr string) (*container.Client, error) {
-	return container.NewClientWithSharedKeyCredential(fmt.Sprintf("%v/%v", az.serviceURL, ctr), az.creds, nil)
+	containerURL := az.getContainerURL(ctr)
+	if az.sasToken != "" {
+		return container.NewClientWithNoCredential(containerURL+"?"+az.sasToken, nil)
+	}
+	return container.NewClientWithSharedKeyCredential(containerURL, az.creds, nil)
 }
 
 func (az *Azure) getBlockBlobClient(container, blob string) (*blockblob.Client, error) {
-	return blockblob.NewClientWithSharedKeyCredential(fmt.Sprintf("%v/%v/%v", az.serviceURL, container, blob), az.creds, nil)
+	blobURL := az.getBlobURL(container, blob)
+	if az.sasToken != "" {
+		return blockblob.NewClientWithNoCredential(blobURL+"?"+az.sasToken, nil)
+	}
+	return blockblob.NewClientWithSharedKeyCredential(blobURL, az.creds, nil)
 }
 
 func parseMetadata(m map[string]string) map[string]*string {
diff --git a/cmd/versitygw/azure.go b/cmd/versitygw/azure.go
index 0fa0e0a1..ccd950f0 100644
--- a/cmd/versitygw/azure.go
+++ b/cmd/versitygw/azure.go
@@ -22,7 +22,7 @@ import (
 )
 
 var (
-	azAccount, azKey, azServiceURL string
+	azAccount, azKey, azServiceURL, azSASToken string
 )
 
 func azureCommand() *cli.Command {
@@ -46,6 +46,13 @@ func azureCommand() *cli.Command {
 				Aliases:     []string{"k"},
 				Destination: &azKey,
 			},
+			&cli.StringFlag{
+				Name:        "sas-token",
+				Usage:       "azure blob storage SAS token",
+				EnvVars:     []string{"AZ_SAS_TOKEN"},
+				Aliases:     []string{"st"},
+				Destination: &azSASToken,
+			},
 			&cli.StringFlag{
 				Name:        "url",
 				Usage:       "azure service URL",
@@ -63,7 +70,7 @@ func runAzure(ctx *cli.Context) error {
 		azServiceURL = fmt.Sprintf("https://%s.blob.core.windows.net/", azAccount)
 	}
 
-	be, err := azure.New(azAccount, azKey, azServiceURL)
+	be, err := azure.New(azAccount, azKey, azServiceURL, azSASToken)
 	if err != nil {
 		return fmt.Errorf("init azure: %v", err)
 	}