diff --git a/cmd/auth/token.go b/cmd/auth/token.go index bbde8dc5eff..0350d7056df 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -50,7 +50,7 @@ and secret is not supported.`, var tokenTimeout time.Duration cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTimeout, - "Timeout for acquiring a token.") + "Timeout for acquiring a token, including any browser login.") var forceRefresh bool cmd.Flags().BoolVar(&forceRefresh, "force-refresh", false, @@ -199,7 +199,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { // resolve the target through environment variables or interactive profile selection. if args.profileName == "" && args.authArguments.Host == "" && len(args.args) == 0 { var resolvedProfile string - resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenStore, args.mode) + resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenStore, args.mode, args.tokenTimeout) if err != nil { return nil, err } @@ -347,7 +347,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { // // Returns the resolved profile name and profile (if any). The host and related // fields on authArgs are updated in place when resolved via environment variables. -func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenStore storage.Store, mode storage.StorageMode) (string, *profile.Profile, error) { +func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenStore storage.Store, mode storage.StorageMode, timeout time.Duration) (string, *profile.Profile, error) { // Step 1: Try DATABRICKS_HOST env var (highest priority). if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" { authArgs.Host = envHost @@ -410,7 +410,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs // Fall through — setHostAndAccountId will prompt for the host. return "", nil, nil case profilePickerCreateNew: - return runInlineLogin(ctx, profiler, tokenStore, mode) + return runInlineLogin(ctx, profiler, tokenStore, mode, timeout) default: p, err := loadProfileByName(ctx, selectedName, profiler) if err != nil { @@ -425,7 +425,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs // .databrickscfg, and returns the new profile name and profile. It rejects // the names of existing sections so that it cannot log in to an existing // profile. -func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore storage.Store, mode storage.StorageMode) (string, *profile.Profile, error) { +func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore storage.Store, mode storage.StorageMode, timeout time.Duration) (string, *profile.Profile, error) { takenNames, err := databrickscfg.GetSectionNames(ctx, env.Get(ctx, "DATABRICKS_CONFIG_FILE")) if err != nil { return "", nil, err @@ -455,7 +455,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s // Challenge uses the context passed to NewPersistentAuth, so the timeout // must be set before it. - ctx, cancel := context.WithTimeout(ctx, defaultTimeout) + ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...) diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index 001aa527034..8bcdd48f363 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "net/http" + "net/http/httptest" "net/url" "runtime" "strings" @@ -984,7 +985,7 @@ host = https://prod.cloud.databricks.test cluster_id = abc `)) tm := termtest.Run(t, func(ctx context.Context) (string, error) { - name, _, err := resolveNoArgsToken(ctx, profile.DefaultProfiler, &auth.AuthArguments{}, nil, storage.StorageModePlaintext) + name, _, err := resolveNoArgsToken(ctx, profile.DefaultProfiler, &auth.AuthArguments{}, nil, storage.StorageModePlaintext, defaultTimeout) return name, err }) tm.WaitFor(profilePickerCreateNewLabel) @@ -998,3 +999,47 @@ cluster_id = abc _, err := tm.Result() assert.EqualError(t, err, "^C") } + +func TestLoadTokenCreateNewProfileAppliesTimeout(t *testing.T) { + // See TestLoginDeviceCodePickerEmptyHost. + if runtime.GOOS == "windows" { + t.Skip("bubbletea cannot cancel reads on non-console input on Windows") + } + // Either one would resolve the profile before the picker is shown. + t.Setenv("DATABRICKS_HOST", "") + t.Setenv("DATABRICKS_CONFIG_PROFILE", "") + t.Setenv("DATABRICKS_CONFIG_FILE", writeTempConfig(t, `[prod] +host = https://prod.cloud.databricks.test +`)) + t.Setenv("BROWSER", "none") + + // Host discovery runs before the login timeout starts and has its own 5s + // limit, so answer it immediately. Every other request blocks until the + // login's context is done. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + w.WriteHeader(http.StatusNotFound) + return + } + <-r.Context().Done() + })) + t.Cleanup(server.Close) + + tm := termtest.Run(t, func(ctx context.Context) (*oauth2.Token, error) { + return loadToken(ctx, loadTokenArgs{ + authArguments: &auth.AuthArguments{}, + tokenTimeout: 100 * time.Millisecond, + profiler: profile.DefaultProfiler, + tokenStore: storage.NewMemoryStore(), + mode: storage.StorageModePlaintext, + }) + }) + tm.WaitFor(profilePickerCreateNewLabel) + tm.Type(termtest.KeyDown + termtest.KeyEnter) + tm.WaitFor("Databricks profile name [DEFAULT]") + tm.Type("new" + termtest.KeyEnter) + tm.WaitFor("Databricks host") + tm.Type(server.URL + termtest.KeyEnter) + _, err := tm.Result() + assert.ErrorIs(t, err, context.DeadlineExceeded) +}