diff --git a/docs/deployment/configuration/auth_setup.rst b/docs/deployment/configuration/auth_setup.rst index 9b19cfd8d60..16fc77da330 100644 --- a/docs/deployment/configuration/auth_setup.rst +++ b/docs/deployment/configuration/auth_setup.rst @@ -428,6 +428,63 @@ Apply OIDC Configuration It should now be possible to go to Flyte UI and be prompted for authentication. Flytectl should automatically pickup the change and start prompting for authentication as well. If you want to use an external OAuth2 provider for App authentication, please continue reading into the next section. +Headless login with device flow and ID tokens +--------------------------------------------- + +Flyte's built-in authorization server does not implement the device authorization grant +(`RFC 8628 `_), so on a host without a browser ``flytectl`` (and any +other client built on the Go admin client) cannot obtain a token from it. Instead, the client can run the device flow +directly against the OIDC provider configured in ``userAuth.openId`` and send the resulting ID token to ``flyteadmin``. +``flyteadmin`` accepts ID tokens from that provider over gRPC with the ``IDToken`` scheme, which is the same identity +the UI login uses, so no ``flyteadmin`` changes are needed. + +The OIDC provider needs: + +* A public client for the CLI that is allowed to use the device authorization grant. +* ID tokens for that client that carry ``flyteadmin``'s own client id (``userAuth.openId.clientId``) in the ``aud`` + claim. With Dex, add the CLI client to the Flyte client's ``trustedPeers`` and request the scope + ``audience:server:client_id:``. Other providers have equivalent audience settings. + +Then configure the client: + +.. code-block:: yaml + + admin: + endpoint: dns:///flyte.example.com + authType: DeviceFlow + # The provider's device_authorization_endpoint and token_endpoint, taken from its + # /.well-known/openid-configuration document. Setting them switches the device flow from + # flyteadmin's authorization server to the provider. + deviceAuthorizationUrl: https://dex.example.com/device/code + tokenUrl: https://dex.example.com/token + # Send the OIDC id_token as "IDToken " instead of the access token as "Bearer ". + tokenType: IDToken + # Client id and scopes registered with the OIDC provider. Both are required with deviceAuthorizationUrl. + clientId: flytectl + scopes: + - openid + - email + - profile + - offline_access + - "audience:server:client_id:flyteadmin" + +For ``authType: Pkce`` set ``authorizationUrl`` (the provider's ``authorization_endpoint``) and ``tokenUrl`` +instead. ``tokenType`` applies to ``authType: ExternalCommand`` too, so a command that prints an ID token can be sent +with the ``IDToken`` scheme. + +Provider notes. The client only needs the standard device authorization grant and a token response with an +``id_token`` (and a ``refresh_token`` for silent renewal), so it works with any compliant provider. What differs is +how the provider adds ``flyteadmin``'s client id to the ID token's ``aud`` claim: + +* **Dex**: list the CLI client in the Flyte client's ``trustedPeers`` and request + ``audience:server:client_id:``. +* **Keycloak**: add an *Audience* protocol mapper to the CLI client with ``flyteadmin``'s client as the included + audience, and enable *OAuth 2.0 Device Authorization Grant* on the CLI client. +* Providers that only ever issue ID tokens with the requesting client as the audience (for example Okta, Auth0, + Microsoft Entra ID, Google) need the CLI to use ``flyteadmin``'s own client id, so the ID token is issued for that + client. This works when the provider lets that client use the device grant with just its client id, as Entra ID + does with "Allow public client flows"; otherwise the provider cannot be used this way today. + *************************** Custom Authorization Server *************************** diff --git a/flyteidl/clients/go/admin/config.go b/flyteidl/clients/go/admin/config.go index 6c8adf35532..4ff69b6d43d 100644 --- a/flyteidl/clients/go/admin/config.go +++ b/flyteidl/clients/go/admin/config.go @@ -46,7 +46,7 @@ type Config struct { PerRetryTimeout config.Duration `json:"perRetryTimeout" pflag:",gRPC per retry timeout"` MaxRetries int `json:"maxRetries" pflag:",Max number of gRPC retries"` MaxMessageSizeBytes int `json:"maxMessageSizeBytes" pflag:",The max size in bytes for incoming gRPC messages"` - AuthType AuthType `json:"authType" pflag:",Type of OAuth2 flow used for communicating with admin.ClientSecret,Pkce,ExternalCommand are valid values"` + AuthType AuthType `json:"authType" pflag:",Type of OAuth2 flow used for communicating with admin.ClientSecret,Pkce,ExternalCommand,DeviceFlow are valid values"` TokenRefreshWindow config.Duration `json:"tokenRefreshWindow" pflag:",Max duration between token refresh attempt and token expiry."` // Deprecated: settings will be discovered dynamically DeprecatedUseAuth bool `json:"useAuth" pflag:",Deprecated: Auth will be enabled/disabled based on admin's dynamically discovered information."` @@ -62,8 +62,29 @@ type Config struct { // Deprecated: This will now be discovered through admin's anonymously accessible metadata. DeprecatedAuthorizationServerURL string `json:"authorizationServerUrl" pflag:",This is the URL to your IdP's authorization server. It'll default to Endpoint"` // If not provided, it'll be discovered through admin's anonymously accessible metadata endpoint. + // Together with AuthorizationURL or DeviceAuthorizationURL it also points the Pkce and DeviceFlow auth types at + // another authorization server. TokenURL string `json:"tokenUrl" pflag:",OPTIONAL: Your IdP's token endpoint. It'll be discovered from flyte admin's OAuth Metadata endpoint if not provided."` + // AuthorizationURL optionally runs the Pkce auth type against this authorization endpoint and TokenURL instead of + // the endpoints admin advertises, for example those of admin's userAuth OIDC provider. ClientID, Scopes and + // Audience from this config are then used instead of admin's public client config, since the client has to be + // registered with that server. The value is the provider's authorization_endpoint (OpenID Connect Discovery 1.0, + // RFC 8414). + AuthorizationURL string `json:"authorizationUrl" pflag:",OPTIONAL: Authorization endpoint to use for the Pkce auth type instead of the one admin advertises. Requires tokenUrl, clientId and scopes."` + // DeviceAuthorizationURL optionally runs the DeviceFlow auth type against this device authorization endpoint + // (RFC 8628) and TokenURL instead of the endpoints admin advertises, for example those of admin's userAuth OIDC + // provider. ClientID, Scopes and Audience from this config are then used instead of admin's public client config, + // since the client has to be registered with that server. The value is the provider's + // device_authorization_endpoint (RFC 8628 section 4). + DeviceAuthorizationURL string `json:"deviceAuthorizationUrl" pflag:",OPTIONAL: Device authorization endpoint to use for the DeviceFlow auth type instead of the one admin advertises. Requires tokenUrl, clientId and scopes."` + // TokenType is the scheme used when sending the token to admin for the Pkce, DeviceFlow and ExternalCommand auth + // types. Bearer (the default) sends the OAuth2 access token. IDToken sends the OIDC id_token from the token response + // instead, which admin validates against its userAuth OIDC provider; this lets a client authenticate through that + // provider directly (see AuthorizationURL and DeviceAuthorizationURL) when admin's own authorization server does not + // offer the flow it needs, such as the device authorization grant. + TokenType string `json:"tokenType" pflag:",OPTIONAL: Scheme used to send the token to admin: Bearer (default) or IDToken. IDToken sends the OIDC id_token instead of the access token."` + // See the implementation of the 'grpcAuthorizationHeader' option in Flyte Admin for more information. But // basically we want to be able to use a different string to pass the token from this client to the the Admin service // because things might be running in a service mesh (like Envoy) that already uses the default 'authorization' header diff --git a/flyteidl/clients/go/admin/config_flags.go b/flyteidl/clients/go/admin/config_flags.go index c344147c8c4..a6d237c585b 100755 --- a/flyteidl/clients/go/admin/config_flags.go +++ b/flyteidl/clients/go/admin/config_flags.go @@ -58,7 +58,7 @@ func (cfg Config) GetPFlagSet(prefix string) *pflag.FlagSet { cmdFlags.String(fmt.Sprintf("%v%v", prefix, "perRetryTimeout"), defaultConfig.PerRetryTimeout.String(), "gRPC per retry timeout") cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "maxRetries"), defaultConfig.MaxRetries, "Max number of gRPC retries") cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "maxMessageSizeBytes"), defaultConfig.MaxMessageSizeBytes, "The max size in bytes for incoming gRPC messages") - cmdFlags.String(fmt.Sprintf("%v%v", prefix, "authType"), defaultConfig.AuthType.String(), "Type of OAuth2 flow used for communicating with admin.ClientSecret, Pkce, ExternalCommand are valid values") + cmdFlags.String(fmt.Sprintf("%v%v", prefix, "authType"), defaultConfig.AuthType.String(), "Type of OAuth2 flow used for communicating with admin.ClientSecret, Pkce, ExternalCommand, DeviceFlow are valid values") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "tokenRefreshWindow"), defaultConfig.TokenRefreshWindow.String(), "Max duration between token refresh attempt and token expiry.") cmdFlags.Bool(fmt.Sprintf("%v%v", prefix, "useAuth"), defaultConfig.DeprecatedUseAuth, "Deprecated: Auth will be enabled/disabled based on admin's dynamically discovered information.") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "clientId"), defaultConfig.ClientID, "Client ID") @@ -69,6 +69,9 @@ func (cfg Config) GetPFlagSet(prefix string) *pflag.FlagSet { cmdFlags.String(fmt.Sprintf("%v%v", prefix, "audience"), defaultConfig.Audience, "Audience to use when initiating OAuth2 authorization requests.") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "authorizationServerUrl"), defaultConfig.DeprecatedAuthorizationServerURL, "This is the URL to your IdP's authorization server. It'll default to Endpoint") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "tokenUrl"), defaultConfig.TokenURL, "OPTIONAL: Your IdP's token endpoint. It'll be discovered from flyte admin's OAuth Metadata endpoint if not provided.") + cmdFlags.String(fmt.Sprintf("%v%v", prefix, "authorizationUrl"), defaultConfig.AuthorizationURL, "OPTIONAL: Authorization endpoint to use for the Pkce auth type instead of the one admin advertises. Requires tokenUrl, clientId and scopes.") + cmdFlags.String(fmt.Sprintf("%v%v", prefix, "deviceAuthorizationUrl"), defaultConfig.DeviceAuthorizationURL, "OPTIONAL: Device authorization endpoint to use for the DeviceFlow auth type instead of the one admin advertises. Requires tokenUrl, clientId and scopes.") + cmdFlags.String(fmt.Sprintf("%v%v", prefix, "tokenType"), defaultConfig.TokenType, "OPTIONAL: Scheme used to send the token to admin: Bearer (default) or IDToken. IDToken sends the OIDC id_token instead of the access token.") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "authorizationHeader"), defaultConfig.AuthorizationHeader, "Custom metadata header to pass JWT") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "pkceConfig.timeout"), defaultConfig.PkceConfig.BrowserSessionTimeout.String(), "Amount of time the browser session would be active for authentication from client app.") cmdFlags.String(fmt.Sprintf("%v%v", prefix, "pkceConfig.refreshTime"), defaultConfig.PkceConfig.TokenRefreshGracePeriod.String(), "grace period from the token expiry after which it would refresh the token.") diff --git a/flyteidl/clients/go/admin/config_flags_test.go b/flyteidl/clients/go/admin/config_flags_test.go index a0732e2b10e..0da43b41aa2 100755 --- a/flyteidl/clients/go/admin/config_flags_test.go +++ b/flyteidl/clients/go/admin/config_flags_test.go @@ -365,6 +365,48 @@ func TestConfig_SetFlags(t *testing.T) { } }) }) + t.Run("Test_authorizationUrl", func(t *testing.T) { + + t.Run("Override", func(t *testing.T) { + testValue := "1" + + cmdFlags.Set("authorizationUrl", testValue) + if vString, err := cmdFlags.GetString("authorizationUrl"); err == nil { + testDecodeJson_Config(t, fmt.Sprintf("%v", vString), &actual.AuthorizationURL) + + } else { + assert.FailNow(t, err.Error()) + } + }) + }) + t.Run("Test_deviceAuthorizationUrl", func(t *testing.T) { + + t.Run("Override", func(t *testing.T) { + testValue := "1" + + cmdFlags.Set("deviceAuthorizationUrl", testValue) + if vString, err := cmdFlags.GetString("deviceAuthorizationUrl"); err == nil { + testDecodeJson_Config(t, fmt.Sprintf("%v", vString), &actual.DeviceAuthorizationURL) + + } else { + assert.FailNow(t, err.Error()) + } + }) + }) + t.Run("Test_tokenType", func(t *testing.T) { + + t.Run("Override", func(t *testing.T) { + testValue := "1" + + cmdFlags.Set("tokenType", testValue) + if vString, err := cmdFlags.GetString("tokenType"); err == nil { + testDecodeJson_Config(t, fmt.Sprintf("%v", vString), &actual.TokenType) + + } else { + assert.FailNow(t, err.Error()) + } + }) + }) t.Run("Test_authorizationHeader", func(t *testing.T) { t.Run("Override", func(t *testing.T) { diff --git a/flyteidl/clients/go/admin/deviceflow/payload.go b/flyteidl/clients/go/admin/deviceflow/payload.go index 38e8b5502fc..cfb316bcd19 100644 --- a/flyteidl/clients/go/admin/deviceflow/payload.go +++ b/flyteidl/clients/go/admin/deviceflow/payload.go @@ -42,4 +42,5 @@ type DeviceAccessTokenResponse struct { oauth2.Token Error string `json:"error"` ExpiresIn int64 `json:"expires_in"` // relative seconds from now + IDToken string `json:"id_token"` // OIDC id_token, present when the openid scope was requested } diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go index e934ec69496..36963d6eb81 100644 --- a/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go +++ b/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go @@ -133,13 +133,17 @@ func (t TokenOrchestrator) PollTokenEndpoint(ctx context.Context, tokReq DeviceA if secs := tokResp.ExpiresIn; secs > 0 { tokResp.Token.Expiry = time.Now().Add(time.Duration(secs) * time.Second) } + token, err := t.ClientConfig.PrepareTokenWithIDToken(&tokResp.Token, tokResp.IDToken) + if err != nil { + return nil, err + } // Got the auth token in the response and save it in the cache - err = t.TokenCache.SaveToken(&tokResp.Token) + err = t.TokenCache.SaveToken(token) // Saving into the cache is only considered to be a warning in this case. if err != nil { logger.Warnf(ctx, "failed to save token in the token cache. Error: %w", err) } - return &tokResp.Token, nil + return token, nil } fmt.Printf("Waiting for %v secs\n", pollInterval.Seconds()) time.Sleep(pollInterval) diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go index 9f20fb3ef59..e67fb443aa8 100644 --- a/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go +++ b/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go @@ -123,3 +123,72 @@ func TestFetchFromAuthFlow(t *testing.T) { assert.True(t, authToken.Expiry.After(time.Now().Add(time.Second*200))) }) } + +func TestFetchFromAuthFlowWithIDToken(t *testing.T) { + ctx := context.Background() + newServer := func(t *testing.T, tokenResponse string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + assert.Nil(t, err) + if strings.Contains(string(body), deviceCode) { + _, err = w.Write([]byte(tokenResponse)) + assert.Nil(t, err) + return + } + dar := DeviceAuthorizationResponse{DeviceCode: "device-code", UserCode: "USER-CODE", VerificationURI: "https://idp/device", Interval: 1} + darBytes, err := json.Marshal(dar) + assert.Nil(t, err) + _, err = w.Write(darBytes) + assert.Nil(t, err) + })) + } + newOrchestrator := func(t *testing.T, server *httptest.Server, tokenType string) TokenOrchestrator { + orchestrator, err := NewDeviceFlowTokenOrchestrator(tokenorchestrator.BaseTokenOrchestrator{ + ClientConfig: &oauth.Config{ + Config: &oauth2.Config{ + ClientID: clientID, + Scopes: []string{"openid", "offline_access"}, + Endpoint: oauth2.Endpoint{TokenURL: server.URL}, + }, + DeviceEndpoint: server.URL, + TokenType: tokenType, + }, + TokenCache: cache.NewTokenCacheInMemoryProvider(), + }, Config{Timeout: config.Duration{Duration: 1 * time.Minute}}) + assert.NoError(t, err) + return orchestrator + } + + t.Run("id token is sent as IDToken", func(t *testing.T) { + server := newServer(t, `{"access_token": "access", "refresh_token": "refresh", "token_type": "bearer", "expires_in": 300, "id_token": "id"}`) + defer server.Close() + orchestrator := newOrchestrator(t, server, oauth.TokenTypeIDToken) + authToken, err := orchestrator.FetchTokenFromAuthFlow(ctx) + assert.NoError(t, err) + assert.Equal(t, "id", authToken.AccessToken) + assert.Equal(t, "refresh", authToken.RefreshToken) + assert.Equal(t, oauth.TokenTypeIDToken, authToken.Type()) + cached, err := orchestrator.TokenCache.GetToken() + assert.NoError(t, err) + assert.Equal(t, "id", cached.AccessToken) + }) + + t.Run("bearer keeps the access token", func(t *testing.T) { + server := newServer(t, `{"access_token": "access", "token_type": "bearer", "expires_in": 300, "id_token": "id"}`) + defer server.Close() + orchestrator := newOrchestrator(t, server, oauth.TokenTypeBearer) + authToken, err := orchestrator.FetchTokenFromAuthFlow(ctx) + assert.NoError(t, err) + assert.Equal(t, "access", authToken.AccessToken) + assert.Equal(t, "Bearer", authToken.Type()) + }) + + t.Run("id token missing", func(t *testing.T) { + server := newServer(t, `{"access_token": "access", "token_type": "bearer", "expires_in": 300}`) + defer server.Close() + orchestrator := newOrchestrator(t, server, oauth.TokenTypeIDToken) + authToken, err := orchestrator.FetchTokenFromAuthFlow(ctx) + assert.Error(t, err) + assert.Nil(t, authToken) + }) +} diff --git a/flyteidl/clients/go/admin/oauth/config.go b/flyteidl/clients/go/admin/oauth/config.go index 94055f678a2..4bb550d905f 100644 --- a/flyteidl/clients/go/admin/oauth/config.go +++ b/flyteidl/clients/go/admin/oauth/config.go @@ -2,18 +2,70 @@ package oauth import ( "context" + "fmt" + "strings" "golang.org/x/oauth2" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" ) +const ( + // TokenTypeBearer sends the OAuth2 access token to admin as "Bearer ". This is the default. + TokenTypeBearer = "Bearer" + // TokenTypeIDToken sends the OIDC id_token from the token response to admin as "IDToken ". Admin validates + // tokens sent with this scheme against its configured userAuth OIDC provider rather than its authorization server. + TokenTypeIDToken = "IDToken" + + idTokenKey = "id_token" +) + // Config oauth2.Config overridden with device endpoint for supporting Device Authorization Grant flow [RFC8268] type Config struct { *oauth2.Config DeviceEndpoint string // Audience value to be passed when requesting access token using device flow.This needs to be passed in the first request of the device flow currently and is configured in admin public client config.Required when auth server hasn't been configured with default audience"` Audience string + // TokenType is the scheme used when sending the token to admin: TokenTypeBearer (default) or TokenTypeIDToken. + TokenType string +} + +// NormalizeTokenType returns the canonical token type for a configured value. An empty value means TokenTypeBearer. +func NormalizeTokenType(tokenType string) (string, error) { + switch strings.ToLower(strings.TrimSpace(tokenType)) { + case "", "bearer": + return TokenTypeBearer, nil + case "idtoken", "id_token": + return TokenTypeIDToken, nil + default: + return "", fmt.Errorf("unsupported tokenType %q, expected %v or %v", tokenType, TokenTypeBearer, TokenTypeIDToken) + } +} + +// PrepareToken adjusts a token returned by the authorization server so that it is sent to admin with the configured +// TokenType. For TokenTypeIDToken the id_token carried in the token response replaces the access token. +func (c *Config) PrepareToken(token *oauth2.Token) (*oauth2.Token, error) { + if token == nil { + return nil, nil + } + + idToken, _ := token.Extra(idTokenKey).(string) + return c.PrepareTokenWithIDToken(token, idToken) +} + +// PrepareTokenWithIDToken is PrepareToken for callers that parsed the token response themselves and hold the id_token. +func (c *Config) PrepareTokenWithIDToken(token *oauth2.Token, idToken string) (*oauth2.Token, error) { + if token == nil || c == nil || c.TokenType != TokenTypeIDToken { + return token, nil + } + + if len(idToken) == 0 { + return nil, fmt.Errorf("tokenType is %v but the token response has no id_token; make sure the openid scope is requested", TokenTypeIDToken) + } + + token.AccessToken = idToken + token.TokenType = TokenTypeIDToken + return token, nil } // BuildConfigFromMetadataService builds OAuth2 config from information retrieved through the anonymous auth metadata service. @@ -40,6 +92,7 @@ func BuildConfigFromMetadataService(ctx context.Context, authMetadataClient serv }, DeviceEndpoint: oauthMetaResp.GetDeviceAuthorizationEndpoint(), Audience: clientResp.GetAudience(), + TokenType: TokenTypeBearer, } return clientConf, nil diff --git a/flyteidl/clients/go/admin/oauth/config_test.go b/flyteidl/clients/go/admin/oauth/config_test.go index cf702809a96..4165d7b535c 100644 --- a/flyteidl/clients/go/admin/oauth/config_test.go +++ b/flyteidl/clients/go/admin/oauth/config_test.go @@ -2,10 +2,12 @@ package oauth import ( "context" + "fmt" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" + "golang.org/x/oauth2" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/mocks" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" @@ -39,3 +41,80 @@ func TestGenerateClientConfig(t *testing.T) { assert.Equal(t, "dummyDeviceEndpoint", oauthConfig.DeviceEndpoint) assert.Equal(t, "dummyAudience", oauthConfig.Audience) } + +const ( + testAccessToken = "access" + testIDToken = "id" + testBearerLower = "bearer" +) + +func TestNormalizeTokenType(t *testing.T) { + tests := []struct { + in string + expected string + wantErr bool + }{ + {in: "", expected: TokenTypeBearer}, + {in: testBearerLower, expected: TokenTypeBearer}, + {in: " Bearer ", expected: TokenTypeBearer}, + {in: "IDToken", expected: TokenTypeIDToken}, + {in: "idtoken", expected: TokenTypeIDToken}, + {in: idTokenKey, expected: TokenTypeIDToken}, + {in: "MAC", wantErr: true}, + } + for _, test := range tests { + t.Run(fmt.Sprintf("%q", test.in), func(t *testing.T) { + got, err := NormalizeTokenType(test.in) + if test.wantErr { + assert.Error(t, err) + return + } + assert.NoError(t, err) + assert.Equal(t, test.expected, got) + }) + } +} + +func TestPrepareToken(t *testing.T) { + t.Run("bearer leaves the token alone", func(t *testing.T) { + cfg := &Config{Config: &oauth2.Config{}, TokenType: TokenTypeBearer} + token := (&oauth2.Token{AccessToken: testAccessToken, TokenType: testBearerLower}).WithExtra(map[string]interface{}{idTokenKey: testIDToken}) + got, err := cfg.PrepareToken(token) + assert.NoError(t, err) + assert.Equal(t, testAccessToken, got.AccessToken) + assert.Equal(t, "Bearer", got.Type()) + }) + + t.Run("id token replaces the access token", func(t *testing.T) { + cfg := &Config{Config: &oauth2.Config{}, TokenType: TokenTypeIDToken} + token := (&oauth2.Token{AccessToken: testAccessToken, RefreshToken: "refresh", TokenType: testBearerLower}).WithExtra(map[string]interface{}{idTokenKey: testIDToken}) + got, err := cfg.PrepareToken(token) + assert.NoError(t, err) + assert.Equal(t, testIDToken, got.AccessToken) + assert.Equal(t, "refresh", got.RefreshToken) + assert.Equal(t, TokenTypeIDToken, got.TokenType) + assert.Equal(t, TokenTypeIDToken, got.Type()) + }) + + t.Run("id token missing from the response", func(t *testing.T) { + cfg := &Config{Config: &oauth2.Config{}, TokenType: TokenTypeIDToken} + got, err := cfg.PrepareToken(&oauth2.Token{AccessToken: testAccessToken}) + assert.Error(t, err) + assert.Nil(t, got) + }) + + t.Run("explicit id token", func(t *testing.T) { + cfg := &Config{Config: &oauth2.Config{}, TokenType: TokenTypeIDToken} + got, err := cfg.PrepareTokenWithIDToken(&oauth2.Token{AccessToken: testAccessToken}, testIDToken) + assert.NoError(t, err) + assert.Equal(t, testIDToken, got.AccessToken) + assert.Equal(t, TokenTypeIDToken, got.TokenType) + }) + + t.Run("nil token", func(t *testing.T) { + cfg := &Config{Config: &oauth2.Config{}, TokenType: TokenTypeIDToken} + got, err := cfg.PrepareToken(nil) + assert.NoError(t, err) + assert.Nil(t, got) + }) +} diff --git a/flyteidl/clients/go/admin/pkce/handle_app_call_back.go b/flyteidl/clients/go/admin/pkce/handle_app_call_back.go index 565a9d54909..2dc2ff2d9a5 100644 --- a/flyteidl/clients/go/admin/pkce/handle_app_call_back.go +++ b/flyteidl/clients/go/admin/pkce/handle_app_call_back.go @@ -51,6 +51,13 @@ func getAuthServerCallbackHandler(c *oauth.Config, codeVerifier string, tokenCha _, _ = rw.Write([]byte(fmt.Sprintf(`

