From 183122ab1e8471bec4206e8ac9ed852410e116c3 Mon Sep 17 00:00:00 2001 From: Archie Lamb Date: Fri, 10 Mar 2023 17:45:04 +0000 Subject: [PATCH 1/4] Add Azure Workload Identity authentication support --- CHANGELOG.md | 1 + docs/sources/tempo/configuration/_index.md | 13 ++++ docs/sources/tempo/configuration/azure.md | 1 + modules/storage/config.go | 1 + tempodb/backend/azure/azure_helpers.go | 81 ++++++++++++++++++++- tempodb/backend/azure/azure_helpers_test.go | 42 +++++++++++ tempodb/backend/azure/config.go | 2 + 7 files changed, 138 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3c99deaa88..3ace9a8c2d7 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/docs/sources/tempo/configuration/_index.md b/docs/sources/tempo/configuration/_index.md index 2645d0e152c..fcf85ddbedb 100644 --- a/docs/sources/tempo/configuration/_index.md +++ b/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: ] + # optional. + # Azure Cloud environment. Supported values are: AzureGlobal, + # AzureChinaCloud, AzureGermanCloud, AzureUSGovernment. + # Used by Azure Workload Identity. + [environment: | default = "AzureGlobal"] + # Name of the azure storage account [storage_account_name: ] @@ -709,6 +715,13 @@ storage: # use Azure Managed Identity to access Azure storage. [use_managed_identity: ] + # 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: | default = false] + # optional. # The Client ID for the user-assigned Azure Managed Identity used to access Azure storage. [user_assigned_id: ] diff --git a/docs/sources/tempo/configuration/azure.md b/docs/sources/tempo/configuration/azure.md index 4a2c16143c3..5ce8a9f5d37 100644 --- a/docs/sources/tempo/configuration/azure.md +++ b/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 diff --git a/modules/storage/config.go b/modules/storage/config.go index 8ae2ebec976..fd40f2c5e9e 100644 --- a/modules/storage/config.go +++ b/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 diff --git a/tempodb/backend/azure/azure_helpers.go b/tempodb/backend/azure/azure_helpers.go index 7e4f9e8db8c..ba3c89e194e 100644 --- a/tempodb/backend/azure/azure_helpers.go +++ b/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) +} diff --git a/tempodb/backend/azure/azure_helpers_test.go b/tempodb/backend/azure/azure_helpers_test.go index 6b3ab65b523..f1cafea2914 100644 --- a/tempodb/backend/azure/azure_helpers_test.go +++ b/tempodb/backend/azure/azure_helpers_test.go @@ -5,6 +5,8 @@ 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" ) @@ -12,6 +14,8 @@ import ( 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") +} diff --git a/tempodb/backend/azure/config.go b/tempodb/backend/azure/config.go index f02a90a88e0..080d818cdc5 100644 --- a/tempodb/backend/azure/config.go +++ b/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"` From 90510c72d1c1fcd9b288e6328bc06f440c85b437 Mon Sep 17 00:00:00 2001 From: Archie Lamb Date: Mon, 13 Mar 2023 10:39:59 +0000 Subject: [PATCH 2/4] Linting --- tempodb/backend/azure/azure_helpers_test.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tempodb/backend/azure/azure_helpers_test.go b/tempodb/backend/azure/azure_helpers_test.go index f1cafea2914..1147b5e7829 100644 --- a/tempodb/backend/azure/azure_helpers_test.go +++ b/tempodb/backend/azure/azure_helpers_test.go @@ -14,8 +14,8 @@ import ( const ( TestStorageAccountName = "foobar" TestStorageAccountKey = "abc123" - TestAzureClientId = "myClientId" - TestAzureTenantId = "myTenantId" + TestAzureClientID = "myClientId" + TestAzureTenantID = "myTenantId" ) // TestGetStorageAccountName* explicitly broken out into @@ -115,9 +115,9 @@ func TestGetContainerURL(t *testing.T) { } func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { - os.Setenv("AZURE_CLIENT_ID", TestAzureClientId) + os.Setenv("AZURE_CLIENT_ID", TestAzureClientID) defer os.Unsetenv("AZURE_CLIENT_ID") - os.Setenv("AZURE_TENANT_ID", TestAzureTenantId) + os.Setenv("AZURE_TENANT_ID", TestAzureTenantID) defer os.Unsetenv("AZURE_TENANT_ID") mockOAuthConfig, _ := adal.NewOAuthConfig("foo", "bar") @@ -130,7 +130,7 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { newOAuthConfigFunc := func(activeDirectoryEndpoint, tenantID string) (*adal.OAuthConfig, error) { assert.Equal(t, azure.PublicCloud.ActiveDirectoryEndpoint, activeDirectoryEndpoint) - assert.Equal(t, TestAzureTenantId, tenantID) + assert.Equal(t, TestAzureTenantID, tenantID) _, err := adal.NewOAuthConfig(activeDirectoryEndpoint, tenantID) assert.NoError(t, err) @@ -140,7 +140,7 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { 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, TestAzureClientID, clientID) assert.Equal(t, "myJwtToken", jwt) assert.Equal(t, "https://bar.blob.core.windows.net", resource) return mockedServicePrincipalToken, nil From d5f3332c85914414aaf26c9abfd3a16a3198da10 Mon Sep 17 00:00:00 2001 From: Archie Lamb Date: Tue, 14 Mar 2023 14:05:12 +0000 Subject: [PATCH 3/4] Automate getting Azure AD Endpoint --- docs/sources/tempo/configuration/_index.md | 12 +++------ docs/sources/tempo/configuration/azure.md | 2 +- modules/storage/config.go | 1 - tempodb/backend/azure/azure_helpers.go | 30 +++++++-------------- tempodb/backend/azure/azure_helpers_test.go | 10 ++++--- tempodb/backend/azure/config.go | 1 - 6 files changed, 19 insertions(+), 37 deletions(-) diff --git a/docs/sources/tempo/configuration/_index.md b/docs/sources/tempo/configuration/_index.md index fcf85ddbedb..329e156a315 100644 --- a/docs/sources/tempo/configuration/_index.md +++ b/docs/sources/tempo/configuration/_index.md @@ -698,12 +698,6 @@ storage: # Azure German(blob.core.cloudapi.de), Azure US Government(blob.core.usgovcloudapi.net). [endpoint_suffix: ] - # optional. - # Azure Cloud environment. Supported values are: AzureGlobal, - # AzureChinaCloud, AzureGermanCloud, AzureUSGovernment. - # Used by Azure Workload Identity. - [environment: | default = "AzureGlobal"] - # Name of the azure storage account [storage_account_name: ] @@ -718,9 +712,9 @@ storage: # 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: | default = false] + # AZURE_TENANT_ID, AZURE_AUTHORITY_HOST and AZURE_FEDERATED_TOKEN_FILE envs to be present + # (these are set automatically when using Azure Workload Identity). + [use_federated_token: ] # optional. # The Client ID for the user-assigned Azure Managed Identity used to access Azure storage. diff --git a/docs/sources/tempo/configuration/azure.md b/docs/sources/tempo/configuration/azure.md index 5ce8a9f5d37..bab86a605d2 100644 --- a/docs/sources/tempo/configuration/azure.md +++ b/docs/sources/tempo/configuration/azure.md @@ -13,7 +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. + - Via Azure Workload Identity. To use Azure Workload Identity, you'll need to enable Azure Workload Identity on your cluster, add the required label and annotation to the service account and the required pod label. ## Azure blocklist polling diff --git a/modules/storage/config.go b/modules/storage/config.go index fd40f2c5e9e..8ae2ebec976 100644 --- a/modules/storage/config.go +++ b/modules/storage/config.go @@ -69,7 +69,6 @@ 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 diff --git a/tempodb/backend/azure/azure_helpers.go b/tempodb/backend/azure/azure_helpers.go index ba3c89e194e..d34eb03cb50 100644 --- a/tempodb/backend/azure/azure_helpers.go +++ b/tempodb/backend/azure/azure_helpers.go @@ -21,13 +21,6 @@ import ( const ( maxRetries = 1 - - // Environment - azureGlobal = "AzureGlobal" - azurePublicCloud = "AzurePublicCloud" - azureChinaCloud = "AzureChinaCloud" - azureGermanCloud = "AzureGermanCloud" - azureUSGovernment = "AzureUSGovernment" ) var ( @@ -207,13 +200,13 @@ func getServicePrincipalToken(cfg *Config, authFunctions authFunctions) (*adal.S resource := fmt.Sprintf("https://%s.%s", cfg.StorageAccountName, endpoint) if cfg.UseFederatedToken { - token, err := servicePrincipalTokenFromFederatedToken(cfg.Environment, resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc) + token, err := servicePrincipalTokenFromFederatedToken(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) + newToken, err := servicePrincipalTokenFromFederatedToken(resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc) if err != nil { return nil, err } @@ -243,20 +236,15 @@ func getServicePrincipalToken(cfg *Config, authFunctions authFunctions) (*adal.S 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 - } - +func servicePrincipalTokenFromFederatedToken(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) { azClientID := os.Getenv("AZURE_CLIENT_ID") azTenantID := os.Getenv("AZURE_TENANT_ID") + azADEndpoint, ok := os.LookupEnv("AZURE_AUTHORITY_HOST") + if !ok { + azADEndpoint = azure.PublicCloud.ActiveDirectoryEndpoint + } + jwtBytes, err := os.ReadFile(os.Getenv("AZURE_FEDERATED_TOKEN_FILE")) if err != nil { return nil, err @@ -264,7 +252,7 @@ func servicePrincipalTokenFromFederatedToken(environment string, resource string jwt := string(jwtBytes) - oauthConfig, err := newOAuthConfigFunc(env.ActiveDirectoryEndpoint, azTenantID) + oauthConfig, err := newOAuthConfigFunc(azADEndpoint, azTenantID) if err != nil { return nil, err } diff --git a/tempodb/backend/azure/azure_helpers_test.go b/tempodb/backend/azure/azure_helpers_test.go index 1147b5e7829..5a40a9ac104 100644 --- a/tempodb/backend/azure/azure_helpers_test.go +++ b/tempodb/backend/azure/azure_helpers_test.go @@ -6,7 +6,6 @@ import ( "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" ) @@ -16,6 +15,7 @@ const ( TestStorageAccountKey = "abc123" TestAzureClientID = "myClientId" TestAzureTenantID = "myTenantId" + TestAzureADEndpoint = "https://example.com/" ) // TestGetStorageAccountName* explicitly broken out into @@ -119,8 +119,10 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { defer os.Unsetenv("AZURE_CLIENT_ID") os.Setenv("AZURE_TENANT_ID", TestAzureTenantID) defer os.Unsetenv("AZURE_TENANT_ID") + os.Setenv("AZURE_AUTHORITY_HOST", TestAzureADEndpoint) + defer os.Unsetenv("AZURE_AUTHORITY_HOST") - mockOAuthConfig, _ := adal.NewOAuthConfig("foo", "bar") + mockOAuthConfig, _ := adal.NewOAuthConfig(TestAzureADEndpoint, "bar") mockedServicePrincipalToken := new(adal.ServicePrincipalToken) tmpDir := t.TempDir() @@ -129,7 +131,7 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { 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, TestAzureADEndpoint, activeDirectoryEndpoint) assert.Equal(t, TestAzureTenantID, tenantID) _, err := adal.NewOAuthConfig(activeDirectoryEndpoint, tenantID) @@ -146,7 +148,7 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { return mockedServicePrincipalToken, nil } - token, err := servicePrincipalTokenFromFederatedToken("AzureGlobal", "https://bar.blob.core.windows.net", newOAuthConfigFunc, servicePrincipalTokenFromFederatedTokenFunc) + token, err := servicePrincipalTokenFromFederatedToken("https://bar.blob.core.windows.net", newOAuthConfigFunc, servicePrincipalTokenFromFederatedTokenFunc) assert.NoError(t, err) assert.True(t, mockedServicePrincipalToken == token, "should return the mocked object") diff --git a/tempodb/backend/azure/config.go b/tempodb/backend/azure/config.go index 080d818cdc5..54eb2cd97c9 100644 --- a/tempodb/backend/azure/config.go +++ b/tempodb/backend/azure/config.go @@ -14,7 +14,6 @@ type Config struct { 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"` From 83b5d32b2d091cd8646ef047d99c26d802e6d8ce Mon Sep 17 00:00:00 2001 From: Archie Lamb Date: Tue, 14 Mar 2023 18:53:37 +0000 Subject: [PATCH 4/4] PR Feedback --- tempodb/backend/azure/azure_helpers.go | 14 +++++++------- tempodb/backend/azure/azure_helpers_test.go | 5 ++++- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/tempodb/backend/azure/azure_helpers.go b/tempodb/backend/azure/azure_helpers.go index d34eb03cb50..88c21ac38f2 100644 --- a/tempodb/backend/azure/azure_helpers.go +++ b/tempodb/backend/azure/azure_helpers.go @@ -166,7 +166,7 @@ func getStorageAccountKey(cfg *Config) string { } func getOAuthToken(cfg *Config) (*blob.TokenCredential, error) { - spt, err := getServicePrincipalToken(cfg, defaultAuthFunctions) + spt, err := getServicePrincipalToken(cfg) if err != nil { return nil, err } @@ -194,19 +194,19 @@ func getOAuthToken(cfg *Config) (*blob.TokenCredential, error) { return &tc, nil } -func getServicePrincipalToken(cfg *Config, authFunctions authFunctions) (*adal.ServicePrincipalToken, error) { +func getServicePrincipalToken(cfg *Config) (*adal.ServicePrincipalToken, error) { endpoint := cfg.Endpoint resource := fmt.Sprintf("https://%s.%s", cfg.StorageAccountName, endpoint) if cfg.UseFederatedToken { - token, err := servicePrincipalTokenFromFederatedToken(resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc) + token, err := servicePrincipalTokenFromFederatedToken(resource, defaultAuthFunctions) if err != nil { return nil, err } var customRefreshFunc adal.TokenRefresh = func(context context.Context, resource string) (*adal.Token, error) { - newToken, err := servicePrincipalTokenFromFederatedToken(resource, authFunctions.NewOAuthConfigFunc, authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc) + newToken, err := servicePrincipalTokenFromFederatedToken(resource, defaultAuthFunctions) if err != nil { return nil, err } @@ -236,7 +236,7 @@ func getServicePrincipalToken(cfg *Config, authFunctions authFunctions) (*adal.S return msiConfig.ServicePrincipalToken() } -func servicePrincipalTokenFromFederatedToken(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) { +func servicePrincipalTokenFromFederatedToken(resource string, authFunctions authFunctions) (*adal.ServicePrincipalToken, error) { azClientID := os.Getenv("AZURE_CLIENT_ID") azTenantID := os.Getenv("AZURE_TENANT_ID") @@ -252,10 +252,10 @@ func servicePrincipalTokenFromFederatedToken(resource string, newOAuthConfigFunc jwt := string(jwtBytes) - oauthConfig, err := newOAuthConfigFunc(azADEndpoint, azTenantID) + oauthConfig, err := authFunctions.NewOAuthConfigFunc(azADEndpoint, azTenantID) if err != nil { return nil, err } - return newServicePrincipalTokenFromFederatedTokenFunc(*oauthConfig, azClientID, jwt, resource) + return authFunctions.NewServicePrincipalTokenFromFederatedTokenFunc(*oauthConfig, azClientID, jwt, resource) } diff --git a/tempodb/backend/azure/azure_helpers_test.go b/tempodb/backend/azure/azure_helpers_test.go index 5a40a9ac104..002acef7376 100644 --- a/tempodb/backend/azure/azure_helpers_test.go +++ b/tempodb/backend/azure/azure_helpers_test.go @@ -148,7 +148,10 @@ func TestServicePrincipalTokenFromFederatedToken(t *testing.T) { return mockedServicePrincipalToken, nil } - token, err := servicePrincipalTokenFromFederatedToken("https://bar.blob.core.windows.net", newOAuthConfigFunc, servicePrincipalTokenFromFederatedTokenFunc) + token, err := servicePrincipalTokenFromFederatedToken("https://bar.blob.core.windows.net", authFunctions{ + newOAuthConfigFunc, + servicePrincipalTokenFromFederatedTokenFunc, + }) assert.NoError(t, err) assert.True(t, mockedServicePrincipalToken == token, "should return the mocked object")