Add Azure Workload Identity authentication support

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## main / unreleased
 
+* [FEATURE] Add support for Azure Workload Identity authentication [#2195](https://github.com/grafana/tempo/pull/2195) (@LambArchie)
 * [FEATURE] Add flag to check configuration [#2131](https://github.com/grafana/tempo/issues/2131) (@robertscherbarth @agrib-01)
 * [FEATURE] Add flag to optionally enable all available Go runtime metrics [#2005](https://github.com/grafana/tempo/pull/2005) (@andreasgerstmayr)
 * [CHANGE] Add support for s3 session token in static config [#2093](https://github.com/grafana/tempo/pull/2093) (@farodin91)

docs/sources/tempo/configuration/_index.md

@@ -698,6 +698,12 @@ storage:
             # Azure German(blob.core.cloudapi.de), Azure US Government(blob.core.usgovcloudapi.net).
             [endpoint_suffix: <string>]
 
+            # optional.
+            # Azure Cloud environment. Supported values are: AzureGlobal,
+            # AzureChinaCloud, AzureGermanCloud, AzureUSGovernment.
+            # Used by Azure Workload Identity.
+            [environment: <string> | default = "AzureGlobal"]
+
             # Name of the azure storage account
             [storage_account_name: <string>]
 
@@ -709,6 +715,13 @@ storage:
             # use Azure Managed Identity to access Azure storage.
             [use_managed_identity: <bool>]
 
+            # optional.
+            # Use a Federated Token to authenticate to the Azure storage account.
+            # Enable if you want to use Azure Workload Identity. Expects AZURE_CLIENT_ID,
+            # AZURE_TENANT_ID and AZURE_FEDERATED_TOKEN_FILE envs to be present (set automatically
+            # when using Azure Workload Identity).
+            [use_federated_token: <boolean> | default = false]
+
             # optional.
             # The Client ID for the user-assigned Azure Managed Identity used to access Azure storage.
             [user_assigned_id: <bool>]

docs/sources/tempo/configuration/azure.md

@@ -13,6 +13,7 @@ Tempo requires the following configuration to authenticate to and access Azure b
   - An Azure Managed Identity; either system or user assigned. To use Azure Managed Identities, you'll need to set `use_managed_identity` to `true` in the configuration file or set `user_assigned_id` to the client ID for the managed identity you'd like to use.  
       - For a system-assigned managed identity, no additional configuration is required.
       - For a user-assigned managed identity, you'll need to set `user_assigned_id` to the client ID for the managed identity in the configuration file.
+  - Via Azure Workload Identity. To use Azure Workload Identity, you'll need to enable Workload Identity on your AKS cluster, add the required label and annotation to the service account and the reqiured pod label. Additionally, ensure environment is set in the configuration.
 
 ## Azure blocklist polling
 

modules/storage/config.go

@@ -69,6 +69,7 @@ func (cfg *Config) RegisterFlagsAndApplyDefaults(prefix string, f *flag.FlagSet)
 	f.Var(&cfg.Trace.Azure.StorageAccountKey, util.PrefixConfig(prefix, "trace.azure.storage_account_key"), "Azure storage access key.")
 	f.StringVar(&cfg.Trace.Azure.ContainerName, util.PrefixConfig(prefix, "trace.azure.container_name"), "", "Azure container name to store blocks in.")
 	f.StringVar(&cfg.Trace.Azure.Endpoint, util.PrefixConfig(prefix, "trace.azure.endpoint"), "blob.core.windows.net", "Azure endpoint to push blocks to.")
+	f.StringVar(&cfg.Trace.Azure.Environment, util.PrefixConfig(prefix, "trace.azure.environment"), "AzureGlobal", "Azure Cloud environment")
 	f.IntVar(&cfg.Trace.Azure.MaxBuffers, util.PrefixConfig(prefix, "trace.azure.max_buffers"), 4, "Number of simultaneous uploads.")
 	cfg.Trace.Azure.BufferSize = 3 * 1024 * 1024
 	cfg.Trace.Azure.HedgeRequestsUpTo = 2

tempodb/backend/azure/azure_helpers.go

@@ -14,14 +14,34 @@ import (
 	"github.com/Azure/azure-pipeline-go/pipeline"
 	blob "github.com/Azure/azure-storage-blob-go/azblob"
 	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/Azure/go-autorest/autorest/azure/auth"
 	"github.com/cristalhq/hedgedhttp"
 )
 
 const (
 	maxRetries = 1
+
+	// Environment
+	azureGlobal       = "AzureGlobal"
+	azurePublicCloud  = "AzurePublicCloud"
+	azureChinaCloud   = "AzureChinaCloud"
+	azureGermanCloud  = "AzureGermanCloud"
+	azureUSGovernment = "AzureUSGovernment"
 )
 
+var (
+	defaultAuthFunctions = authFunctions{
+		NewOAuthConfigFunc: adal.NewOAuthConfig,
+		NewServicePrincipalTokenFromFederatedTokenFunc: adal.NewServicePrincipalTokenFromFederatedToken,
+	}
+)
+
+type authFunctions struct {
+	NewOAuthConfigFunc                             func(activeDirectoryEndpoint, tenantID string) (*adal.OAuthConfig, error)
+	NewServicePrincipalTokenFromFederatedTokenFunc func(oauthConfig adal.OAuthConfig, clientID string, jwt string, resource string, callbacks ...adal.TokenRefreshCallback) (*adal.ServicePrincipalToken, error)
+}
+
 func GetContainerURL(ctx context.Context, cfg *Config, hedge bool) (blob.ContainerURL, error) {
 	var err error
 	var p pipeline.Pipeline
@@ -71,7 +91,7 @@ func GetContainerURL(ctx context.Context, cfg *Config, hedge bool) (blob.Contain
 		HTTPSender: httpSender,
 	}
 
-	if !cfg.UseManagedIdentity && cfg.UserAssignedID == "" {
+	if !cfg.UseFederatedToken && !cfg.UseManagedIdentity && cfg.UserAssignedID == "" {
 		credential, err := blob.NewSharedKeyCredential(getStorageAccountName(cfg), getStorageAccountKey(cfg))
 		if err != nil {
 			return blob.ContainerURL{}, err
@@ -153,7 +173,7 @@ func getStorageAccountKey(cfg *Config) string {
 }
 
 func getOAuthToken(cfg *Config) (*blob.TokenCredential, error) {
-	spt, err := getServicePrincipalToken(cfg)
+	spt, err := getServicePrincipalToken(cfg, defaultAuthFunctions)
 	if err != nil {
 		return nil, err
 	}
@@ -181,11 +201,37 @@ func getOAuthToken(cfg *Config) (*blob.TokenCredential, error) {
 	return &tc, nil
 }
 
-func getServicePrincipalToken(cfg *Config) (*adal.ServicePrincipalToken, error) {
+func getServicePrincipalToken(cfg *Config, authFunctions authFunctions) (*adal.ServicePrincipalToken, error) {
 	endpoint := cfg.Endpoint
 
 	resource := fmt.Sprintf("https://%s.%s", cfg.StorageAccountName, endpoint)
 
+	if cfg.UseFederatedToken {
+		token, err := servicePrincipalTokenFromFederatedToken(cfg.Environment, resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc)
+		if err != nil {
+			return nil, err
+		}
+
+		var customRefreshFunc adal.TokenRefresh = func(context context.Context, resource string) (*adal.Token, error) {
+			newToken, err := servicePrincipalTokenFromFederatedToken(cfg.Environment, resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc)
+			if err != nil {
+				return nil, err
+			}
+
+			err = newToken.Refresh()
+			if err != nil {
+				return nil, err
+			}
+
+			token := newToken.Token()
+
+			return &token, nil
+		}
+
+		token.SetCustomRefreshFunc(customRefreshFunc)
+		return token, err
+	}
+
 	msiConfig := auth.MSIConfig{
 		Resource: resource,
 	}
@@ -196,3 +242,32 @@ func getServicePrincipalToken(cfg *Config) (*adal.ServicePrincipalToken, error)
 
 	return msiConfig.ServicePrincipalToken()
 }
+
+func servicePrincipalTokenFromFederatedToken(environment string, resource string, newOAuthConfigFunc func(activeDirectoryEndpoint, tenantID string) (*adal.OAuthConfig, error), newServicePrincipalTokenFromFederatedTokenFunc func(oauthConfig adal.OAuthConfig, clientID string, jwt string, resource string, callbacks ...adal.TokenRefreshCallback) (*adal.ServicePrincipalToken, error)) (*adal.ServicePrincipalToken, error) {
+	environmentName := azurePublicCloud
+	if environment != azureGlobal {
+		environmentName = environment
+	}
+
+	env, err := azure.EnvironmentFromName(environmentName)
+	if err != nil {
+		return nil, err
+	}
+
+	azClientID := os.Getenv("AZURE_CLIENT_ID")
+	azTenantID := os.Getenv("AZURE_TENANT_ID")
+
+	jwtBytes, err := os.ReadFile(os.Getenv("AZURE_FEDERATED_TOKEN_FILE"))
+	if err != nil {
+		return nil, err
+	}
+
+	jwt := string(jwtBytes)
+
+	oauthConfig, err := newOAuthConfigFunc(env.ActiveDirectoryEndpoint, azTenantID)
+	if err != nil {
+		return nil, err
+	}
+
+	return newServicePrincipalTokenFromFederatedTokenFunc(*oauthConfig, azClientID, jwt, resource)
+}

tempodb/backend/azure/azure_helpers_test.go

@@ -5,13 +5,17 @@ import (
 	"os"
 	"testing"
 
+	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/grafana/dskit/flagext"
 	"github.com/stretchr/testify/assert"
 )
 
 const (
 	TestStorageAccountName = "foobar"
 	TestStorageAccountKey  = "abc123"
+	TestAzureClientID      = "myClientId"
+	TestAzureTenantID      = "myTenantId"
 )
 
 // TestGetStorageAccountName* explicitly broken out into
@@ -109,3 +113,41 @@ func TestGetContainerURL(t *testing.T) {
 		})
 	}
 }
+
+func TestServicePrincipalTokenFromFederatedToken(t *testing.T) {
+	os.Setenv("AZURE_CLIENT_ID", TestAzureClientID)
+	defer os.Unsetenv("AZURE_CLIENT_ID")
+	os.Setenv("AZURE_TENANT_ID", TestAzureTenantID)
+	defer os.Unsetenv("AZURE_TENANT_ID")
+
+	mockOAuthConfig, _ := adal.NewOAuthConfig("foo", "bar")
+	mockedServicePrincipalToken := new(adal.ServicePrincipalToken)
+
+	tmpDir := t.TempDir()
+	_ = os.WriteFile(tmpDir+"/jwtToken", []byte("myJwtToken"), 0666)
+	os.Setenv("AZURE_FEDERATED_TOKEN_FILE", tmpDir+"/jwtToken")
+	defer os.Unsetenv("AZURE_FEDERATED_TOKEN_FILE")
+
+	newOAuthConfigFunc := func(activeDirectoryEndpoint, tenantID string) (*adal.OAuthConfig, error) {
+		assert.Equal(t, azure.PublicCloud.ActiveDirectoryEndpoint, activeDirectoryEndpoint)
+		assert.Equal(t, TestAzureTenantID, tenantID)
+
+		_, err := adal.NewOAuthConfig(activeDirectoryEndpoint, tenantID)
+		assert.NoError(t, err)
+
+		return mockOAuthConfig, nil
+	}
+
+	servicePrincipalTokenFromFederatedTokenFunc := func(oauthConfig adal.OAuthConfig, clientID string, jwt string, resource string, callbacks ...adal.TokenRefreshCallback) (*adal.ServicePrincipalToken, error) {
+		assert.True(t, *mockOAuthConfig == oauthConfig, "should return the mocked object")
+		assert.Equal(t, TestAzureClientID, clientID)
+		assert.Equal(t, "myJwtToken", jwt)
+		assert.Equal(t, "https://bar.blob.core.windows.net", resource)
+		return mockedServicePrincipalToken, nil
+	}
+
+	token, err := servicePrincipalTokenFromFederatedToken("AzureGlobal", "https://bar.blob.core.windows.net", newOAuthConfigFunc, servicePrincipalTokenFromFederatedTokenFunc)
+
+	assert.NoError(t, err)
+	assert.True(t, mockedServicePrincipalToken == token, "should return the mocked object")
+}

tempodb/backend/azure/config.go

@@ -10,9 +10,11 @@ type Config struct {
 	StorageAccountName string         `yaml:"storage_account_name"`
 	StorageAccountKey  flagext.Secret `yaml:"storage_account_key"`
 	UseManagedIdentity bool           `yaml:"use_managed_identity"`
+	UseFederatedToken  bool           `yaml:"use_federated_token"`
 	UserAssignedID     string         `yaml:"user_assigned_id"`
 	ContainerName      string         `yaml:"container_name"`
 	Endpoint           string         `yaml:"endpoint_suffix"`
+	Environment        string         `yaml:"environment"`
 	MaxBuffers         int            `yaml:"max_buffers"`
 	BufferSize         int            `yaml:"buffer_size"`
 	HedgeRequestsAt    time.Duration  `yaml:"hedge_requests_at"`