Couldn't get access token due to error: %s

`, err.Error()))) return } + + token, err = c.PrepareToken(token) + if err != nil { + errorChannel <- err + _, _ = rw.Write([]byte(fmt.Sprintf(`

Couldn't get access token due to error: %s

`, err.Error()))) + return + } _, _ = rw.Write([]byte(`

Cool! Your authentication was successful and you can close the window.

`)) tokenChannel <- token } diff --git a/flyteidl/clients/go/admin/token_source_provider.go b/flyteidl/clients/go/admin/token_source_provider.go index 2b54ca98ed6..91eb99656bd 100644 --- a/flyteidl/clients/go/admin/token_source_provider.go +++ b/flyteidl/clients/go/admin/token_source_provider.go @@ -18,6 +18,7 @@ import ( "github.com/flyteorg/flyte/flyteidl/clients/go/admin/cache" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/deviceflow" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/externalprocess" + "github.com/flyteorg/flyte/flyteidl/clients/go/admin/oauth" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/pkce" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/tokenorchestrator" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/utils" @@ -43,10 +44,18 @@ type TokenSourceProvider interface { func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.TokenCache, authClient service.AuthMetadataServiceClient) (TokenSourceProvider, error) { + tokenType, err := oauth.NormalizeTokenType(cfg.TokenType) + if err != nil { + return nil, err + } + var tokenProvider TokenSourceProvider - var err error switch cfg.AuthType { case AuthTypeClientSecret: + if tokenType != oauth.TokenTypeBearer { + return nil, fmt.Errorf("tokenType %v is not supported with authType %v: the client credentials grant does not issue an id_token", tokenType, cfg.AuthType) + } + tokenURL := cfg.TokenURL if len(tokenURL) == 0 { metadata, err := authClient.GetOAuth2Metadata(ctx, &service.OAuth2MetadataRequest{}) @@ -80,7 +89,7 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T return nil, err } case AuthTypePkce: - baseTokenOrchestrator, err := tokenorchestrator.NewBaseTokenOrchestrator(ctx, tokenCache, authClient) + baseTokenOrchestrator, err := newBaseTokenOrchestrator(ctx, cfg, tokenType, tokenCache, authClient) if err != nil { return nil, err } @@ -90,12 +99,12 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T return nil, err } case AuthTypeExternalCommand: - tokenProvider, err = NewExternalTokenSourceProvider(cfg.Command) + tokenProvider, err = NewExternalTokenSourceProviderWithTokenType(cfg.Command, tokenType) if err != nil { return nil, err } case AuthTypeDeviceFlow: - baseTokenOrchestrator, err := tokenorchestrator.NewBaseTokenOrchestrator(ctx, tokenCache, authClient) + baseTokenOrchestrator, err := newBaseTokenOrchestrator(ctx, cfg, tokenType, tokenCache, authClient) if err != nil { return nil, err } @@ -111,12 +120,71 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T return tokenProvider, nil } +// newBaseTokenOrchestrator builds the OAuth2 client config shared by the Pkce and DeviceFlow auth types. Everything is +// discovered from admin's anonymously accessible metadata unless the endpoint for the auth type in use is configured +// explicitly (AuthorizationURL for Pkce, DeviceAuthorizationURL for DeviceFlow, each together with TokenURL), in which +// case those endpoints are used and the client id, scopes and audience come from cfg because the client has to be +// registered with that server. +func newBaseTokenOrchestrator(ctx context.Context, cfg *Config, tokenType string, tokenCache cache.TokenCache, + authClient service.AuthMetadataServiceClient) (tokenorchestrator.BaseTokenOrchestrator, error) { + orchestrator, err := tokenorchestrator.NewBaseTokenOrchestrator(ctx, tokenCache, authClient) + if err != nil { + return tokenorchestrator.BaseTokenOrchestrator{}, err + } + + var endpointKey, endpoint string + switch cfg.AuthType { + case AuthTypePkce: + endpointKey, endpoint = "authorizationUrl", cfg.AuthorizationURL + case AuthTypeDeviceFlow: + endpointKey, endpoint = "deviceAuthorizationUrl", cfg.DeviceAuthorizationURL + } + + if len(endpoint) > 0 { + if len(cfg.TokenURL) == 0 { + return tokenorchestrator.BaseTokenOrchestrator{}, fmt.Errorf("tokenUrl must be configured together with %v", endpointKey) + } + + if len(cfg.Scopes) == 0 { + return tokenorchestrator.BaseTokenOrchestrator{}, fmt.Errorf("scopes must be configured together with %v", endpointKey) + } + + if cfg.ClientID == DefaultClientID { + logger.Warnf(ctx, "clientId is the default %q; with %v it must be a client registered with that authorization server", DefaultClientID, endpointKey) + } + + orchestrator.ClientConfig.Endpoint.TokenURL = cfg.TokenURL + orchestrator.ClientConfig.Endpoint.AuthURL = cfg.AuthorizationURL + orchestrator.ClientConfig.DeviceEndpoint = cfg.DeviceAuthorizationURL + orchestrator.ClientConfig.ClientID = cfg.ClientID + orchestrator.ClientConfig.Scopes = cfg.Scopes + if len(cfg.Audience) > 0 { + orchestrator.ClientConfig.Audience = cfg.Audience + } + } + + orchestrator.ClientConfig.TokenType = tokenType + return orchestrator, nil +} + type ExternalTokenSourceProvider struct { - command []string + command []string + tokenType string } func NewExternalTokenSourceProvider(command []string) (TokenSourceProvider, error) { - return &ExternalTokenSourceProvider{command: command}, nil + return NewExternalTokenSourceProviderWithTokenType(command, oauth.TokenTypeBearer) +} + +// NewExternalTokenSourceProviderWithTokenType returns a provider that sends the command's output with the given +// token type (see oauth.NormalizeTokenType). +func NewExternalTokenSourceProviderWithTokenType(command []string, tokenType string) (TokenSourceProvider, error) { + tokenType, err := oauth.NormalizeTokenType(tokenType) + if err != nil { + return nil, err + } + + return &ExternalTokenSourceProvider{command: command, tokenType: tokenType}, nil } func (e ExternalTokenSourceProvider) GetTokenSource(ctx context.Context) (oauth2.TokenSource, error) { @@ -127,7 +195,7 @@ func (e ExternalTokenSourceProvider) GetTokenSource(ctx context.Context) (oauth2 return oauth2.StaticTokenSource(&oauth2.Token{ AccessToken: strings.Trim(string(output), "\t \n"), - TokenType: "bearer", + TokenType: e.tokenType, }), nil } diff --git a/flyteidl/clients/go/admin/token_source_provider_test.go b/flyteidl/clients/go/admin/token_source_provider_test.go index fb439b0ec77..10a4aa84259 100644 --- a/flyteidl/clients/go/admin/token_source_provider_test.go +++ b/flyteidl/clients/go/admin/token_source_provider_test.go @@ -13,6 +13,7 @@ import ( tokenCacheMocks "github.com/flyteorg/flyte/flyteidl/clients/go/admin/cache/mocks" adminMocks "github.com/flyteorg/flyte/flyteidl/clients/go/admin/mocks" + "github.com/flyteorg/flyte/flyteidl/clients/go/admin/oauth" "github.com/flyteorg/flyte/flyteidl/clients/go/admin/utils" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service" ) @@ -163,3 +164,150 @@ func TestCustomTokenSource_Token(t *testing.T) { }) } } + +const ( + testIdpClientID = "flytectl" + testAdminAuthURL = "https://admin/oauth2/authorize" + testAdminTokenURL = "https://admin/oauth2/token" //nolint:gosec // test fixture URL, not a credential + testScopeOpenID = "openid" + testScopeOffline = "offline" + testCommandArgument = "token-from-command" + testEcho = "echo" + testIdpAuthURL = "https://idp/auth" + testIdpTokenURL = "https://idp/token" //nolint:gosec // test fixture URL, not a credential + testIdpDeviceURL = "https://idp/device/code" + testIgnored = "ignored" +) + +func TestNewTokenSourceProvider_TokenType(t *testing.T) { + ctx := context.Background() + + t.Run("external command defaults to bearer", func(t *testing.T) { + cfg := &Config{AuthType: AuthTypeExternalCommand, Command: []string{testEcho, testCommandArgument}} + provider, err := NewTokenSourceProvider(ctx, cfg, nil, nil) + assert.NoError(t, err) + tokenSource, err := provider.GetTokenSource(ctx) + assert.NoError(t, err) + token, err := tokenSource.Token() + assert.NoError(t, err) + assert.Equal(t, testCommandArgument, token.AccessToken) + assert.Equal(t, "Bearer", token.Type()) + }) + + t.Run("external command with IDToken", func(t *testing.T) { + cfg := &Config{AuthType: AuthTypeExternalCommand, Command: []string{testEcho, "id-token-from-command"}, TokenType: oauth.TokenTypeIDToken} + provider, err := NewTokenSourceProvider(ctx, cfg, nil, nil) + assert.NoError(t, err) + tokenSource, err := provider.GetTokenSource(ctx) + assert.NoError(t, err) + token, err := tokenSource.Token() + assert.NoError(t, err) + assert.Equal(t, "id-token-from-command", token.AccessToken) + assert.Equal(t, oauth.TokenTypeIDToken, token.Type()) + }) + + t.Run("unsupported token type", func(t *testing.T) { + cfg := &Config{AuthType: AuthTypeExternalCommand, Command: []string{testEcho, "x"}, TokenType: "MAC"} + _, err := NewTokenSourceProvider(ctx, cfg, nil, nil) + assert.Error(t, err) + }) + + t.Run("device flow against explicitly configured endpoints", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{ + AuthorizationEndpoint: testAdminAuthURL, TokenEndpoint: testAdminTokenURL}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{ + ClientId: "flytectl-on-admin", RedirectUri: "http://localhost:53593/callback", Scopes: []string{testScopeOffline, "all"}, Audience: "admin-audience"}, nil) + + cfg := &Config{ + AuthType: AuthTypeDeviceFlow, + DeviceAuthorizationURL: testIdpDeviceURL, + TokenURL: testIdpTokenURL, + ClientID: "flytectl-on-idp", + Scopes: []string{testScopeOpenID, "offline_access"}, + TokenType: oauth.TokenTypeIDToken, + } + provider, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.NoError(t, err) + deviceFlowProvider, ok := provider.(DeviceFlowTokenSourceProvider) + assert.True(t, ok) + clientConfig := deviceFlowProvider.tokenOrchestrator.ClientConfig + assert.Equal(t, testIdpDeviceURL, clientConfig.DeviceEndpoint) + assert.Equal(t, testIdpTokenURL, clientConfig.Endpoint.TokenURL) + assert.Equal(t, "flytectl-on-idp", clientConfig.ClientID) + assert.Equal(t, []string{testScopeOpenID, "offline_access"}, clientConfig.Scopes) + assert.Equal(t, "admin-audience", clientConfig.Audience) + assert.Equal(t, "http://localhost:53593/callback", clientConfig.RedirectURL) + assert.Equal(t, oauth.TokenTypeIDToken, clientConfig.TokenType) + }) + + t.Run("pkce against explicitly configured endpoints", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{}, nil) + cfg := &Config{AuthType: AuthTypePkce, AuthorizationURL: testIdpAuthURL, TokenURL: testIdpTokenURL, ClientID: testIdpClientID, Scopes: []string{testScopeOpenID}, Audience: "idp-audience"} + provider, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.NoError(t, err) + clientConfig := provider.(PKCETokenSourceProvider).tokenOrchestrator.ClientConfig + assert.Equal(t, testIdpAuthURL, clientConfig.Endpoint.AuthURL) + assert.Equal(t, testIdpTokenURL, clientConfig.Endpoint.TokenURL) + assert.Equal(t, testIdpClientID, clientConfig.ClientID) + assert.Equal(t, "idp-audience", clientConfig.Audience) + assert.Equal(t, oauth.TokenTypeBearer, clientConfig.TokenType) + }) + + t.Run("client secret rejects IDToken", func(t *testing.T) { + cfg := &Config{AuthType: AuthTypeClientSecret, TokenType: oauth.TokenTypeIDToken, TokenURL: testAdminTokenURL, Scopes: []string{"all"}} + _, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, &adminMocks.AuthMetadataServiceClient{}) + assert.Error(t, err) + }) + + t.Run("explicit endpoints require tokenUrl", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{}, nil) + cfg := &Config{AuthType: AuthTypeDeviceFlow, DeviceAuthorizationURL: testIdpDeviceURL, ClientID: testIdpClientID, Scopes: []string{testScopeOpenID}} + _, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.ErrorContains(t, err, "tokenUrl") + }) + + t.Run("explicit endpoints require scopes", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{}, nil) + cfg := &Config{AuthType: AuthTypePkce, AuthorizationURL: testIdpAuthURL, TokenURL: testIdpTokenURL, ClientID: testIdpClientID} + _, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.ErrorContains(t, err, "scopes") + }) + + t.Run("a stale tokenUrl alone does not change pkce or device flow", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{ + TokenEndpoint: testAdminTokenURL, DeviceAuthorizationEndpoint: "https://admin/oauth2/device"}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{ + ClientId: testIdpClientID, Scopes: []string{testScopeOffline, "all"}}, nil) + cfg := &Config{AuthType: AuthTypeDeviceFlow, TokenURL: testIdpTokenURL, ClientID: testIgnored, Scopes: []string{testIgnored}} + provider, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.NoError(t, err) + clientConfig := provider.(DeviceFlowTokenSourceProvider).tokenOrchestrator.ClientConfig + assert.Equal(t, testAdminTokenURL, clientConfig.Endpoint.TokenURL) + assert.Equal(t, "https://admin/oauth2/device", clientConfig.DeviceEndpoint) + assert.Equal(t, testIdpClientID, clientConfig.ClientID) + }) + + t.Run("admin discovery is unchanged without a separate authorization server", func(t *testing.T) { + metadataClient := &adminMocks.AuthMetadataServiceClient{} + metadataClient.EXPECT().GetOAuth2Metadata(mock.Anything, mock.Anything).Return(&service.OAuth2MetadataResponse{ + TokenEndpoint: testAdminTokenURL, DeviceAuthorizationEndpoint: "https://admin/oauth2/device"}, nil) + metadataClient.EXPECT().GetPublicClientConfig(mock.Anything, mock.Anything).Return(&service.PublicClientAuthConfigResponse{ + ClientId: testIdpClientID, Scopes: []string{testScopeOffline, "all"}}, nil) + cfg := &Config{AuthType: AuthTypeDeviceFlow, ClientID: testIgnored, Scopes: []string{testIgnored}} + provider, err := NewTokenSourceProvider(ctx, cfg, &tokenCacheMocks.TokenCache{}, metadataClient) + assert.NoError(t, err) + clientConfig := provider.(DeviceFlowTokenSourceProvider).tokenOrchestrator.ClientConfig + assert.Equal(t, testIdpClientID, clientConfig.ClientID) + assert.Equal(t, []string{testScopeOffline, "all"}, clientConfig.Scopes) + assert.Equal(t, "https://admin/oauth2/device", clientConfig.DeviceEndpoint) + assert.Equal(t, oauth.TokenTypeBearer, clientConfig.TokenType) + }) +} diff --git a/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator.go b/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator.go index 441127ce073..11c43d04f6c 100644 --- a/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator.go +++ b/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator.go @@ -31,6 +31,12 @@ func (t BaseTokenOrchestrator) RefreshToken(ctx context.Context, token *oauth2.T return nil, err } + refreshedToken, err = t.ClientConfig.PrepareToken(refreshedToken) + if err != nil { + logger.Warnf(ctx, "refreshed token is unusable due to %v and will be doing re-auth", err) + return nil, err + } + if refreshedToken != nil { logger.Debugf(ctx, "got a response from the refresh grant for old expiry %v with new expiry %v", token.Expiry, refreshedToken.Expiry) diff --git a/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator_test.go b/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator_test.go index 6247dadda63..53bd38d10d8 100644 --- a/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator_test.go +++ b/flyteidl/clients/go/admin/tokenorchestrator/base_token_orchestrator_test.go @@ -2,6 +2,8 @@ package tokenorchestrator import ( "context" + "net/http" + "net/http/httptest" "testing" "time" @@ -115,3 +117,33 @@ func TestFetchFromCache(t *testing.T) { mockTokenCacheProvider.AssertNotCalled(t, "SaveToken") }) } + +func TestRefreshTheTokenWithIDToken(t *testing.T) { + ctx := context.Background() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.NoError(t, r.ParseForm()) + assert.Equal(t, "refresh_token", r.Form.Get("grant_type")) + assert.Equal(t, "old-refresh", r.Form.Get("refresh_token")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token": "new-access", "refresh_token": "new-refresh", "token_type": "bearer", "expires_in": 3600, "id_token": "new-id"}`)) + })) + defer server.Close() + + tokenCacheProvider := cache.NewTokenCacheInMemoryProvider() + orchestrator := BaseTokenOrchestrator{ + ClientConfig: &oauth.Config{ + Config: &oauth2.Config{ClientID: "flytectl", Endpoint: oauth2.Endpoint{TokenURL: server.URL}}, + TokenType: oauth.TokenTypeIDToken, + }, + TokenCache: tokenCacheProvider, + } + expired := &oauth2.Token{AccessToken: "old-id", RefreshToken: "old-refresh", TokenType: oauth.TokenTypeIDToken, Expiry: time.Now().Add(-time.Hour)} + refreshed, err := orchestrator.RefreshToken(ctx, expired) + assert.NoError(t, err) + assert.Equal(t, "new-id", refreshed.AccessToken) + assert.Equal(t, "new-refresh", refreshed.RefreshToken) + assert.Equal(t, oauth.TokenTypeIDToken, refreshed.Type()) + cached, err := tokenCacheProvider.GetToken() + assert.NoError(t, err) + assert.Equal(t, "new-id", cached.AccessToken) +}