feat(): add ManagedIdentity in Azure Blob Storage

docs/sources/configuration/_index.md

@@ -706,6 +706,10 @@ The `azure_storage_config` configures Azure as a general storage for different d
 # Maximum time to wait before retrying a request.
 # CLI flag: -<prefix>.azure.max-retry-delay
 [max_retry_delay: <duration> | default = 500ms]
+
+# Use Managed Identity or not.
+# CLI flag: -ruler.storage.azure.use-managed-identity
+[use_managed_identity: <boolean> | default = false]
 ```
 
 ## gcs_storage_config

docs/sources/storage/_index.md

@@ -307,6 +307,7 @@ storage_config:
     # See https://docs.microsoft.com/en-us/azure/storage/blobs/storage-blobs-introduction#containers
     container_name: <container-name>
     request_timeout: 0
+    use_managed_identity: <true|false>
   boltdb_shipper:
     active_index_directory: /data/loki/boltdb-shipper-active
     cache_location: /data/loki/boltdb-shipper-cache

go.mod

@@ -6,6 +6,7 @@ require (
 	cloud.google.com/go/bigtable v1.3.0
 	cloud.google.com/go/pubsub v1.3.1
 	cloud.google.com/go/storage v1.10.0
+	github.com/Azure/go-autorest/autorest/adal v0.9.14
 	github.com/Azure/azure-pipeline-go v0.2.3
 	github.com/Azure/azure-storage-blob-go v0.13.0
 	github.com/Masterminds/sprig/v3 v3.2.2

pkg/storage/chunk/azure/blob_storage_client.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/Azure/azure-pipeline-go/pipeline"
 	"github.com/Azure/azure-storage-blob-go/azblob"
+	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/mattn/go-ieproxy"
 
 	cortex_azure "github.com/cortexproject/cortex/pkg/chunk/azure"
@@ -90,6 +91,7 @@ type BlobStorageConfig struct {
 	MaxRetries         int            `yaml:"max_retries"`
 	MinRetryDelay      time.Duration  `yaml:"min_retry_delay"`
 	MaxRetryDelay      time.Duration  `yaml:"max_retry_delay"`
+	UseManagedIdentity bool           `yaml:"use_managed_identity"`
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet
@@ -110,6 +112,7 @@ func (c *BlobStorageConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagS
 	f.IntVar(&c.MaxRetries, prefix+"azure.max-retries", 5, "Number of retries for a request which times out.")
 	f.DurationVar(&c.MinRetryDelay, prefix+"azure.min-retry-delay", 10*time.Millisecond, "Minimum time to wait before retrying a request.")
 	f.DurationVar(&c.MaxRetryDelay, prefix+"azure.max-retry-delay", 500*time.Millisecond, "Maximum time to wait before retrying a request.")
+	f.BoolVar(&c.UseManagedIdentity, prefix+"azure.use-managed-identity", false, "Use Managed Identity or not.")
 }
 
 func (c *BlobStorageConfig) ToCortexAzureConfig() cortex_azure.BlobStorageConfig {
@@ -234,39 +237,97 @@ func (b *BlobStorage) buildContainerURL() (azblob.ContainerURL, error) {
 }
 
 func (b *BlobStorage) newPipeline(hedging bool) (pipeline.Pipeline, error) {
-	credential, err := azblob.NewSharedKeyCredential(b.cfg.AccountName, b.cfg.AccountKey.Value)
+	if b.cfg.UseManagedIdentity == false {
+		credential, err := azblob.NewSharedKeyCredential(b.cfg.AccountName, b.cfg.AccountKey.Value)
+		if err != nil {
+			return nil, err
+		}
+
+		opts := azblob.PipelineOptions{
+			Retry: azblob.RetryOptions{
+				Policy:        azblob.RetryPolicyExponential,
+				MaxTries:      (int32)(b.cfg.MaxRetries),
+				TryTimeout:    b.cfg.RequestTimeout,
+				RetryDelay:    b.cfg.MinRetryDelay,
+				MaxRetryDelay: b.cfg.MaxRetryDelay,
+			},
+		}
+
+		if hedging {
+			opts.HTTPSender = pipeline.FactoryFunc(func(next pipeline.Policy, po *pipeline.PolicyOptions) pipeline.PolicyFunc {
+				client := b.hedgingCfg.Client(defaultClient)
+				return func(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
+					resp, err := client.Do(request.WithContext(ctx))
+					return pipeline.NewHTTPResponse(resp), err
+				}
+			})
+		}
+
+		return azblob.NewPipeline(credential, opts), nil
+	} else {
+		tokenCredential, err := b.getOAuthToken()
+		if err != nil {
+			return nil, err
+		}
+
+		return azblob.NewPipeline(*tokenCredential, azblob.PipelineOptions{
+			Retry: azblob.RetryOptions{
+				Policy:        azblob.RetryPolicyExponential,
+				MaxTries:      (int32)(b.cfg.MaxRetries),
+				TryTimeout:    b.cfg.RequestTimeout,
+				RetryDelay:    b.cfg.MinRetryDelay,
+				MaxRetryDelay: b.cfg.MaxRetryDelay,
+			},
+		}), nil
+	}
+
+}
+
+func (b *BlobStorage) getOAuthToken() (*azblob.TokenCredential, error) {
+	spt, err := b.fetchMSIToken()
 	if err != nil {
 		return nil, err
 	}
 
-	opts := azblob.PipelineOptions{
-		Retry: azblob.RetryOptions{
-			Policy:        azblob.RetryPolicyExponential,
-			MaxTries:      (int32)(b.cfg.MaxRetries),
-			TryTimeout:    b.cfg.RequestTimeout,
-			RetryDelay:    b.cfg.MinRetryDelay,
-			MaxRetryDelay: b.cfg.MaxRetryDelay,
-		},
+	// Refresh obtains a fresh token
+	err = spt.Refresh()
+	if err != nil {
+		return nil, err
 	}
 
-	opts.HTTPSender = pipeline.FactoryFunc(func(next pipeline.Policy, po *pipeline.PolicyOptions) pipeline.PolicyFunc {
-		return func(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
-			resp, err := defaultClient.Do(request.WithContext(ctx))
-			return pipeline.NewHTTPResponse(resp), err
+	tc := azblob.NewTokenCredential(spt.Token().AccessToken, func(tc azblob.TokenCredential) time.Duration {
+		err := spt.Refresh()
+		if err != nil {
+			// something went wrong, prevent the refresher from being triggered again
+			return 0
 		}
+
+		// set the new token value
+		tc.SetToken(spt.Token().AccessToken)
+
+		// get the next token slightly before the current one expires
+		return time.Until(spt.Token().Expires()) - 10*time.Second
 	})
 
-	if hedging {
-		opts.HTTPSender = pipeline.FactoryFunc(func(next pipeline.Policy, po *pipeline.PolicyOptions) pipeline.PolicyFunc {
-			client := b.hedgingCfg.Client(defaultClient)
-			return func(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
-				resp, err := client.Do(request.WithContext(ctx))
-				return pipeline.NewHTTPResponse(resp), err
-			}
-		})
+	return &tc, nil
+}
+
+func (b *BlobStorage) fetchMSIToken() (*adal.ServicePrincipalToken, error) {
+	// msiEndpoint is the well known endpoint for getting MSI authentications tokens
+	// msiEndpoint := "http://169.254.169.254/metadata/identity/oauth2/token" for production Jobs
+	msiEndpoint, _ := adal.GetMSIVMEndpoint()
+
+	var spt *adal.ServicePrincipalToken
+	var err error
+
+	// both can be empty, systemAssignedMSI scenario
+	spt, err = adal.NewServicePrincipalTokenFromMSI(msiEndpoint, "https://storage.azure.com/")
+
+	if err != nil {
+		return nil, err
 	}
 
-	return azblob.NewPipeline(credential, opts), nil
+	return spt, spt.Refresh()
 }
 
 // List implements chunk.ObjectClient.