Let clients run the device flow against the OIDC provider and send ID tokens

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 <https://datatracker.ietf.org/doc/html/rfc8628>`_), 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:<flyteadmin 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 <jwt>" instead of the access token as "Bearer <jwt>".
+     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:<flyteadmin 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
 ***************************

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

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
 }

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)

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)
+	})
+}

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 <token>". This is the default.
+	TokenTypeBearer = "Bearer"
+	// TokenTypeIDToken sends the OIDC id_token from the token response to admin as "IDToken <token>". 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