Added UseAudienceFromAdmin property to force pull audience from admin…

… config. Default is false and expects clients to pass it

Signed-off-by: pmahindrakar-oss <[email protected]>

clients/go/admin/config.go

@@ -53,6 +53,7 @@ type Config struct {
 	ClientSecretLocation string   `json:"clientSecretLocation" pflag:",File containing the client secret"`
 	ClientSecretEnvVar   string   `json:"clientSecretEnvVar" pflag:",Environment variable containing the client secret"`
 	Scopes               []string `json:"scopes" pflag:",List of scopes to request"`
+	UseAudienceFromAdmin bool     `json:"useAudienceFromAdmin" pflag:",Use Audience configured from admins public endpoint config."`
 	Audience             string   `json:"audience" pflag:",Audience to use when initiating OAuth2 authorization requests."`
 
 	// There are two ways to get the token URL. If the authorization server url is provided, the client will try to use RFC 8414 to

clients/go/admin/token_source_provider.go

@@ -53,7 +53,7 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T
 		scopes := cfg.Scopes
 		audienceValue := cfg.Audience
 
-		if len(scopes) == 0 || len(audienceValue) == 0 {
+		if len(scopes) == 0 || cfg.UseAudienceFromAdmin {
 			publicClientConfig, err := authClient.GetPublicClientConfig(ctx, &service.PublicClientAuthConfigRequest{})
 			if err != nil {
 				return nil, fmt.Errorf("failed to fetch client metadata. Error: %v", err)
@@ -63,9 +63,7 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T
 				scopes = publicClientConfig.Scopes
 			}
 			// Update audience from publicClientConfig
-			if len(audienceValue) == 0 {
-				audienceValue = publicClientConfig.Audience
-			}
+			audienceValue = publicClientConfig.Audience
 		}
 
 		tokenProvider, err = NewClientCredentialsTokenSourceProvider(ctx, cfg, scopes, tokenURL, audienceValue)

clients/go/admin/token_source_test.go

@@ -38,44 +38,45 @@ func TestNewTokenSourceProvider(t *testing.T) {
 		name                 string
 		audienceCfg          string
 		scopesCfg            []string
+		useAudienceFromAdmin bool
 		clientConfigResponse service.PublicClientAuthConfigResponse
 		expectedAudience     string
 		expectedScopes       []string
 	}{
 		{
 			name:                 "audience from client config",
-			audienceCfg:          "aud",
+			audienceCfg:          "clientConfiguredAud",
 			scopesCfg:            []string{"all"},
 			clientConfigResponse: service.PublicClientAuthConfigResponse{},
-			expectedAudience:     "aud",
+			expectedAudience:     "clientConfiguredAud",
 			expectedScopes:       []string{"all"},
 		},
 		{
 			name:                 "audience from public client response",
-			audienceCfg:          "",
-			scopesCfg:            []string{},
-			clientConfigResponse: service.PublicClientAuthConfigResponse{Audience: "aud", Scopes: []string{"all"}},
-			expectedAudience:     "aud",
+			audienceCfg:          "clientConfiguredAud",
+			useAudienceFromAdmin: true,
+			scopesCfg:            []string{"all"},
+			clientConfigResponse: service.PublicClientAuthConfigResponse{Audience: "AdminConfiguredAud", Scopes: []string{}},
+			expectedAudience:     "AdminConfiguredAud",
 			expectedScopes:       []string{"all"},
 		},
 	}
 	for _, test := range tests {
-		t.Run("audience from client config", func(t *testing.T) {
-			cfg := GetConfig(ctx)
-			tokenCache := &tokenCacheMocks.TokenCache{}
-			metadataClient := &adminMocks.AuthMetadataServiceClient{}
-			metadataClient.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil)
-			metadataClient.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&test.clientConfigResponse, nil)
-			cfg.AuthType = AuthTypeClientSecret
-			cfg.Audience = test.audienceCfg
-			cfg.Scopes = test.scopesCfg
-			flyteTokenSource, err := NewTokenSourceProvider(ctx, cfg, tokenCache, metadataClient)
-			assert.NoError(t, err)
-			assert.NotNil(t, flyteTokenSource)
-			clientCredSourceProvider, ok := flyteTokenSource.(ClientCredentialsTokenSourceProvider)
-			assert.True(t, ok)
-			assert.Equal(t, test.expectedScopes, clientCredSourceProvider.ccConfig.Scopes)
-			assert.Equal(t, url.Values{audienceKey: {test.expectedAudience}}, clientCredSourceProvider.ccConfig.EndpointParams)
-		})
+		cfg := GetConfig(ctx)
+		tokenCache := &tokenCacheMocks.TokenCache{}
+		metadataClient := &adminMocks.AuthMetadataServiceClient{}
+		metadataClient.OnGetOAuth2MetadataMatch(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil)
+		metadataClient.OnGetPublicClientConfigMatch(mock.Anything, mock.Anything).Return(&test.clientConfigResponse, nil)
+		cfg.AuthType = AuthTypeClientSecret
+		cfg.Audience = test.audienceCfg
+		cfg.Scopes = test.scopesCfg
+		cfg.UseAudienceFromAdmin = test.useAudienceFromAdmin
+		flyteTokenSource, err := NewTokenSourceProvider(ctx, cfg, tokenCache, metadataClient)
+		assert.NoError(t, err)
+		assert.NotNil(t, flyteTokenSource)
+		clientCredSourceProvider, ok := flyteTokenSource.(ClientCredentialsTokenSourceProvider)
+		assert.True(t, ok)
+		assert.Equal(t, test.expectedScopes, clientCredSourceProvider.ccConfig.Scopes)
+		assert.Equal(t, url.Values{audienceKey: {test.expectedAudience}}, clientCredSourceProvider.ccConfig.EndpointParams)
 	}
 }