diff --git a/acceptance/cmd/auth/login/device-code/client-not-allowed/out.test.toml b/acceptance/cmd/auth/login/device-code/client-not-allowed/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/client-not-allowed/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/auth/login/device-code/client-not-allowed/output.txt b/acceptance/cmd/auth/login/device-code/client-not-allowed/output.txt new file mode 100644 index 00000000000..4b880689589 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/client-not-allowed/output.txt @@ -0,0 +1,3 @@ + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code --client-id custom-client-id +Error: device authorization: Client is not allowed to use device code login (error code: unauthorized_client) diff --git a/acceptance/cmd/auth/login/device-code/client-not-allowed/script b/acceptance/cmd/auth/login/device-code/client-not-allowed/script new file mode 100644 index 00000000000..4197d2fdf24 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/client-not-allowed/script @@ -0,0 +1,3 @@ +sethome "./home" + +trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code --client-id custom-client-id diff --git a/acceptance/cmd/auth/login/device-code/nominal/out.databrickscfg b/acceptance/cmd/auth/login/device-code/nominal/out.databrickscfg new file mode 100644 index 00000000000..02b3737c466 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/nominal/out.databrickscfg @@ -0,0 +1,16 @@ +; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified. +[DEFAULT] + +[test] +host = [DATABRICKS_URL] +workspace_id = [WORKSPACE_ID] +auth_type = databricks-cli + +[__settings__] +default_profile = test + +[scoped] +host = [DATABRICKS_URL] +workspace_id = [WORKSPACE_ID] +scopes = sql +auth_type = databricks-cli diff --git a/acceptance/cmd/auth/login/device-code/nominal/out.test.toml b/acceptance/cmd/auth/login/device-code/nominal/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/nominal/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/auth/login/device-code/nominal/output.txt b/acceptance/cmd/auth/login/device-code/nominal/output.txt new file mode 100644 index 00000000000..2216533b304 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/nominal/output.txt @@ -0,0 +1,56 @@ + +=== Login with a host + +>>> [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Profile test was successfully saved + +>>> [CLI] auth token --profile test +json.access_token = "oauth-token"; + +>>> [CLI] auth profiles +Name Host Valid +test (Default) [DATABRICKS_URL] YES + +=== Login again with the host from the saved profile + +>>> [CLI] auth login --profile test --use-device-code +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Profile test was successfully saved + +>>> [CLI] auth token --profile test +json.access_token = "oauth-token"; + +>>> print_requests.py --oneline //oidc/v1/device/auth //oidc/v1/token +{"method": "POST", "path": "/oidc/v1/device/auth", "raw_body": "client_id=databricks-cli&scope=offline_access+all-apis"} +{"method": "POST", "path": "/oidc/v1/token", "raw_body": "client_id=databricks-cli&confirmation_code=FTFAQQKY&device_code=device-code-1&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code"} +{"method": "POST", "path": "/oidc/v1/device/auth", "raw_body": "client_id=databricks-cli&scope=offline_access+all-apis"} +{"method": "POST", "path": "/oidc/v1/token", "raw_body": "client_id=databricks-cli&confirmation_code=FTFAQQKY&device_code=device-code-2&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code"} + +=== Login with custom scopes + +>>> [CLI] auth login --host [DATABRICKS_URL] --profile scoped --use-device-code --scopes sql +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Profile scoped was successfully saved + +>>> print_requests.py --oneline //oidc/v1/device/auth //oidc/v1/token +{"method": "POST", "path": "/oidc/v1/device/auth", "raw_body": "client_id=databricks-cli&scope=offline_access+sql"} +{"method": "POST", "path": "/oidc/v1/token", "raw_body": "client_id=databricks-cli&confirmation_code=FTFAQQKY&device_code=device-code-3&grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code"} diff --git a/acceptance/cmd/auth/login/device-code/nominal/script b/acceptance/cmd/auth/login/device-code/nominal/script new file mode 100644 index 00000000000..f1d6c04f59a --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/nominal/script @@ -0,0 +1,20 @@ +sethome "./home" + +title "Login with a host\n" +printf "FTFA-QQKY\n" | trace $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code +trace $CLI auth token --profile test | gron.py | grep access_token +trace $CLI auth profiles + +title "Login again with the host from the saved profile\n" +printf "FTFA-QQKY\n" | trace $CLI auth login --profile test --use-device-code +trace $CLI auth token --profile test | gron.py | grep access_token +trace print_requests.py --oneline //oidc/v1/device/auth //oidc/v1/token + +# The token gets the scopes requested at device authorization, so --scopes is +# sent there and not with the token request. +title "Login with custom scopes\n" +printf "FTFA-QQKY\n" | trace $CLI auth login --host $DATABRICKS_HOST --profile scoped --use-device-code --scopes sql +trace print_requests.py --oneline //oidc/v1/device/auth //oidc/v1/token + +# Track the .databrickscfg file that was created to surface changes. +mv "./home/.databrickscfg" "./out.databrickscfg" diff --git a/acceptance/cmd/auth/login/device-code/nominal/test.toml b/acceptance/cmd/auth/login/device-code/nominal/test.toml new file mode 100644 index 00000000000..159efe02696 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/nominal/test.toml @@ -0,0 +1 @@ +RecordRequests = true diff --git a/acceptance/cmd/auth/login/device-code/precondition/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/device-code/precondition/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..328fc4d2cc4 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/precondition/home/.databrickscfg.tmpl @@ -0,0 +1,4 @@ +[scoped] +host = $DATABRICKS_HOST +auth_type = databricks-cli +resources = https://resource.test diff --git a/acceptance/cmd/auth/login/device-code/precondition/out.test.toml b/acceptance/cmd/auth/login/device-code/precondition/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/precondition/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/auth/login/device-code/precondition/output.txt b/acceptance/cmd/auth/login/device-code/precondition/output.txt new file mode 100644 index 00000000000..0b34cd722e0 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/precondition/output.txt @@ -0,0 +1,16 @@ + +=== Device code login is disabled by default + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Error: --use-device-code is experimental. Set DATABRICKS_CLI_DEVICE_CODE=true to enable it + +=== Inputs device code login cannot honor are rejected + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code --resource https://resource.test +Error: --use-device-code cannot be combined with --resource + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code --skip-workspace +Error: --use-device-code cannot be combined with --skip-workspace + +>>> musterr [CLI] auth login --profile scoped --use-device-code +Error: --use-device-code cannot be used with profile scoped because its login is scoped to resources. Log in without --use-device-code diff --git a/acceptance/cmd/auth/login/device-code/precondition/script b/acceptance/cmd/auth/login/device-code/precondition/script new file mode 100644 index 00000000000..c3a7a3208c0 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/precondition/script @@ -0,0 +1,15 @@ +sethome "./home" + +title "Device code login is disabled by default\n" +( + unset DATABRICKS_CLI_DEVICE_CODE + trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code +) + +title "Inputs device code login cannot honor are rejected\n" +trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code --resource https://resource.test +trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code --skip-workspace + +# A profile scoped to resources is rejected like --resource. +envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg" +trace musterr $CLI auth login --profile scoped --use-device-code diff --git a/acceptance/cmd/auth/login/device-code/test.toml b/acceptance/cmd/auth/login/device-code/test.toml new file mode 100644 index 00000000000..c8487e95d6e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/test.toml @@ -0,0 +1,4 @@ +Env.DATABRICKS_CLI_DEVICE_CODE = "true" +Ignore = [ + "home" +] diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/device-code/unsupported-host/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..90baae2ed29 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/home/.databrickscfg.tmpl @@ -0,0 +1,6 @@ +[workspace] +host = $DATABRICKS_HOST +auth_type = databricks-cli + +[account] +host = https://accounts.cloud.databricks.test diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/out.test.toml b/acceptance/cmd/auth/login/device-code/unsupported-host/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt b/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt new file mode 100644 index 00000000000..5d327ce2db2 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt @@ -0,0 +1,31 @@ + +=== No host + +>>> musterr [CLI] auth login --profile new-profile --use-device-code +Error: --use-device-code requires a workspace host. Pass it as an argument or with --host + +=== Account host + +>>> musterr [CLI] auth login --host https://accounts.cloud.databricks.test --profile test --use-device-code +Warn: [hostmetadata] failed to fetch host metadata for https://accounts.cloud.databricks.test, will skip for 1m0s +Error: --use-device-code is only supported for workspace hosts + +>>> musterr [CLI] auth login --profile account --use-device-code +Error: --use-device-code is only supported for workspace hosts + +=== Unified host + +>>> DATABRICKS_DISCOVERY_URL=[DATABRICKS_URL]/oidc/accounts/spog-acct-123/.well-known/oauth-authorization-server musterr [CLI] auth login --host [DATABRICKS_URL] --use-device-code +Error: --use-device-code is only supported for workspace hosts + +=== Workspace without device code login + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Error: device code login is not enabled for this workspace + +>>> musterr [CLI] auth login --profile workspace --use-device-code +Error: device code login is not enabled for this workspace + +=== Device authorization requests + +>>> print_requests.py //oidc/v1/device/auth diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/script b/acceptance/cmd/auth/login/device-code/unsupported-host/script new file mode 100644 index 00000000000..dd8a9ddca47 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/script @@ -0,0 +1,27 @@ +sethome "./home" + +title "No host\n" +# The profile does not exist yet and no host is given. auth login does not read +# DATABRICKS_HOST, so the usual login would fall back to the browser-based +# discovery flow, which device code login does not support. +trace musterr $CLI auth login --profile new-profile --use-device-code + +envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg" + +# The host is rejected before the CLI would prompt for an account ID, whether +# it is passed directly or comes from a profile without one. +title "Account host\n" +trace musterr $CLI auth login --host https://accounts.cloud.databricks.test --profile test --use-device-code +trace musterr $CLI auth login --profile account --use-device-code + +# An account-scoped discovery URL marks the host as a unified host. +title "Unified host\n" +trace DATABRICKS_DISCOVERY_URL="$DATABRICKS_HOST/oidc/accounts/spog-acct-123/.well-known/oauth-authorization-server" musterr $CLI auth login --host $DATABRICKS_HOST --use-device-code + +title "Workspace without device code login\n" +trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code +trace musterr $CLI auth login --profile workspace --use-device-code + +# Every host is rejected before the device authorization request. +title "Device authorization requests\n" +trace print_requests.py //oidc/v1/device/auth diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/test.toml b/acceptance/cmd/auth/login/device-code/unsupported-host/test.toml new file mode 100644 index 00000000000..9eb77236bab --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/test.toml @@ -0,0 +1,12 @@ +RecordRequests = true + +# Workspaces that are not enrolled in the device code preview do not advertise +# a device_authorization_endpoint in their OIDC discovery document. +[[Server]] +Pattern = "GET /oidc/.well-known/oauth-authorization-server" +Response.Body = ''' +{ + "authorization_endpoint": "https://workspace.test/oidc/v1/authorize", + "token_endpoint": "https://workspace.test/oidc/v1/token" +} +''' diff --git a/acceptance/cmd/auth/login/device-code/wrong-code/out.test.toml b/acceptance/cmd/auth/login/device-code/wrong-code/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/wrong-code/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/auth/login/device-code/wrong-code/output.txt b/acceptance/cmd/auth/login/device-code/wrong-code/output.txt new file mode 100644 index 00000000000..2ba7f3ea23e --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/wrong-code/output.txt @@ -0,0 +1,40 @@ + +=== Login fails when no confirmation code is provided + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Error: reading confirmation code: EOF + +=== Login fails after 3 wrong confirmation codes + +>>> musterr [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Invalid confirmation code (error code: invalid_request). Enter the confirmation code again (2 attempts left). +Confirmation code: Invalid confirmation code (error code: invalid_request). Enter the confirmation code again (1 attempt left). +Confirmation code: Error: Invalid confirmation code (error code: invalid_request). Start login from scratch + +=== Login succeeds after re-entering the confirmation code + +>>> [CLI] auth login --host [DATABRICKS_URL] --profile test --use-device-code +Using a browser on another device, visit: +[DATABRICKS_URL]/device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: Invalid confirmation code (error code: invalid_request). Enter the confirmation code again (2 attempts left). +Confirmation code: Invalid confirmation code (error code: invalid_request). Enter the confirmation code again (1 attempt left). +Confirmation code: Profile test was successfully saved diff --git a/acceptance/cmd/auth/login/device-code/wrong-code/script b/acceptance/cmd/auth/login/device-code/wrong-code/script new file mode 100644 index 00000000000..3adb0de81a3 --- /dev/null +++ b/acceptance/cmd/auth/login/device-code/wrong-code/script @@ -0,0 +1,13 @@ +sethome "./home" + +title "Login fails when no confirmation code is provided\n" +trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code < /dev/null + +title "Login fails after 3 wrong confirmation codes\n" +printf "AAAA-AAAA\nBBBB-BBBB\nCCCC-CCCC\n" | trace musterr $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code + +# A wrong confirmation code leaves the device code valid, so the user can +# enter it again. A blank line counts as a wrong confirmation code. +# The confirmation code is accepted the way the browser displays it, in any case. +title "Login succeeds after re-entering the confirmation code\n" +printf "AAAA-AAAA\n\nftfa-qqky\n" | trace $CLI auth login --host $DATABRICKS_HOST --profile test --use-device-code diff --git a/cmd/auth/README.md b/cmd/auth/README.md index 3185ee63d66..11d43bc14c4 100644 --- a/cmd/auth/README.md +++ b/cmd/auth/README.md @@ -49,3 +49,58 @@ sequenceDiagram CLI -X User: no auth configured end ``` + +# Device code login (experimental) + +`databricks auth login --use-device-code` logs in from a machine without a browser. +It is experimental: the flag is hidden and requires +`DATABRICKS_CLI_DEVICE_CODE=true`. Only workspace hosts are supported. See +[docs/device-code-login.md](../../docs/device-code-login.md) for the user guide. + +Databricks implements a variant of [RFC 8628](https://datatracker.ietf.org/doc/html/rfc8628): +instead of polling the token endpoint until the user approves the login, the +browser shows the user a confirmation code, which the CLI sends together with +the device code in a single token request. + +```mermaid +sequenceDiagram + autonumber + actor User + + User ->> CLI: type `databricks auth login HOST --use-device-code` + CLI ->>+ HOST: request OIDC endpoints + HOST ->>- CLI: token & device_authorization endpoints + CLI ->>+ Device Auth Endpoint: client_id + scopes + Device Auth Endpoint ->>- CLI: DEVICE CODE + USER CODE + verification URI + CLI ->> User: show verification URI + USER CODE + + User ->>+ Verification Page: open on another device, go through SSO + User ->> Verification Page: enter USER CODE and approve + Verification Page ->>- User: CONFIRMATION CODE + + User ->> CLI: type CONFIRMATION CODE + CLI ->>+ Token Endpoint: Exchange: DEVICE CODE + CONFIRMATION CODE + Token Endpoint ->>- CLI: Access Token (JWT) + refresh + expiry + CLI ->> Token cache: Save Access Token (JWT) + refresh + expiry + CLI ->> User: success +``` + +Notes on the server behavior the CLI relies on: + +- Workspaces advertise `device_authorization_endpoint` in + `/oidc/.well-known/oauth-authorization-server` only when device code login + is enabled for them. The CLI fails before requesting a device code if it is + missing. +- The token request carries the confirmation code in the `confirmation_code` + form parameter. A wrong confirmation code does not consume the device code; + a successful exchange does. The server answers a wrong confirmation code + with HTTP 400 and `invalid_request`, and the CLI lets the user enter it up + to 3 times in total. Any other error ends the login, including an unknown + device code, which the server answers the same way but with the + description `Invalid device code`. +- Approving the login in the browser extends the device code's lifetime, so + the CLI does not enforce the `expires_in` from the device authorization + response. +- The refresh token is single-use (rotating) and has an absolute lifetime of + 24 hours by default. After that, `databricks auth token` fails with a + refresh error and the user needs to log in again. diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 6019e8b90a0..13a5a8a61e0 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -49,6 +49,9 @@ const ( // host used by the discovery login flow. Intended for testing and // development against non-production environments. discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST" + // deviceCodeEnvVar enables the experimental device code login flow + // (--use-device-code) while it is experimental. + deviceCodeEnvVar = "DATABRICKS_CLI_DEVICE_CODE" ) // discoveryErr wraps an error (or creates a new one) and appends the @@ -151,6 +154,7 @@ a new profile is created. var scopes string var clientID string var resources []string + var useDeviceCode bool cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout, "Timeout for completing login challenge in the browser") cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false, @@ -165,6 +169,10 @@ a new profile is created. "OAuth client ID to use for U2M authentication") cmd.Flags().StringArrayVar(&resources, "resource", nil, "RFC 8707 resource indicator to scope the login to (repeatable). Requires --host.") + cmd.Flags().BoolVar(&useDeviceCode, "use-device-code", false, + "Log in with a device code, using a browser on another device (experimental)") + // Hidden while device code login is experimental. + cmd.Flags().MarkHidden("use-device-code") cmd.PreRunE = profileHostConflictCheck @@ -197,6 +205,12 @@ a new profile is created. } } + if useDeviceCode { + if err := validateDeviceCodeLogin(ctx, cmd); err != nil { + return err + } + } + // Resolve the cache before the browser step so an unavailable // keyring surfaces here rather than after OAuth. The probe also // triggers the OS unlock prompt, which the user can answer during @@ -284,11 +298,19 @@ a new profile is created. } if len(resources) == 0 { resources = u2mResourcesFromProfile(existingProfile) + // Device code login cannot scope the token to these resources, so + // the profile would claim a scoping its token does not have. + if useDeviceCode && len(resources) > 0 { + return fmt.Errorf("--use-device-code cannot be used with profile %s because its login is scoped to resources. Log in without --use-device-code", profileName) + } } // If no host is available from any source, use the discovery flow // via login.databricks.com. if shouldUseDiscovery(authArguments.Host, args, existingProfile) { + if useDeviceCode { + return errDeviceCodeNeedsHost + } if err := validateDiscoveryFlagCompatibility(cmd); err != nil { return err } @@ -305,7 +327,18 @@ a new profile is created. }) } - err = setHostAndAccountId(ctx, existingProfile, authArguments, args) + err = setHost(ctx, existingProfile, authArguments, args) + if err != nil { + return err + } + // Check the host once it is resolved, before setAccountID prompts for + // an account ID for account and unified hosts. + if useDeviceCode { + if err := validateDeviceCodeHost(authArguments); err != nil { + return err + } + } + err = setAccountID(ctx, existingProfile, authArguments) if err != nil { return err } @@ -350,7 +383,12 @@ a new profile is created. } defer persistentAuth.Close() - token, err := persistentAuth.Challenge() + var token *oauth2.Token + if useDeviceCode { + token, err = persistentAuth.DeviceCodeChallenge(ctx, promptForConfirmationCode) + } else { + token, err = persistentAuth.Challenge() + } if err != nil { return err } @@ -480,18 +518,20 @@ func shouldPromptWorkspace(authArguments *auth.AuthArguments, existingProfile *p return true } +// setHostAndAccountId calls setHost and then setAccountID. +func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error { + if err := setHost(ctx, existingProfile, authArguments, args); err != nil { + return err + } + return setAccountID(ctx, existingProfile, authArguments) +} + // Sets the host in the persistentAuth object based on the provided arguments and flags. // Follows the following precedence: // 1. [HOST] (first positional argument) or --host flag. Error if both are specified. // 2. Profile host, if available. // 3. Prompt the user for the host. -// -// Set the account in the persistentAuth object based on the flags. -// Follows the following precedence: -// 1. --account-id flag. -// 2. account-id from the specified profile, if available. -// 3. Prompt the user for the account-id. -func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error { +func setHost(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error { // If both [HOST] and --host are provided, return an error. host := authArguments.Host if len(args) > 0 && host != "" { @@ -541,7 +581,15 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, // .well-known/databricks-config endpoint. This is best-effort: failures // are logged as warnings and never block login. runHostDiscovery(ctx, authArguments) + return nil +} +// Set the account in the persistentAuth object based on the flags. +// Follows the following precedence: +// 1. --account-id flag. +// 2. account-id from the specified profile, if available. +// 3. Prompt the user for the account-id. +func setAccountID(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments) error { if needsAccountIDPrompt(authArguments.Host, authArguments.DiscoveryURL) && authArguments.AccountID == "" { if existingProfile != nil && existingProfile.AccountID != "" { authArguments.AccountID = existingProfile.AccountID @@ -561,6 +609,13 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, // for OAuth URL construction. True for classic account hosts (accounts.*) and // for unified hosts detected via account-scoped DiscoveryURL. func needsAccountIDPrompt(host, discoveryURL string) bool { + return isAccountOrUnifiedHost(host, discoveryURL) +} + +// isAccountOrUnifiedHost reports whether host is a classic account host +// (accounts.* or accounts-dod.*) or a unified host with an account-scoped +// discoveryURL. +func isAccountOrUnifiedHost(host, discoveryURL string) bool { canonicalHost := (&config.Config{Host: host}).CanonicalHostName() if auth.IsClassicAccountHost(canonicalHost) { return true diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go new file mode 100644 index 00000000000..93455ad9dd7 --- /dev/null +++ b/cmd/auth/login_device_code.go @@ -0,0 +1,96 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/env" + "github.com/spf13/cobra" +) + +// deviceCodeIncompatibleFlags lists flags that device code login cannot honor. +var deviceCodeIncompatibleFlags = []string{ + // Resource indicators are sent on the browser authorization request, + // which the device code flow does not make. + "resource", + // Device code login is only supported for workspace hosts. + "skip-workspace", +} + +var ( + errDeviceCodeNeedsHost = errors.New("--use-device-code requires a workspace host. Pass it as an argument or with --host") + errDeviceCodeWorkspaceOnly = errors.New("--use-device-code is only supported for workspace hosts") +) + +// validateDeviceCodeHost returns an error if the host resolved by setHost is +// an account or unified host. A workspace without device code login enabled +// is rejected by DeviceCodeChallenge before it creates a device code. +func validateDeviceCodeHost(authArguments *auth.AuthArguments) error { + if isAccountOrUnifiedHost(authArguments.Host, authArguments.DiscoveryURL) { + return errDeviceCodeWorkspaceOnly + } + return nil +} + +// validateDeviceCodeLogin returns an error if device code login is not enabled +// or is combined with flags it cannot honor. +func validateDeviceCodeLogin(ctx context.Context, cmd *cobra.Command) error { + if enabled, _ := env.GetBool(ctx, deviceCodeEnvVar); !enabled { + return fmt.Errorf("--use-device-code is experimental. Set %s=true to enable it", deviceCodeEnvVar) + } + for _, name := range deviceCodeIncompatibleFlags { + // Compare with the default so that, e.g., --skip-workspace=false is + // accepted. + if f := cmd.Flag(name); f.Changed && f.Value.String() != f.DefValue { + return fmt.Errorf("--use-device-code cannot be combined with --%s", name) + } + } + return nil +} + +// promptForConfirmationCode shows the user where to approve the login and +// reads the confirmation code from stdin. If the server rejected the previous +// code, it shows the server's error and the number of attempts left instead +// and asks again. +func promptForConfirmationCode(ctx context.Context, da u2m.DeviceAuthorization, lastErr error, attemptsLeft int) (string, error) { + if lastErr == nil { + cmdio.LogString(ctx, strings.Join([]string{ + "Using a browser on another device, visit:", + da.VerificationURI, + "", + "And enter the code:", + formatCode(da.UserCode), + "", + "Once finished, enter the confirmation code provided in your browser.", + }, "\n")) + } else { + attempts := "attempts" + if attemptsLeft == 1 { + attempts = "attempt" + } + cmdio.LogString(ctx, fmt.Sprintf("%v. Enter the confirmation code again (%d %s left).", lastErr, attemptsLeft, attempts)) + } + + // A blank line is sent like any other code, so the server rejects it and + // it counts as one of the attempts. A cancelled read fails the login, so + // stdin is not read again after it. + code, err := cmdio.AskCancellable(ctx, "Confirmation code", "") + if err != nil { + return "", fmt.Errorf("reading confirmation code: %w", err) + } + return code, nil +} + +// formatCode displays an 8-character device code as two groups of four +// separated by a dash, matching how the browser displays codes. +func formatCode(code string) string { + if len(code) != 8 { + return code + } + return code[:4] + "-" + code[4:] +} diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go new file mode 100644 index 00000000000..84fdc3f0bbd --- /dev/null +++ b/cmd/auth/login_device_code_test.go @@ -0,0 +1,206 @@ +package auth + +import ( + "bytes" + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateDeviceCodeLogin(t *testing.T) { + tests := []struct { + name string + envVal string + setFlag string + flagVal string + wantErr string + }{ + { + name: "disabled by default", + wantErr: "--use-device-code is experimental. Set DATABRICKS_CLI_DEVICE_CODE=true to enable it", + }, + { + name: "explicitly disabled", + envVal: "false", + wantErr: "--use-device-code is experimental. Set DATABRICKS_CLI_DEVICE_CODE=true to enable it", + }, + { + name: "enabled", + envVal: "true", + }, + { + name: "resource is incompatible", + envVal: "1", + setFlag: "resource", + flagVal: "https://workspace.test/ai-gateway/mcp-services/system.ai.github", + wantErr: "--use-device-code cannot be combined with --resource", + }, + { + name: "skip-workspace is incompatible", + envVal: "1", + setFlag: "skip-workspace", + flagVal: "true", + wantErr: "--use-device-code cannot be combined with --skip-workspace", + }, + { + name: "skip-workspace set to its default", + envVal: "1", + setFlag: "skip-workspace", + flagVal: "false", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := env.Set(t.Context(), deviceCodeEnvVar, tt.envVal) + cmd := &cobra.Command{} + cmd.Flags().StringArray("resource", nil, "") + cmd.Flags().Bool("skip-workspace", false, "") + if tt.setFlag != "" { + require.NoError(t, cmd.Flags().Set(tt.setFlag, tt.flagVal)) + } + + err := validateDeviceCodeLogin(ctx, cmd) + if tt.wantErr != "" { + assert.EqualError(t, err, tt.wantErr) + } else { + assert.NoError(t, err) + } + }) + } +} + +func TestPromptForConfirmationCode(t *testing.T) { + tests := []struct { + name string + input string + want string + wantErr string + }{ + {name: "code", input: "FTFA-QQKY\n", want: "FTFA-QQKY"}, + {name: "code without newline", input: "FTFA-QQKY", want: "FTFA-QQKY"}, + {name: "blank line", input: " \nFTFA-QQKY\n", want: " "}, + {name: "no input", input: "", wantErr: "reading confirmation code: EOF"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var stderr bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, strings.NewReader(tt.input), &bytes.Buffer{}, &stderr, "", "")) + + code, err := promptForConfirmationCode(ctx, u2m.DeviceAuthorization{ + UserCode: "WDJBMJHT", + VerificationURI: "https://workspace.test/device/verification", + }, nil, 3) + if tt.wantErr != "" { + assert.EqualError(t, err, tt.wantErr) + } else { + require.NoError(t, err) + assert.Equal(t, tt.want, code) + } + assert.True(t, strings.HasPrefix(stderr.String(), strings.Join([]string{ + "Using a browser on another device, visit:", + "https://workspace.test/device/verification", + "", + "And enter the code:", + "WDJB-MJHT", + "", + "Once finished, enter the confirmation code provided in your browser.", + "Confirmation code: ", + }, "\n")), stderr.String()) + }) + } +} + +func TestPromptForConfirmationCode_Retry(t *testing.T) { + tests := []struct { + attemptsLeft int + want string + }{ + {attemptsLeft: 2, want: "wrong code. Enter the confirmation code again (2 attempts left).\nConfirmation code: "}, + {attemptsLeft: 1, want: "wrong code. Enter the confirmation code again (1 attempt left).\nConfirmation code: "}, + } + for _, tt := range tests { + var stderr bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, strings.NewReader("FTFA-QQKY\n"), &bytes.Buffer{}, &stderr, "", "")) + + code, err := promptForConfirmationCode(ctx, u2m.DeviceAuthorization{ + UserCode: "WDJBMJHT", + VerificationURI: "https://workspace.test/device/verification", + }, errors.New("wrong code"), tt.attemptsLeft) + require.NoError(t, err) + assert.Equal(t, "FTFA-QQKY", code) + // The instructions are only shown before the first attempt. + assert.Equal(t, tt.want, stderr.String()) + } +} + +func TestPromptForConfirmationCode_Timeout(t *testing.T) { + // The pipe is never written to, so reading the code blocks until the + // context is done. + stdin, _ := io.Pipe() + defer stdin.Close() + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, stdin, &bytes.Buffer{}, &bytes.Buffer{}, "", "")) + ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer cancel() + + _, err := promptForConfirmationCode(ctx, u2m.DeviceAuthorization{ + UserCode: "WDJBMJHT", + VerificationURI: "https://workspace.test/device/verification", + }, nil, 3) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.EqualError(t, err, "reading confirmation code: context deadline exceeded") +} + +func TestFormatCode(t *testing.T) { + assert.Equal(t, "WDJB-MJHT", formatCode("WDJBMJHT")) + assert.Equal(t, "ABC", formatCode("ABC")) +} + +func TestValidateDeviceCodeHost(t *testing.T) { + workspace := newDiscoveryServer(t, map[string]any{ + "account_id": "acct-123", + "workspace_id": "12345", + "oidc_endpoint": "https://workspace.test/oidc", + }) + unified := newDiscoveryServer(t, map[string]any{ + "account_id": "acct-123", + "oidc_endpoint": "https://spog.test/oidc/accounts/acct-123", + }) + // Without an account ID, host discovery cannot fill in the placeholder + // and does not resolve the account-scoped discovery URL. + unifiedWithoutAccountID := newDiscoveryServer(t, map[string]any{ + "oidc_endpoint": "https://spog.test/oidc/accounts/{account_id}", + }) + + tests := []struct { + name string + args auth.AuthArguments + wantErr error + }{ + {name: "workspace host", args: auth.AuthArguments{Host: workspace.URL}}, + {name: "workspace host with query params", args: auth.AuthArguments{Host: workspace.URL + "?o=12345"}}, + {name: "classic account host", args: auth.AuthArguments{Host: "https://accounts.cloud.databricks.test"}, wantErr: errDeviceCodeWorkspaceOnly}, + {name: "unified host", args: auth.AuthArguments{Host: unified.URL}, wantErr: errDeviceCodeWorkspaceOnly}, + {name: "unified host with query params", args: auth.AuthArguments{Host: unified.URL + "?o=12345"}, wantErr: errDeviceCodeWorkspaceOnly}, + {name: "unified host with account ID in query params", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL + "?a=acct-123"}, wantErr: errDeviceCodeWorkspaceOnly}, + {name: "unified host with account ID flag", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL, AccountID: "acct-123"}, wantErr: errDeviceCodeWorkspaceOnly}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + args := tt.args + require.NoError(t, setHost(t.Context(), nil, &args, nil)) + assert.ErrorIs(t, validateDeviceCodeHost(&args), tt.wantErr) + }) + } +} diff --git a/docs/device-code-login.md b/docs/device-code-login.md new file mode 100644 index 00000000000..87c96c03c16 --- /dev/null +++ b/docs/device-code-login.md @@ -0,0 +1,68 @@ +# Device code login (experimental) + +Device code login authenticates the Databricks CLI on machines that cannot open +a browser, such as headless servers, SSH sessions, and web terminals. You +approve the login in a browser on another device (for example your laptop or +phone), and type the confirmation code it shows you back into the CLI. + +> [!WARNING] +> Device code login is experimental. It is disabled by default, may change +> without notice, and only supports workspace hosts. +> Account hosts and unified (SPOG) hosts are not supported yet. + +## Prerequisites + +- The workspace must have device code login enabled. If it is not, the CLI + fails with `device code login is not enabled for this workspace`. +- Enable the feature in the CLI by setting `DATABRICKS_CLI_DEVICE_CODE=true`. + +## Log in from a headless machine + +```console +$ export DATABRICKS_CLI_DEVICE_CODE=true +$ databricks auth login --host https:// --profile my-profile --use-device-code +Using a browser on another device, visit: +https:///device/verification + +And enter the code: +WDJB-MJHT + +Once finished, enter the confirmation code provided in your browser. +Confirmation code: FTFA-QQKY +Profile my-profile was successfully saved +``` + +1. Run `databricks auth login` with `--use-device-code`. The host can be passed as + an argument, with `--host`, or come from an existing profile. +2. On another device, open the URL that the CLI shows, sign in to the + workspace if needed, and enter the code. The code expires after about 10 + minutes. +3. Confirm that you started this login yourself. Only approve a login that you + initiated: approving it grants the CLI access to the Databricks API as you. +4. The browser shows a confirmation code. Type it into the CLI. Codes are not + case-sensitive, and the dash is optional. If the code is rejected, the CLI + asks for it again and shows how many attempts are left, up to 3 attempts in + total. Never share the confirmation code with anyone. + +The CLI saves the profile and the token the same way as the browser-based +`databricks auth login`, so other commands use it via `--profile my-profile`. + +## Session lifetime + +The token is renewed automatically for up to 24 hours. After that, commands +fail with a token refresh error and you need to run the login again. + +## Troubleshooting + +| Error or symptom | Likely cause | What to do | +| --- | --- | --- | +| `--use-device-code is experimental. Set DATABRICKS_CLI_DEVICE_CODE=true to enable it` | The feature is not enabled in the CLI. | Set `DATABRICKS_CLI_DEVICE_CODE=true`. | +| `device code login is not enabled for this workspace` | The workspace does not advertise device code login. | Use a workspace where the feature is enabled. | +| `--use-device-code is only supported for workspace hosts` | The host is an account or unified host. | Log in to a workspace host instead. | +| The browser says the code is invalid or expired | The code was mistyped, or more than ~10 minutes passed. | Re-enter the code carefully, or run the login again for a new code. Codes never contain `0`, `1`, `2`, `O`, `Z`, or `I`. | +| `Invalid confirmation code` | The confirmation code was mistyped or left blank. | Re-enter the code. After 3 rejected codes, the login fails with `Invalid confirmation code (error code: invalid_request). Start login from scratch` and you need to run it again. | +| `token exchange: Invalid device code` | The login expired. A confirmation code entered after the login expired is first rejected as `Invalid confirmation code`, and the workspace then discards the login. | Run the login again to get a new code. | +| `token exchange: Device grant rate limit exceeded` | Too many confirmation codes were sent for this login. | Run the login again to get a new code. | +| `unauthorized_client` or `invalid_client` | The OAuth client is not allowed to use device code login in this workspace. | Use the default `databricks-cli` client, or ask your admin to enable the custom client. | +| `reading confirmation code: EOF` | stdin was closed before a code was entered. | Keep stdin open until you write the confirmation code. | +| `reading confirmation code: context deadline exceeded` | No confirmation code was entered before the login timed out (1 hour by default). | Run the login again. Use `--timeout` to change the limit. | diff --git a/libs/auth/u2m/device_code.go b/libs/auth/u2m/device_code.go new file mode 100644 index 00000000000..51f77f8e189 --- /dev/null +++ b/libs/auth/u2m/device_code.go @@ -0,0 +1,244 @@ +package u2m + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "golang.org/x/oauth2" +) + +// confirmationCodeParam is the Databricks extension to the RFC 8628 token +// request. The user receives the confirmation code in the browser after +// approving the login, and the client submits it together with the device +// code instead of polling the token endpoint until the login is approved. +const confirmationCodeParam = "confirmation_code" + +// deviceCodeGrantType is the grant type of the RFC 8628 token request. +const deviceCodeGrantType = "urn:ietf:params:oauth:grant-type:device_code" + +// invalidDeviceCodeDescription is the error_description the server sends with +// HTTP 400 and invalid_request for a device code it does not know, e.g. one +// it discarded after the login expired. +const invalidDeviceCodeDescription = "Invalid device code" + +// ErrDeviceCodeNotSupported is returned when the workspace does not advertise +// a device authorization endpoint in its OIDC discovery document. +var ErrDeviceCodeNotSupported = errors.New("device code login is not enabled for this workspace") + +// DeviceAuthorization holds what the user needs to approve a device code +// login from another device. +type DeviceAuthorization struct { + // UserCode is the code the user enters at VerificationURI. + UserCode string + + // VerificationURI is the page where the user approves the login. + VerificationURI string +} + +// maxConfirmationAttempts is how many confirmation codes the user can enter +// for a single device code before the login fails. +const maxConfirmationAttempts = 3 + +// ConfirmFunc shows the device authorization to the user and returns the +// confirmation code that the user received after approving the login. When +// the server rejects a confirmation code, it is called again with the +// server's error in lastErr so that the user can re-enter the code. lastErr +// is nil on the first call. attemptsLeft is the number of confirmation codes +// the user can still enter, including this one. +type ConfirmFunc func(ctx context.Context, da DeviceAuthorization, lastErr error, attemptsLeft int) (string, error) + +// deviceCodeConfig returns the OAuth config for device code login, or an +// error if the host does not support it. +func (a *PersistentAuth) deviceCodeConfig() (*oauth2.Config, error) { + if _, ok := a.oAuthArgument.(WorkspaceOAuthArgument); !ok { + return nil, errors.New("device code login is only supported for workspace hosts") + } + // Resource indicators are only sent on the browser authorization request, + // so the token would not be scoped to them. + if len(a.resources) > 0 { + return nil, errors.New("device code login does not support resource indicators") + } + cfg, err := a.oauth2Config() + if err != nil { + return nil, fmt.Errorf("fetching oauth config: %w", err) + } + if cfg.Endpoint.DeviceAuthURL == "" { + return nil, ErrDeviceCodeNotSupported + } + return cfg, nil +} + +// DeviceCodeChallenge logs in using the Databricks variant of the RFC 8628 +// device authorization grant. It requests a device authorization, calls +// confirm with the user code and verification URI, and exchanges the device +// code together with the returned confirmation code for a token. A rejected +// confirmation code leaves the device code valid, so the user can re-enter +// it up to maxConfirmationAttempts times. Only workspace hosts are supported. +// The caller is responsible for storing the returned token. +func (a *PersistentAuth) DeviceCodeChallenge(ctx context.Context, confirm ConfirmFunc) (*oauth2.Token, error) { + cfg, err := a.deviceCodeConfig() + if err != nil { + return nil, err + } + + da, err := a.requestDeviceAuthorization(ctx, cfg) + if err != nil { + return nil, fmt.Errorf("device authorization: %w", oauthError(err)) + } + + authorization := DeviceAuthorization{ + UserCode: da.UserCode, + VerificationURI: da.VerificationURI, + } + var lastErr error + for attempt := range maxConfirmationAttempts { + code, err := confirm(ctx, authorization, lastErr, maxConfirmationAttempts-attempt) + if err != nil { + return nil, err + } + + t, err := a.exchangeDeviceCode(ctx, cfg, da.DeviceCode, normalizeCode(code)) + if err == nil { + return t, nil + } + lastErr = oauthError(err) + // The server answers a wrong confirmation code with HTTP 400 and + // invalid_request. Any other error cannot be fixed by entering the + // code again, including invalid_request with HTTP 429 when the server + // rate limits the token requests for this device code. An unknown + // device code also gets HTTP 400 and invalid_request, and only the + // description tells it apart. If the description changes, the user is + // asked for the code again until the attempts run out. + status, errCode, description := oauthErrorResponse(err) + if status != http.StatusBadRequest || errCode != "invalid_request" || description == invalidDeviceCodeDescription { + return nil, fmt.Errorf("token exchange: %w", lastErr) + } + } + return nil, fmt.Errorf("%w. Start login from scratch", lastErr) +} + +// requestDeviceAuthorization makes the RFC 8628 device authorization request. +// Unlike oauth2.Config.DeviceAuth, it sends the request with ctx, so that the +// login timeout and cancellation stop a stalled request. +func (a *PersistentAuth) requestDeviceAuthorization(ctx context.Context, cfg *oauth2.Config) (*oauth2.DeviceAuthResponse, error) { + body, err := a.postForm(ctx, cfg.Endpoint.DeviceAuthURL, url.Values{ + "client_id": {cfg.ClientID}, + "scope": {strings.Join(cfg.Scopes, " ")}, + }) + if err != nil { + return nil, err + } + var da oauth2.DeviceAuthResponse + if err := json.Unmarshal(body, &da); err != nil { + return nil, fmt.Errorf("parsing device authorization response: %w", err) + } + // The user needs user_code and verification_uri to approve the login, and + // the token request needs device_code. + switch { + case da.DeviceCode == "": + return nil, errors.New("server response missing device_code") + case da.UserCode == "": + return nil, errors.New("server response missing user_code") + case da.VerificationURI == "": + return nil, errors.New("server response missing verification_uri") + } + return &da, nil +} + +// exchangeDeviceCode makes a single RFC 8628 device access token request. +// The server issues the token as soon as it receives a valid confirmation +// code, so unlike oauth2.Config.DeviceAccessToken this does not wait before +// the request or poll while the login is pending. It also does not enforce +// the expiry from the device authorization response, because the server +// extends the device code's lifetime when the user approves the login. +func (a *PersistentAuth) exchangeDeviceCode(ctx context.Context, cfg *oauth2.Config, deviceCode, confirmationCode string) (*oauth2.Token, error) { + // No scope: the token gets the scopes of the device authorization + // request, and the server ignores a scope sent here. + body, err := a.postForm(ctx, cfg.Endpoint.TokenURL, url.Values{ + "grant_type": {deviceCodeGrantType}, + "client_id": {cfg.ClientID}, + "device_code": {deviceCode}, + confirmationCodeParam: {confirmationCode}, + }) + if err != nil { + return nil, err + } + + var tokenResp struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + RefreshToken string `json:"refresh_token"` + ExpiresIn int64 `json:"expires_in"` + } + var raw map[string]any + if err := json.Unmarshal(body, &tokenResp); err != nil { + return nil, fmt.Errorf("parsing token response: %w", err) + } + if err := json.Unmarshal(body, &raw); err != nil { + return nil, fmt.Errorf("parsing token response: %w", err) + } + if tokenResp.AccessToken == "" { + return nil, errors.New("server response missing access_token") + } + t := &oauth2.Token{ + AccessToken: tokenResp.AccessToken, + TokenType: tokenResp.TokenType, + RefreshToken: tokenResp.RefreshToken, + ExpiresIn: tokenResp.ExpiresIn, + } + if tokenResp.ExpiresIn > 0 { + t.Expiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + } + return t.WithExtra(raw), nil +} + +// postForm sends a form-encoded OAuth request and returns the body of a +// successful response. An error response is returned as an +// *oauth2.RetrieveError with the RFC 6749 §5.2 error fields. +func (a *PersistentAuth) postForm(ctx context.Context, endpoint string, form url.Values) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode())) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + + resp, err := a.client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return nil, fmt.Errorf("reading response: %w", err) + } + + if resp.StatusCode < 200 || resp.StatusCode > 299 { + retrieveErr := &oauth2.RetrieveError{Response: resp, Body: body} + var errResp struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + if json.Unmarshal(body, &errResp) == nil { + retrieveErr.ErrorCode = errResp.Error + retrieveErr.ErrorDescription = errResp.ErrorDescription + } + return nil, retrieveErr + } + return body, nil +} + +// normalizeCode lets users enter a code the way it is displayed. Databricks +// codes consist of 8 uppercase alphanumeric characters and are compared +// case-sensitively, but are often shown as two groups of four separated by a +// dash (e.g. FTFA-QQKY). +func normalizeCode(code string) string { + return strings.NewReplacer("-", "", " ", "").Replace(strings.ToUpper(strings.TrimSpace(code))) +} diff --git a/libs/auth/u2m/device_code_test.go b/libs/auth/u2m/device_code_test.go new file mode 100644 index 00000000000..f976ef5555d --- /dev/null +++ b/libs/auth/u2m/device_code_test.go @@ -0,0 +1,417 @@ +package u2m + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +const ( + testDeviceCode = "device-code-123" + testUserCode = "WDJBMJHT" + testConfirmationCode = "FTFAQQKY" +) + +// deviceCodeEndpointSupplier serves workspace OIDC endpoints from a test +// server, optionally advertising the device authorization endpoint. +type deviceCodeEndpointSupplier struct { + MockOAuthEndpointSupplier + host string + deviceAuthAdvertise bool +} + +func (s deviceCodeEndpointSupplier) GetWorkspaceOAuthEndpoints(_ context.Context, _ string) (*OAuthAuthorizationServer, error) { + endpoints := &OAuthAuthorizationServer{ + AuthorizationEndpoint: s.host + "/oidc/v1/authorize", + TokenEndpoint: s.host + "/oidc/v1/token", + } + if s.deviceAuthAdvertise { + endpoints.DeviceAuthorizationEndpoint = s.host + "/oidc/v1/device/auth" + } + return endpoints, nil +} + +type fakeDeviceServer struct { + deviceAuthStatus int + deviceAuthBody string + // deviceAuthStalls makes the device authorization request hang until + // the client gives up on it. + deviceAuthStalls bool + tokenStatus int + tokenBody string + expiresIn int + + tokenRequests atomic.Int32 +} + +func (f *fakeDeviceServer) handler(t *testing.T) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("POST /oidc/v1/device/auth", func(w http.ResponseWriter, r *http.Request) { + assert.NoError(t, r.ParseForm()) + assert.Equal(t, "databricks-cli", r.PostForm.Get("client_id")) + assert.Equal(t, "offline_access all-apis", r.PostForm.Get("scope")) + if f.deviceAuthStalls { + <-r.Context().Done() + return + } + w.Header().Set("Content-Type", "application/json") + if f.deviceAuthStatus != 0 { + w.WriteHeader(f.deviceAuthStatus) + _, _ = w.Write([]byte(f.deviceAuthBody)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "device_code": testDeviceCode, + "user_code": testUserCode, + "verification_uri": "https://workspace.test/device/verification", + "expires_in": f.expiresIn, + "interval": 0, + }) + }) + mux.HandleFunc("POST /oidc/v1/token", func(w http.ResponseWriter, r *http.Request) { + f.tokenRequests.Add(1) + assert.NoError(t, r.ParseForm()) + assert.Equal(t, "urn:ietf:params:oauth:grant-type:device_code", r.PostForm.Get("grant_type")) + assert.Equal(t, "databricks-cli", r.PostForm.Get("client_id")) + assert.Equal(t, testDeviceCode, r.PostForm.Get("device_code")) + w.Header().Set("Content-Type", "application/json") + if f.tokenStatus != 0 { + w.WriteHeader(f.tokenStatus) + _, _ = w.Write([]byte(f.tokenBody)) + return + } + if r.PostForm.Get("confirmation_code") != testConfirmationCode { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error": "invalid_request", "error_description": "Invalid confirmation code"}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-token", + "refresh_token": "refresh-token", + "token_type": "Bearer", + "expires_in": 3600, + }) + }) + return mux +} + +func newDeviceCodeAuth(t *testing.T, f *fakeDeviceServer, advertise bool, opts ...PersistentAuthOption) *PersistentAuth { + srv := httptest.NewServer(f.handler(t)) + t.Cleanup(srv.Close) + arg, err := NewBasicWorkspaceOAuthArgument("https://workspace.test") + require.NoError(t, err) + p, err := NewPersistentAuth(t.Context(), append([]PersistentAuthOption{ + WithOAuthArgument(arg), + WithOAuthEndpointSupplier(deviceCodeEndpointSupplier{host: srv.URL, deviceAuthAdvertise: advertise}), + }, opts...)...) + require.NoError(t, err) + t.Cleanup(func() { p.Close() }) + return p +} + +func TestDeviceCodeChallenge(t *testing.T) { + f := &fakeDeviceServer{expiresIn: 600} + p := newDeviceCodeAuth(t, f, true) + + var got DeviceAuthorization + tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, da DeviceAuthorization, _ error, _ int) (string, error) { + got = da + return " ftfa-qqky\t", nil + }) + require.NoError(t, err) + + assert.Equal(t, DeviceAuthorization{ + UserCode: testUserCode, + VerificationURI: "https://workspace.test/device/verification", + }, got) + assert.Equal(t, "access-token", tok.AccessToken) + assert.Equal(t, "refresh-token", tok.RefreshToken) + assert.WithinDuration(t, time.Now().Add(time.Hour), tok.Expiry, time.Minute) + assert.Equal(t, int32(1), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_ConfirmationAfterInitialExpiry(t *testing.T) { + // The server extends the device code's lifetime once the user approves + // the login, so a confirmation code entered after the initial expires_in + // must still be exchanged. + f := &fakeDeviceServer{expiresIn: 1} + p := newDeviceCodeAuth(t, f, true) + + tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + time.Sleep(1100 * time.Millisecond) + return testConfirmationCode, nil + }) + require.NoError(t, err) + assert.Equal(t, "access-token", tok.AccessToken) +} + +func TestDeviceCodeChallenge_Errors(t *testing.T) { + tests := []struct { + name string + server *fakeDeviceServer + wantErr string + }{ + { + name: "device authorization rejected", + server: &fakeDeviceServer{ + deviceAuthStatus: http.StatusBadRequest, + deviceAuthBody: `{"error": "unauthorized_client", "error_description": "Client is not allowed to use device code"}`, + }, + wantErr: "device authorization: Client is not allowed to use device code (error code: unauthorized_client)", + }, + { + name: "token request rejected", + server: &fakeDeviceServer{ + tokenStatus: http.StatusBadRequest, + tokenBody: `{"error": "invalid_client", "error_description": "Client is not allowed to use device code"}`, + }, + wantErr: "token exchange: Client is not allowed to use device code (error code: invalid_client)", + }, + { + name: "token request rate limited", + server: &fakeDeviceServer{ + tokenStatus: http.StatusTooManyRequests, + tokenBody: `{"error": "invalid_request", "error_description": "Device grant rate limit exceeded"}`, + }, + wantErr: "token exchange: Device grant rate limit exceeded (error code: invalid_request)", + }, + { + name: "unknown device code", + server: &fakeDeviceServer{ + tokenStatus: http.StatusBadRequest, + tokenBody: `{"error": "invalid_request", "error_description": "Invalid device code"}`, + }, + wantErr: "token exchange: Invalid device code (error code: invalid_request)", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tc.server.expiresIn = 600 + p := newDeviceCodeAuth(t, tc.server, true) + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + return testConfirmationCode, nil + }) + assert.EqualError(t, err, tc.wantErr) + // None of these errors asks for the confirmation code again. + assert.LessOrEqual(t, tc.server.tokenRequests.Load(), int32(1)) + }) + } +} + +func TestDeviceCodeChallenge_IncompleteDeviceAuthorization(t *testing.T) { + for _, missing := range []string{"device_code", "user_code", "verification_uri"} { + t.Run(missing, func(t *testing.T) { + resp := map[string]any{ + "device_code": testDeviceCode, + "user_code": testUserCode, + "verification_uri": "https://workspace.test/device/verification", + "expires_in": 600, + } + delete(resp, missing) + body, err := json.Marshal(resp) + require.NoError(t, err) + f := &fakeDeviceServer{deviceAuthStatus: http.StatusOK, deviceAuthBody: string(body)} + p := newDeviceCodeAuth(t, f, true) + + _, err = p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + t.Fatal("confirm must not be called") + return "", nil + }) + assert.EqualError(t, err, "device authorization: server response missing "+missing) + assert.Equal(t, int32(0), f.tokenRequests.Load()) + }) + } +} + +func TestDeviceCodeChallenge_ErrorKeepsHTTPError(t *testing.T) { + f := &fakeDeviceServer{ + expiresIn: 600, + tokenStatus: http.StatusBadRequest, + tokenBody: `{"error": "invalid_client", "error_description": "Client is not allowed to use device code"}`, + } + confirm := func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + return testConfirmationCode, nil + } + + _, err := newDeviceCodeAuth(t, f, true).DeviceCodeChallenge(t.Context(), confirm) + httpErr, ok := errors.AsType[*httpclient.HttpError](err) + require.True(t, ok, "want *httpclient.HttpError in %v", err) + assert.Equal(t, http.StatusBadRequest, httpErr.StatusCode) + + // A plain HTTP client hands the error response to the CLI as is. + _, err = newDeviceCodeAuth(t, f, true, WithHttpClient(&http.Client{})).DeviceCodeChallenge(t.Context(), confirm) + retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err) + require.True(t, ok, "want *oauth2.RetrieveError in %v", err) + assert.Equal(t, http.StatusBadRequest, retrieveErr.Response.StatusCode) + assert.EqualError(t, err, "token exchange: Client is not allowed to use device code (error code: invalid_client)") +} + +func TestDeviceCodeChallenge_RetriesWrongConfirmationCode(t *testing.T) { + f := &fakeDeviceServer{expiresIn: 600} + p := newDeviceCodeAuth(t, f, true) + + var lastErrs []string + tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, lastErr error, _ int) (string, error) { + if lastErr == nil { + lastErrs = append(lastErrs, "") + return "AAAAAAAA", nil + } + lastErrs = append(lastErrs, lastErr.Error()) + return testConfirmationCode, nil + }) + require.NoError(t, err) + assert.Equal(t, "access-token", tok.AccessToken) + assert.Equal(t, []string{"", "Invalid confirmation code (error code: invalid_request)"}, lastErrs) + assert.Equal(t, int32(2), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_TooManyWrongConfirmationCodes(t *testing.T) { + f := &fakeDeviceServer{expiresIn: 600} + p := newDeviceCodeAuth(t, f, true) + + var attemptsLeft []int + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, left int) (string, error) { + attemptsLeft = append(attemptsLeft, left) + return "AAAAAAAA", nil + }) + assert.EqualError(t, err, "Invalid confirmation code (error code: invalid_request). Start login from scratch") + assert.Equal(t, []int{3, 2, 1}, attemptsLeft) + assert.Equal(t, int32(maxConfirmationAttempts), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_ConfirmErrorOnRetry(t *testing.T) { + f := &fakeDeviceServer{expiresIn: 600} + p := newDeviceCodeAuth(t, f, true) + confirmErr := errors.New("no input") + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, lastErr error, _ int) (string, error) { + if lastErr != nil { + return "", confirmErr + } + return "AAAAAAAA", nil + }) + assert.ErrorIs(t, err, confirmErr) + assert.Equal(t, int32(1), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_DoesNotPoll(t *testing.T) { + // A plain HTTP client hands the error response to the CLI as is, rather + // than turning it into an HttpError. A pending login must fail right away + // instead of being polled until it is approved. + f := &fakeDeviceServer{ + expiresIn: 600, + tokenStatus: http.StatusBadRequest, + tokenBody: `{"error": "authorization_pending", "error_description": "Login is not approved yet"}`, + } + p := newDeviceCodeAuth(t, f, true, WithHttpClient(&http.Client{})) + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + return testConfirmationCode, nil + }) + assert.EqualError(t, err, "token exchange: Login is not approved yet (error code: authorization_pending)") + assert.Equal(t, int32(1), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_MissingAccessToken(t *testing.T) { + f := &fakeDeviceServer{ + expiresIn: 600, + tokenStatus: http.StatusOK, + tokenBody: `{"token_type": "Bearer"}`, + } + p := newDeviceCodeAuth(t, f, true) + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + return testConfirmationCode, nil + }) + assert.EqualError(t, err, "token exchange: server response missing access_token") +} + +func TestDeviceCodeChallenge_ConfirmError(t *testing.T) { + f := &fakeDeviceServer{expiresIn: 600} + p := newDeviceCodeAuth(t, f, true) + confirmErr := errors.New("no input") + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + return "", confirmErr + }) + assert.ErrorIs(t, err, confirmErr) + assert.Equal(t, int32(0), f.tokenRequests.Load()) +} + +func TestDeviceCodeChallenge_DeviceAuthorizationCancelled(t *testing.T) { + p := newDeviceCodeAuth(t, &fakeDeviceServer{deviceAuthStalls: true}, true) + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := p.DeviceCodeChallenge(ctx, func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + t.Fatal("confirm must not be called") + return "", nil + }) + assert.ErrorIs(t, err, context.DeadlineExceeded) + // The HTTP client's own 30 second timeout also ends the request with + // DeadlineExceeded, so check that the context ended it. + assert.Less(t, time.Since(start), 10*time.Second) +} + +func TestDeviceCodeChallenge_NotAdvertised(t *testing.T) { + p := newDeviceCodeAuth(t, &fakeDeviceServer{}, false) + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + t.Fatal("confirm must not be called") + return "", nil + }) + assert.ErrorIs(t, err, ErrDeviceCodeNotSupported) +} + +func TestDeviceCodeChallenge_Resources(t *testing.T) { + p := newDeviceCodeAuth(t, &fakeDeviceServer{}, true, WithResources([]string{"https://resource.test"})) + + _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + t.Fatal("confirm must not be called") + return "", nil + }) + assert.EqualError(t, err, "device code login does not support resource indicators") +} + +func TestDeviceCodeChallenge_AccountHost(t *testing.T) { + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + require.NoError(t, err) + p, err := NewPersistentAuth(t.Context(), + WithOAuthArgument(arg), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + ) + require.NoError(t, err) + defer p.Close() + + _, err = p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { + t.Fatal("confirm must not be called") + return "", nil + }) + assert.EqualError(t, err, "device code login is only supported for workspace hosts") +} + +func TestNormalizeCode(t *testing.T) { + tests := []struct { + in string + want string + }{ + {"FTFAQQKY", "FTFAQQKY"}, + {"FTFA-QQKY", "FTFAQQKY"}, + {"ftfa-qqky", "FTFAQQKY"}, + {" FTFA QQKY \n", "FTFAQQKY"}, + } + for _, tc := range tests { + assert.Equal(t, tc.want, normalizeCode(tc.in), tc.in) + } +} diff --git a/libs/auth/u2m/endpoint_supplier.go b/libs/auth/u2m/endpoint_supplier.go index f5065aa1c61..987f75ead24 100644 --- a/libs/auth/u2m/endpoint_supplier.go +++ b/libs/auth/u2m/endpoint_supplier.go @@ -78,4 +78,8 @@ type OAuthAuthorizationServer struct { // TokenEndpoint is the URL to exchange an authorization code for an access token. // It typically ends with /v1/token. TokenEndpoint string `json:"token_endpoint"` + + // DeviceAuthorizationEndpoint is the RFC 8628 device authorization endpoint. + // Workspaces only advertise it when device code login is enabled for them. + DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint"` } diff --git a/libs/auth/u2m/error.go b/libs/auth/u2m/error.go index 51ba9036551..d8469409fb3 100644 --- a/libs/auth/u2m/error.go +++ b/libs/auth/u2m/error.go @@ -1,6 +1,13 @@ package u2m -import "errors" +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/databricks/databricks-sdk-go/httpclient" + "golang.org/x/oauth2" +) // ErrMissingRefreshToken is returned when a token refresh is requested but the // cached OAuth token does not include a refresh token. @@ -12,3 +19,52 @@ var ErrMissingRefreshToken = errors.New("cached token has no refresh token") type InvalidRefreshTokenError struct { error } + +// oauthError extracts the RFC 6749 §5.2 error response from err so that users +// see the server's explanation instead of a raw HTTP response body. +func oauthError(err error) error { + _, code, description := oauthErrorResponse(err) + if code == "" { + return err + } + return &oauthResponseError{code: code, description: description, err: err} +} + +// oauthResponseError shows an RFC 6749 §5.2 error response and keeps the +// error that carried it, so that callers can still inspect the HTTP response. +type oauthResponseError struct { + code string + description string + err error +} + +func (e *oauthResponseError) Error() string { + return fmt.Sprintf("%s (error code: %s)", e.description, e.code) +} + +func (e *oauthResponseError) Unwrap() error { + return e.err +} + +// oauthErrorResponse returns the HTTP status code, and the error code and +// description of the RFC 6749 §5.2 error response carried by err. The status +// is 0 if err does not carry an HTTP response, and the code is empty if err +// does not carry an error response. +func oauthErrorResponse(err error) (status int, code, description string) { + // The default RoundTripper of httpclient.ApiClient returns an error for + // non-2xx responses, so the oauth2 library never parses the body itself. + if httpErr, ok := errors.AsType[*httpclient.HttpError](err); ok { + var resp struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + if json.Unmarshal([]byte(httpErr.Message), &resp) != nil { + return httpErr.StatusCode, "", "" + } + return httpErr.StatusCode, resp.Error, resp.ErrorDescription + } + if retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok { + return retrieveErr.Response.StatusCode, retrieveErr.ErrorCode, retrieveErr.ErrorDescription + } + return 0, "", "" +} diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index a87e37783f4..de0df97d47e 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -5,7 +5,6 @@ import ( "crypto/rand" "crypto/sha256" "encoding/base64" - "encoding/json" "errors" "fmt" "net" @@ -426,38 +425,12 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) expired.Expiry = time.Now().Add(-time.Minute) t, err := cfg.TokenSource(ctx, &expired).Token() if err != nil { - // The default RoundTripper of our httpclient.ApiClient returns an error - // if the response status code is not 2xx. This isn't compliant with the - // RoundTripper interface, so this error isn't handled by the oauth2 - // library. We need to handle it here. - if internalHttpError, ok := errors.AsType[*httpclient.HttpError](err); ok { - // error fields - // https://datatracker.ietf.org/doc/html/rfc6749#section-5.2 - var errResponse struct { - Error string `json:"error"` - ErrorDescription string `json:"error_description"` - } - if unmarshalErr := json.Unmarshal([]byte(internalHttpError.Message), &errResponse); unmarshalErr != nil { - return nil, fmt.Errorf("unmarshal: %w", unmarshalErr) - } - // Invalid refresh tokens get their own error type so they can be - // better presented to users. - if errResponse.ErrorDescription == "Refresh token is invalid" { - return nil, &InvalidRefreshTokenError{err} - } - return nil, fmt.Errorf("%s (error code: %s)", errResponse.ErrorDescription, errResponse.Error) + // Invalid refresh tokens get their own error type so they can be + // better presented to users. + if _, _, description := oauthErrorResponse(err); description == "Refresh token is invalid" { + return nil, &InvalidRefreshTokenError{err} } - - // Handle responses from well-behaved *http.Client implementations. - if httpErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok { - // Invalid refresh tokens get their own error type so they can be - // better presented to users. - if httpErr.ErrorDescription == "Refresh token is invalid" { - return nil, &InvalidRefreshTokenError{err} - } - return nil, fmt.Errorf("%s (error code: %s)", httpErr.ErrorDescription, httpErr.ErrorCode) - } - return nil, err + return nil, oauthError(err) } err = a.store.Put(a.oAuthArgument.GetCacheKey(), storage.Entry{Token: t}) if err != nil { @@ -647,9 +620,10 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) { return &oauth2.Config{ ClientID: a.clientID, Endpoint: oauth2.Endpoint{ - AuthURL: authURL, - TokenURL: endpoints.TokenEndpoint, - AuthStyle: oauth2.AuthStyleInParams, + AuthURL: authURL, + TokenURL: endpoints.TokenEndpoint, + DeviceAuthURL: endpoints.DeviceAuthorizationEndpoint, + AuthStyle: oauth2.AuthStyleInParams, }, RedirectURL: "http://" + a.redirectAddr, Scopes: scopes, diff --git a/libs/auth/u2m/persistent_auth_test.go b/libs/auth/u2m/persistent_auth_test.go index 831fac8402c..1817ff778ae 100644 --- a/libs/auth/u2m/persistent_auth_test.go +++ b/libs/auth/u2m/persistent_auth_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "golang.org/x/oauth2" ) @@ -572,6 +573,39 @@ func TestToken_ReturnsInvalidRefreshTokenError(t *testing.T) { } } +func TestToken_ReturnsErrorWithoutErrorCode(t *testing.T) { + // The default ApiClient transport returns an error response as an + // *httpclient.HttpError, which is returned as is without an error code. + for _, body := range []string{"Bad Gateway", `{"error_description": "Bad Gateway"}`} { + t.Run(body, func(t *testing.T) { + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenStore(&tokenStoreMock{lookup: func(string) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "expired", RefreshToken: "cde", Expiry: time.Now().Add(-time.Minute)}, nil + }}), + WithHttpClient(&http.Client{Transport: httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.SliceTransport{{Method: "POST", Resource: "/oidc/accounts/xyz/v1/token", Response: body, Status: 502}}, + })}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + _, err = p.Token() + want := `token refresh: Post "https://accounts.cloud.databricks.test/oidc/accounts/xyz/v1/token": http 502: ` + body + if err == nil || err.Error() != want { + t.Errorf("p.Token(): want error %q, got %v", want, err) + } + }) + } +} + func TestToken_ReturnsMissingRefreshTokenErrorWhenExpired(t *testing.T) { c := &tokenStoreMock{ lookup: func(key string) (*oauth2.Token, error) { @@ -1007,6 +1041,9 @@ func TestForceRefreshToken_FailsOnRefreshError(t *testing.T) { if got := err.Error(); got != want { t.Errorf("ForceRefreshToken(): want error %q, got %q", want, got) } + if _, ok := errors.AsType[*oauth2.RetrieveError](err); !ok { + t.Errorf("ForceRefreshToken(): want error wrapping *oauth2.RetrieveError, got %v", err) + } } func TestForceRefreshToken_InvalidRefreshTokenError(t *testing.T) { diff --git a/libs/cmdio/ask.go b/libs/cmdio/ask.go index 909076d5e60..82772c6d38d 100644 --- a/libs/cmdio/ask.go +++ b/libs/cmdio/ask.go @@ -65,6 +65,30 @@ func Ask(ctx context.Context, question, defaultVal string) (string, error) { return ans, nil } +// AskCancellable is like Ask, but stops waiting for the answer when ctx is +// done and returns ctx.Err(). Reading the input cannot be interrupted, so the +// read keeps running in the background and its result is discarded. The +// caller must not read the input again after a cancelled call, because the +// abandoned read would consume the next line. +func AskCancellable(ctx context.Context, question, defaultVal string) (string, error) { + type answer struct { + ans string + err error + } + answers := make(chan answer, 1) + go func() { + ans, err := Ask(ctx, question, defaultVal) + answers <- answer{ans, err} + }() + + select { + case <-ctx.Done(): + return "", ctx.Err() + case a := <-answers: + return a.ans, a.err + } +} + // AskYesOrNo prompts the user with a question and returns true if the answer // is "y" or "yes" (case-insensitive). Any other answer, including an empty // one, returns false. diff --git a/libs/cmdio/ask_test.go b/libs/cmdio/ask_test.go index 75fc520f8d5..75fc64d7c54 100644 --- a/libs/cmdio/ask_test.go +++ b/libs/cmdio/ask_test.go @@ -1,6 +1,7 @@ package cmdio import ( + "context" "errors" "io" "strings" @@ -177,3 +178,29 @@ func TestAskYesOrNo(t *testing.T) { }) } } + +func TestAskCancellable(t *testing.T) { + ctx := t.Context() + ctx = InContext(ctx, NewIO(ctx, flags.OutputText, strings.NewReader("answer\n"), io.Discard, io.Discard, "", "")) + got, err := AskCancellable(ctx, "Question", "") + require.NoError(t, err) + assert.Equal(t, "answer", got) +} + +func TestAskCancellable_ReadError(t *testing.T) { + ctx := t.Context() + ctx = InContext(ctx, NewIO(ctx, flags.OutputText, strings.NewReader(""), io.Discard, io.Discard, "", "")) + _, err := AskCancellable(ctx, "Question", "") + assert.ErrorIs(t, err, io.EOF) +} + +func TestAskCancellable_Cancelled(t *testing.T) { + // The pipe is never written to, so the read blocks until ctx is done. + stdin, _ := io.Pipe() + defer stdin.Close() + ctx, cancel := context.WithCancel(t.Context()) + ctx = InContext(ctx, NewIO(ctx, flags.OutputText, stdin, io.Discard, io.Discard, "", "")) + cancel() + _, err := AskCancellable(ctx, "Question", "") + assert.ErrorIs(t, err, context.Canceled) +} diff --git a/libs/testserver/fake_oidc.go b/libs/testserver/fake_oidc.go index db6a682c9bd..492b1967056 100644 --- a/libs/testserver/fake_oidc.go +++ b/libs/testserver/fake_oidc.go @@ -3,18 +3,112 @@ package testserver import ( "net/http" "net/url" + "slices" + "strconv" + "strings" + "sync" +) + +const ( + deviceCodeGrantType = "urn:ietf:params:oauth:grant-type:device_code" + + // Device code login is approved on another device, which acceptance tests + // cannot drive, so every device authorization is treated as approved with + // a fixed confirmation code that tests can type into the CLI. + fakeDeviceUserCode = "WDJBMJHT" + fakeDeviceConfirmationCode = "FTFAQQKY" + + // Like the real server, only allowlisted OAuth clients can use device code + // login. The fake allows only the CLI's default client. + fakeDeviceClientID = "databricks-cli" ) // FakeOidc holds OAuth state for acceptance tests. type FakeOidc struct { url string + + mu sync.Mutex + + // deviceCodes maps device codes that have not been exchanged yet to + // whether offline_access was requested for them. + deviceCodes map[string]bool + lastDeviceCode int } func (s *FakeOidc) OidcEndpoints() Response { return Response{ Body: map[string]string{ - "authorization_endpoint": s.url + "/oidc/v1/authorize", - "token_endpoint": s.url + "/oidc/v1/token", + "authorization_endpoint": s.url + "/oidc/v1/authorize", + "token_endpoint": s.url + "/oidc/v1/token", + "device_authorization_endpoint": s.url + "/oidc/v1/device/auth", + }, + } +} + +// OidcDeviceAuth implements the RFC 8628 device authorization endpoint. +func (s *FakeOidc) OidcDeviceAuth(req Request) Response { + form, err := url.ParseQuery(string(req.Body)) + if err != nil || form.Get("client_id") == "" { + return oauthErrorResponse("invalid_request", "Bad Request") + } + if form.Get("client_id") != fakeDeviceClientID { + return oauthErrorResponse("unauthorized_client", "Client is not allowed to use device code login") + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.deviceCodes == nil { + s.deviceCodes = map[string]bool{} + } + s.lastDeviceCode++ + deviceCode := "device-code-" + strconv.Itoa(s.lastDeviceCode) + s.deviceCodes[deviceCode] = slices.Contains(strings.Fields(form.Get("scope")), "offline_access") + + return Response{ + Body: map[string]any{ + "device_code": deviceCode, + "user_code": fakeDeviceUserCode, + "verification_uri": s.url + "/device/verification", + "expires_in": 600, + "interval": 0, + }, + } +} + +// deviceCodeToken exchanges a device code and confirmation code for a token. +// Like the real server, a wrong confirmation code leaves the device code +// valid, and a successful exchange consumes it. +func (s *FakeOidc) deviceCodeToken(form url.Values) Response { + s.mu.Lock() + defer s.mu.Unlock() + deviceCode := form.Get("device_code") + offlineAccess, ok := s.deviceCodes[deviceCode] + if !ok { + return oauthErrorResponse("invalid_request", "Invalid device code") + } + if form.Get("confirmation_code") != fakeDeviceConfirmationCode { + return oauthErrorResponse("invalid_request", "Invalid confirmation code") + } + delete(s.deviceCodes, deviceCode) + + body := map[string]any{ + "access_token": "oauth-token", + "expires_in": 3600, + "scope": "all-apis", + "token_type": "Bearer", + } + if offlineAccess { + body["refresh_token"] = "oauth-refresh-token" + } + return Response{Body: body} +} + +func oauthErrorResponse(code, description string) Response { + return Response{ + StatusCode: http.StatusBadRequest, + Body: map[string]string{ + "error": code, + "error_description": description, }, } } @@ -49,6 +143,10 @@ func (s *FakeOidc) OidcAuthorize(req Request) Response { } func (s *FakeOidc) OidcToken(req Request) Response { + form, err := url.ParseQuery(string(req.Body)) + if err == nil && form.Get("grant_type") == deviceCodeGrantType { + return s.deviceCodeToken(form) + } return Response{ Body: map[string]string{ "access_token": "oauth-token", diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index e1a14b9800b..8ea561ca2ba 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -551,6 +551,10 @@ func AddDefaultHandlers(server *Server) { return server.fakeOidc.OidcToken(req) }) + server.Handle("POST", "/oidc/v1/device/auth", func(req Request) any { + return server.fakeOidc.OidcDeviceAuth(req) + }) + server.Handle("POST", "/telemetry-ext", func(_ Request) any { return map[string]any{ "errors": []string{},