diff --git a/cmd/auth/login/cmd.go b/cmd/auth/login/cmd.go index 156cde747..59b3de56e 100644 --- a/cmd/auth/login/cmd.go +++ b/cmd/auth/login/cmd.go @@ -78,8 +78,22 @@ func RunE(cmd *cobra.Command, args []string) error { //nolint: cyclop noBrowser, _ := cmd.Flags().GetBool("no-browser") + // Only pass an override when the user actually said something. Left nil, + // DATAROBOT_OAUTH_ENABLED decides, and its default is off. + var oauthOverride *bool + + switch { + case cmd.Flags().Changed("oauth"): + oauth, _ := cmd.Flags().GetBool("oauth") + oauthOverride = &oauth + case cmd.Flags().Changed("no-oauth"): + off := false + oauthOverride = &off + } + key, err := auth.RunBrowserLoginWith(cmd.Context(), datarobotHost, auth.LoginOptions{ NoBrowser: noBrowser, + OAuth: oauthOverride, }) if err != nil { log.Error(err) @@ -119,7 +133,13 @@ This command will: 3. Securely store your API key for future CLI operations. If the browser cannot be opened, the CLI prints a link to open yourself. Pass ---no-browser to skip the browser launch entirely, which is useful over SSH.`, +--no-browser to skip the browser launch entirely, which is useful over SSH. + +Deployments that front their own OAuth2 authorization server can be logged +into with --oauth, which runs a standard authorization-code flow with PKCE +instead of the DataRobot hand-off. The access token then never travels in a +URL. Set DATAROBOT_OAUTH_ENABLED=true to make that the default for a shell; +--no-oauth forces the hand-off back on.`, SilenceErrors: true, SilenceUsage: true, RunE: RunE, @@ -129,5 +149,15 @@ If the browser cannot be opened, the CLI prints a link to open yourself. Pass // per-invocation flag and must never be persisted to drconfig.yaml. cmd.Flags().Bool("no-browser", false, "print the login link instead of opening a browser") + // Same reasoning as --no-browser: transient, never persisted. + // + // Two separate booleans, because pflag does not synthesise --no- variants + // for a bool flag. Read via Flags().Changed so "unset" stays + // distinguishable from an explicit --no-oauth, which is what lets the flag + // override DATAROBOT_OAUTH_ENABLED in both directions. + cmd.Flags().Bool("oauth", false, "log in with OAuth2 authorization-code + PKCE (default: DATAROBOT_OAUTH_ENABLED)") + cmd.Flags().Bool("no-oauth", false, "force the DataRobot hand-off even if DATAROBOT_OAUTH_ENABLED is set") + cmd.MarkFlagsMutuallyExclusive("oauth", "no-oauth") + return cmd } diff --git a/cmd/auth/logout/cmd.go b/cmd/auth/logout/cmd.go index 047451c79..7f894442a 100644 --- a/cmd/auth/logout/cmd.go +++ b/cmd/auth/logout/cmd.go @@ -28,6 +28,11 @@ import ( func RunE(_ *cobra.Command, _ []string) error { viperx.Set(config.DataRobotAPIKey, "") + // Clearing the access token alone would leave a working refresh token on + // disk — logout has to drop the means of getting a new one too, or it does + // not log anyone out. + auth.ClearOAuthState() + err := auth.WriteConfigFile() if err != nil { log.Error(fmt.Errorf("failed to write config: %w", err)) diff --git a/internal/auth/auth.go b/internal/auth/auth.go index 0ed525fd2..acd57557c 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -371,6 +371,15 @@ func EnsureAuthenticated(ctx context.Context) bool { //nolint: cyclop return true } + // A REJECTED token may be renewable without a browser — the reason + // `offline_access` is requested at login. Only on a judged rejection: an + // unjudged failure (404, 5xx, a network blip) is not evidence of expiry, + // and spending a rotate-on-use refresh token against a server that never + // rejected anything would turn an outage into a forced re-login. + if tokenWasRejected(viperErr) && renewStoredToken(ctx) { + return true + } + skipAuthFlow := false // Everything this gate prints goes to stderr: PreRunE runs before the command, @@ -648,3 +657,52 @@ func GetBaseURLOrAsk() string { return datarobotHost } + +// renewStoredToken swaps an expired access token for a fresh one and reports +// whether the profile ends up usable. false means "carry on to the interactive +// login", including the ordinary case of nothing to renew with. +func renewStoredToken(ctx context.Context) bool { + if _, err := RefreshAccessToken(ctx); err != nil { + if !errors.Is(err, ErrNoRefreshToken) { + log.Debugf("Could not renew the access token: %v", err) + } + + return false + } + + // Verify rather than trust: a server handing back a token it will not + // accept would otherwise loop us silently. + if _, err := config.GetAPIKey(ctx); err != nil { + log.Debug("Renewed token did not verify; falling back to interactive login") + ClearOAuthState() + + return false + } + + if err := WriteConfigFileSilent(); err != nil { + log.Error("Failed to write config file.", "error", err) + + return false + } + + log.Debug("Renewed the access token with the stored refresh token") + + return true +} + +// tokenWasRejected reports whether a failure was the server judging the +// credential (401/403) rather than being unable to answer — the same +// distinction fprintServerStatus makes. +func tokenWasRejected(err error) bool { + if err == nil { + return false + } + + var statusErr *config.HTTPStatusError + if !errors.As(err, &statusErr) { + return false + } + + return statusErr.StatusCode == http.StatusUnauthorized || + statusErr.StatusCode == http.StatusForbidden +} diff --git a/internal/auth/browserflow.go b/internal/auth/browserflow.go index b5eef1242..f62487482 100644 --- a/internal/auth/browserflow.go +++ b/internal/auth/browserflow.go @@ -66,14 +66,54 @@ type BrowserFlow struct { keyCh chan string timeout time.Duration + // OAuth mode only (nil for the legacy ?key= hand-off). When set, the + // callback carries a code that handleCallback exchanges before publishing + // on keyCh, so Wait's contract is identical in both modes. + oauthMeta *OAuthMetadata + oauthPKCE *pkce + redirectURI string + + // errCh reports failures from inside the callback handler. It exists + // because an empty string on keyCh already means "another process wants + // this port" (see Wait), so errors cannot be signalled that way. + errCh chan error + + // refreshToken is whatever the token exchange returned, for the caller to + // persist. Empty in legacy mode and whenever the server issues none. + refreshToken string + closeOnce sync.Once closeErr error } // NewBrowserFlow binds the callback listener and prepares the browser login for -// datarobotHost. The caller must Close the returned flow. +// datarobotHost. The caller must Close the returned flow. Prefer +// NewBrowserFlowContext where a context is available. func NewBrowserFlow(datarobotHost string) (*BrowserFlow, error) { - return newBrowserFlowOn(CallbackAddr, datarobotHost) + return NewBrowserFlowContext(context.Background(), datarobotHost, nil) +} + +// NewBrowserFlowContext binds the callback listener, choosing between the +// legacy `?key=` hand-off and PKCE. +// +// Unless someone opts in this issues no discovery request at all, which keeps +// hosts that serve a discovery document without supporting the flow working as +// before. Opted in but undiscoverable returns ErrOAuthNotSupported rather than +// falling back, which would hand over a different credential while looking +// like success. +func NewBrowserFlowContext(ctx context.Context, datarobotHost string, oauthOverride *bool) (*BrowserFlow, error) { + if !OAuthRequested(oauthOverride) { + return newBrowserFlowOn(CallbackAddr, datarobotHost) + } + + meta, err := DiscoverOAuth(ctx, datarobotHost) + if err != nil { + return nil, err + } + + log.Debugf("Using OAuth2 authorization-code + PKCE against %s", meta.Issuer) + + return newBrowserFlowOAuthOn(CallbackAddr, meta) } // newBrowserFlowOn is NewBrowserFlow with a configurable address so tests can @@ -84,10 +124,17 @@ func newBrowserFlowOn(addr, datarobotHost string) (*BrowserFlow, error) { return nil, err } + return newFlow(addr, listener, AuthCallbackURL(datarobotHost)), nil +} + +// newFlow assembles a flow and its callback server. Both constructors go +// through here, differing only in the URL and the OAuth fields set after. +func newFlow(addr string, listener net.Listener, authURL string) *BrowserFlow { flow := &BrowserFlow{ - authURL: AuthCallbackURL(datarobotHost), + authURL: authURL, listener: listener, keyCh: make(chan string, 1), + errCh: make(chan error, 1), timeout: DefaultLoginTimeout, } @@ -100,6 +147,55 @@ func newBrowserFlowOn(addr, datarobotHost string) (*BrowserFlow, error) { ReadHeaderTimeout: 10 * time.Second, } + return flow +} + +// newBrowserFlowOAuthOn binds the listener for a PKCE login. The redirect URI +// is the loopback listener itself (RFC 8252) on CallbackAddr's port, so the +// server must have http://localhost:51164/ registered for OAuthClientID. +func newBrowserFlowOAuthOn(addr string, meta *OAuthMetadata) (*BrowserFlow, error) { + p, err := newPKCE() + if err != nil { + return nil, err + } + + listener, err := listenReclaimingPort(addr) + if err != nil { + return nil, err + } + + // Hostname from addr, port from the listener. listener.Addr() resolves + // "localhost" to 127.0.0.1 and servers match redirect_uri exactly, so the + // IP literal is rejected; the port must be the one actually bound because + // tests bind :0. + redirectHost, _, err := net.SplitHostPort(addr) + if err != nil { + listener.Close() + + return nil, fmt.Errorf("parsing callback address %q: %w", addr, err) + } + + _, boundPort, err := net.SplitHostPort(listener.Addr().String()) + if err != nil { + listener.Close() + + return nil, fmt.Errorf("reading bound callback port: %w", err) + } + + redirectURI := "http://" + net.JoinHostPort(redirectHost, boundPort) + "/" + + authURL, err := authorizeURL(meta, p, redirectURI) + if err != nil { //nolint:wsl // grouped with the construction above + listener.Close() + + return nil, err + } + + flow := newFlow(addr, listener, authURL) + flow.oauthMeta = meta + flow.oauthPKCE = p + flow.redirectURI = redirectURI + return flow, nil } @@ -108,6 +204,21 @@ func (f *BrowserFlow) AuthURL() string { return f.authURL } +// RefreshToken is from the OAuth exchange, or "" in legacy mode / when the +// server issued none. Only meaningful after Wait succeeds. +func (f *BrowserFlow) RefreshToken() string { + return f.refreshToken +} + +// TokenEndpoint is where a renewal is sent, or "" in legacy mode. +func (f *BrowserFlow) TokenEndpoint() string { + if f.oauthMeta == nil { + return "" + } + + return f.oauthMeta.TokenEndpoint +} + // localAddr reports the address the listener actually bound, which differs from // CallbackAddr only in tests. func (f *BrowserFlow) localAddr() string { @@ -148,6 +259,13 @@ func (f *BrowserFlow) Wait(ctx context.Context) (string, error) { return apiKey, nil + case err := <-f.errCh: + // The callback arrived but could not be turned into a credential — a + // state mismatch, a refused authorization, or a rejected token + // exchange. Distinct from the empty-key sentinel above so a real + // failure does not masquerade as an interruption. + return "", err + case <-ctx.Done(): if errors.Is(ctx.Err(), context.DeadlineExceeded) { return "", fmt.Errorf("timed out after %s waiting for browser authorization: %w", f.timeout, ctx.Err()) @@ -176,10 +294,49 @@ func (f *BrowserFlow) Close() error { return f.closeErr } -// handleCallback receives the redirect from the DataRobot web app, which carries -// the API key as the "key" query parameter. +// handleCallback receives the redirect that ends the browser login: either +// `?key=` (legacy) or `?code=…&state=…`, which this exchanges so Wait +// returns a usable credential either way. +// +// ORDER MATTERS. The OAuth branch is checked first because a keyless request is +// the port-reclaim interrupt sentinel (see Wait and listenReclaimingPort) — an +// OAuth callback falling through to it would abort the login. func (f *BrowserFlow) handleCallback(w http.ResponseWriter, r *http.Request) { - apiKey := r.URL.Query().Get("key") + query := r.URL.Query() + + if f.oauthMeta != nil { + if code := query.Get("code"); code != "" { + f.handleOAuthCallback(w, r, code, query.Get("state")) + + return + } + + // The authorization server can also report a failure on the redirect, + // e.g. the user declining consent. Surface it rather than sitting until + // the five-minute timeout. + if oauthErr := query.Get("error"); oauthErr != "" { + f.failCallback(w, fmt.Errorf("authorization was refused: %s: %s", oauthErr, query.Get("error_description"))) + + return + } + + // A `key` cannot authenticate THIS flow. Falling through to the legacy + // handler would let any local process inject a credential during the + // login window, bypassing the `state` check that is the whole reason + // this flow exists. + // + // Refused without publishing on either channel, so the login keeps + // waiting for the real callback: neither an injection nor a way to + // cancel someone else's login. + if query.Get("key") != "" { + log.Debug("Ignoring a ?key= callback: this login is an OAuth flow") + http.Error(w, "Unexpected credential on an OAuth callback.", http.StatusBadRequest) + + return + } + } + + apiKey := query.Get("key") w.Header().Set("Content-Type", "text/html; charset=utf-8") @@ -196,11 +353,68 @@ func (f *BrowserFlow) handleCallback(w http.ResponseWriter, r *http.Request) { } } +// handleOAuthCallback validates the redirect and exchanges the code. +func (f *BrowserFlow) handleOAuthCallback(w http.ResponseWriter, r *http.Request, code, state string) { + // Constant-time is unnecessary: state is single-use and anyone who can read + // it already has the code. What matters is that a mismatch is fatal — this + // callback is otherwise open to any local process. + if state != f.oauthPKCE.state { + f.failCallback(w, errors.New("OAuth state mismatch — ignoring a callback this login did not start")) + + return + } + + tok, err := exchangeCode(r.Context(), f.oauthMeta, f.oauthPKCE, f.redirectURI, code) + if err != nil { + f.failCallback(w, err) + + return + } + + // For the caller to persist after Wait. + f.refreshToken = tok.RefreshToken + + if tok.RefreshToken == "" { + log.Debug("Authorization server issued no refresh token; the credential expires without renewal") + } + + w.Header().Set("Content-Type", "text/html; charset=utf-8") + + if writeErr := assets.Write(w, "templates/success.html"); writeErr != nil { + log.Debugf("Failed to render auth success page: %v", writeErr) + } + + select { + case f.keyCh <- tok.AccessToken: + default: + log.Debug("Discarding duplicate auth callback; a credential was already received") + } +} + +// failCallback tells the browser the login failed and hands the reason to Wait. +// It deliberately does not publish an empty string on keyCh: that means +// "release the port" and would turn a failure into a silent interruption. +func (f *BrowserFlow) failCallback(w http.ResponseWriter, err error) { + log.Debugf("Auth callback failed: %v", err) + + http.Error(w, "Login failed: "+err.Error()+"\n\nReturn to the terminal for details.", http.StatusBadRequest) + + select { + case f.errCh <- err: + default: + log.Debug("Discarding duplicate auth failure; one was already reported") + } +} + // LoginOptions tunes the interactive browser login. type LoginOptions struct { // NoBrowser skips launching a browser and shows the link instead. Useful over // SSH or anywhere the CLI cannot reach a usable browser. NoBrowser bool + + // OAuth forces the OAuth2 authorization-code + PKCE flow on or off. nil + // means "not specified", deferring to DATAROBOT_OAUTH_ENABLED (default off). + OAuth *bool } // RunBrowserLogin opens the browser, tells the user what is happening, and blocks @@ -217,7 +431,7 @@ func RunBrowserLogin(ctx context.Context, datarobotHost string) (string, error) // RunBrowserLoginWith is RunBrowserLogin with explicit options. func RunBrowserLoginWith(ctx context.Context, datarobotHost string, opts LoginOptions) (string, error) { - flow, err := NewBrowserFlow(datarobotHost) + flow, err := NewBrowserFlowContext(ctx, datarobotHost, opts.OAuth) if err != nil { return "", err } @@ -291,6 +505,15 @@ func runLoginWithFlow(ctx context.Context, flow *BrowserFlow, opts LoginOptions) return "", err } + // Record what a later renewal needs, or clear stale material when this + // login took the legacy path so an old refresh token is never renewed + // against a different instance. Both call sites persist straight after. + if flow.TokenEndpoint() != "" && flow.RefreshToken() != "" { + StoreOAuthState(flow.RefreshToken(), flow.TokenEndpoint()) + } else { + ClearOAuthState() + } + return apiKey, nil } diff --git a/internal/auth/oauth.go b/internal/auth/oauth.go new file mode 100644 index 000000000..80d396d76 --- /dev/null +++ b/internal/auth/oauth.go @@ -0,0 +1,334 @@ +// Copyright 2026 DataRobot, Inc. and its affiliates. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package auth + +// OAuth2 authorization-code + PKCE login (RFC 8252), for deployments that front +// their own authorization server rather than the SaaS hand-off. The token +// arrives in a POST response body rather than a URL, and the callback carries a +// verifiable `state`. Opt-in, not auto-detected: some hosts serve a discovery +// document without supporting the flow, so its presence proves nothing. + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/datarobot/cli/internal/config" + "github.com/datarobot/cli/internal/config/viperx" + "github.com/datarobot/cli/internal/log" +) + +// OAuthEnabledEnv gates the flow. Unset means no discovery request is made at +// all, so behavior is byte-identical to before. +const OAuthEnabledEnv = "DATAROBOT_OAUTH_ENABLED" + +// OAuthClientID is the public client this CLI presents. No secret — there is +// nowhere safe to keep one on a user's machine — so PKCE is what protects it. +const OAuthClientID = "datarobot-cli" + +const discoveryPath = "/.well-known/openid-configuration" + +// `offline_access` buys renewal without a browser; a server that declines it +// simply returns no refresh token, degrading to re-login on expiry. +var oauthScopes = []string{"openid", "offline_access", "email", "profile", "groups"} + +// The probe runs before anything visible happens, so it must stay short for +// hosts that do not implement discovery. +const discoveryTimeout = 5 * time.Second + +const tokenExchangeTimeout = 30 * time.Second + +// ErrOAuthNotSupported means no usable discovery document. Returned rather than +// falling back: the user asked for OAuth, and quietly issuing a different kind +// of credential would look like success. +var ErrOAuthNotSupported = errors.New("endpoint does not advertise an OAuth2 authorization server") + +// OAuthMetadata is the subset of the discovery document this flow needs. +type OAuthMetadata struct { + Issuer string `json:"issuer"` + AuthorizationEndpoint string `json:"authorization_endpoint"` + TokenEndpoint string `json:"token_endpoint"` +} + +// OAuthRequested reports whether the caller asked for the OAuth flow. override +// is the --oauth/--no-oauth flag and wins; nil defers to the environment. +func OAuthRequested(override *bool) bool { + if override != nil { + return *override + } + + switch strings.ToLower(strings.TrimSpace(os.Getenv(OAuthEnabledEnv))) { + case "1", "true", "yes": + return true + default: + return false + } +} + +// DiscoverOAuth fetches and validates the discovery document. Strict on +// purpose: a single-page app answers 200 with HTML for unknown paths, so +// requiring parseable JSON with both endpoints is what separates a real +// authorization server from a catch-all route. +func DiscoverOAuth(ctx context.Context, datarobotHost string) (*OAuthMetadata, error) { + base := strings.TrimRight(datarobotHost, "/") + + ctx, cancel := context.WithTimeout(ctx, discoveryTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+discoveryPath, nil) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrOAuthNotSupported, err) + } + + req.Header.Set("Accept", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, fmt.Errorf("%w: %s: %w", ErrOAuthNotSupported, base+discoveryPath, err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("%w: %s returned HTTP %d", ErrOAuthNotSupported, base+discoveryPath, resp.StatusCode) + } + + // Bounded read: this is an untrusted endpoint and the document is small. + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return nil, fmt.Errorf("%w: reading %s: %w", ErrOAuthNotSupported, base+discoveryPath, err) + } + + var meta OAuthMetadata + if err := json.Unmarshal(body, &meta); err != nil { + return nil, fmt.Errorf("%w: %s did not return JSON (an SPA catch-all route?)", ErrOAuthNotSupported, base+discoveryPath) + } + + if meta.AuthorizationEndpoint == "" || meta.TokenEndpoint == "" { + return nil, fmt.Errorf("%w: %s is missing authorization_endpoint or token_endpoint", ErrOAuthNotSupported, base+discoveryPath) + } + + log.Debugf("OAuth discovery at %s: issuer=%s", base+discoveryPath, meta.Issuer) + + return &meta, nil +} + +// pkce is one attempt's verifier and CSRF state. +type pkce struct { + verifier string + challenge string + state string +} + +func newPKCE() (*pkce, error) { + verifier, err := randomURLSafe(32) + if err != nil { + return nil, fmt.Errorf("generating PKCE verifier: %w", err) + } + + state, err := randomURLSafe(16) + if err != nil { + return nil, fmt.Errorf("generating OAuth state: %w", err) + } + + sum := sha256.Sum256([]byte(verifier)) + + return &pkce{ + verifier: verifier, + challenge: base64.RawURLEncoding.EncodeToString(sum[:]), + state: state, + }, nil +} + +func randomURLSafe(n int) (string, error) { + buf := make([]byte, n) + if _, err := rand.Read(buf); err != nil { + return "", err + } + + return base64.RawURLEncoding.EncodeToString(buf), nil +} + +func authorizeURL(meta *OAuthMetadata, p *pkce, redirectURI string) (string, error) { + u, err := url.Parse(meta.AuthorizationEndpoint) + if err != nil { + return "", fmt.Errorf("parsing authorization_endpoint %q: %w", meta.AuthorizationEndpoint, err) + } + + q := u.Query() + q.Set("client_id", OAuthClientID) + q.Set("response_type", "code") + q.Set("redirect_uri", redirectURI) + q.Set("scope", strings.Join(oauthScopes, " ")) + q.Set("state", p.state) + q.Set("code_challenge", p.challenge) + q.Set("code_challenge_method", "S256") + u.RawQuery = q.Encode() + + return u.String(), nil +} + +type tokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + Error string `json:"error"` + ErrorDescription string `json:"error_description"` +} + +// exchangeCode trades the authorization code for an access token. This is the +// step the legacy flow lacks, and why the token never appears in a URL. +func exchangeCode(ctx context.Context, meta *OAuthMetadata, p *pkce, redirectURI, code string) (*tokenResponse, error) { + form := url.Values{} + form.Set("grant_type", "authorization_code") + form.Set("code", code) + form.Set("redirect_uri", redirectURI) + form.Set("code_verifier", p.verifier) + + tok, err := postToken(ctx, meta.TokenEndpoint, form) + if err != nil { + return nil, fmt.Errorf("token exchange failed: %w", err) + } + + return tok, nil +} + +// ErrNoRefreshToken means the profile has nothing to renew with: the legacy +// hand-off, or a server that issued none. +var ErrNoRefreshToken = errors.New("no refresh token stored for this profile") + +// RefreshAccessToken renews the stored access token without a browser. Writes +// into viper but does not persist — the caller owns drconfig.yaml. +// ErrNoRefreshToken means "fall back to interactive login", not failure. +func RefreshAccessToken(ctx context.Context) (string, error) { + refresh := viperx.GetString(config.OAuthRefreshToken) + endpoint := viperx.GetString(config.OAuthTokenEndpoint) + + if refresh == "" || endpoint == "" { + return "", ErrNoRefreshToken + } + + tok, err := postRefresh(ctx, endpoint, refresh) + if err != nil { + return "", err + } + + viperx.Set(config.DataRobotAPIKey, tok.AccessToken) + + // Rotating servers invalidate the old value on use, so a rotated one MUST + // replace what is stored or the next renewal fails. + if tok.RefreshToken != "" { + viperx.Set(config.OAuthRefreshToken, tok.RefreshToken) + } + + return tok.AccessToken, nil +} + +// StoreOAuthState records what a later renewal needs. Persisting is the +// caller's job. +func StoreOAuthState(refreshToken, tokenEndpoint string) { + viperx.Set(config.OAuthRefreshToken, refreshToken) + viperx.Set(config.OAuthTokenEndpoint, tokenEndpoint) +} + +// ClearOAuthState drops the renewal material. Called on logout, and whenever a +// login takes the legacy path — a refresh token left beside a hand-off +// credential would renew against whatever instance issued it, not necessarily +// the one now configured. +func ClearOAuthState() { + viperx.Set(config.OAuthRefreshToken, "") + viperx.Set(config.OAuthTokenEndpoint, "") +} + +// postRefresh performs the refresh_token grant. +// +// A non-200 clears the stored material: a rejected refresh token is spent, so +// keeping it would make every later command retry something that cannot work. +func postRefresh(ctx context.Context, endpoint, refresh string) (*tokenResponse, error) { + form := url.Values{} + form.Set("grant_type", "refresh_token") + form.Set("refresh_token", refresh) + + tok, err := postToken(ctx, endpoint, form) + if err != nil { + // A rejected refresh token is spent, so keeping it would make every + // later command retry something that cannot work. + ClearOAuthState() + + return nil, fmt.Errorf("refresh rejected: %w", err) + } + + return tok, nil +} + +// postToken performs a token-endpoint grant. Shared by the authorization-code +// and refresh grants, which differ only in the form and what rejection means. +func postToken(ctx context.Context, endpoint string, form url.Values) (*tokenResponse, error) { + form.Set("client_id", OAuthClientID) + + ctx, cancel := context.WithTimeout(ctx, tokenExchangeTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode())) + if err != nil { + return nil, fmt.Errorf("building token request: %w", err) + } + + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, fmt.Errorf("calling token endpoint %s: %w", endpoint, err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return nil, fmt.Errorf("reading token response: %w", err) + } + + var tok tokenResponse + + // Decode first: error replies are JSON too, and error_description is the + // actionable half. + unmarshalErr := json.Unmarshal(body, &tok) + + if resp.StatusCode != http.StatusOK { + if tok.Error != "" { + return &tok, fmt.Errorf("%s: %s", tok.Error, tok.ErrorDescription) + } + + return &tok, fmt.Errorf("HTTP %d", resp.StatusCode) + } + + if unmarshalErr != nil { + return nil, fmt.Errorf("token endpoint returned unparseable JSON: %w", unmarshalErr) + } + + if tok.AccessToken == "" { + return nil, errors.New("token endpoint returned no access_token") + } + + return &tok, nil +} diff --git a/internal/auth/oauth_test.go b/internal/auth/oauth_test.go new file mode 100644 index 000000000..e3a1fec24 --- /dev/null +++ b/internal/auth/oauth_test.go @@ -0,0 +1,403 @@ +// Copyright 2026 DataRobot, Inc. and its affiliates. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package auth + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/datarobot/cli/internal/config" + "github.com/datarobot/cli/internal/config/viperx" +) + +// testMeta is metadata pointing at a stub token endpoint. +func testMeta(tokenURL string) *OAuthMetadata { + return &OAuthMetadata{ + AuthorizationEndpoint: "https://as.example/oauth2/auth", + TokenEndpoint: tokenURL, + } +} + +// newTestOAuthFlow binds an OAuth flow on an ephemeral port against meta. +func newTestOAuthFlow(t *testing.T, meta *OAuthMetadata) *BrowserFlow { + t.Helper() + + flow, err := newBrowserFlowOAuthOn("127.0.0.1:0", meta) + require.NoError(t, err) + + t.Cleanup(func() { _ = flow.Close() }) + + return flow +} + +// tokenServer stands in for an authorization server's token endpoint. +func tokenServer(t *testing.T, status int, body string) *httptest.Server { + t.Helper() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + + t.Cleanup(srv.Close) + + return srv +} + +// Strict on purpose: a single-page app answers 200 with HTML for unknown paths, +// which would otherwise look like a successful discovery. +func TestDiscoverOAuth(t *testing.T) { + const valid = `{"issuer":"https://as.example","authorization_endpoint":"https://as.example/oauth2/auth","token_endpoint":"https://as.example/oauth2/token"}` + + tests := []struct { + name, contentType, body string + wantOK bool + }{ + {"HTML catch-all", "text/html", "app shell", false}, + {"missing token_endpoint", "application/json", `{"authorization_endpoint":"https://as.example/auth"}`, false}, + {"complete document", "application/json", valid, true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", tc.contentType) + _, _ = w.Write([]byte(tc.body)) + })) + defer srv.Close() + + meta, err := DiscoverOAuth(context.Background(), srv.URL) + + if tc.wantOK { + require.NoError(t, err) + assert.Equal(t, "https://as.example/oauth2/token", meta.TokenEndpoint) + + return + } + + require.Error(t, err) + assert.ErrorIs(t, err, ErrOAuthNotSupported) + }) + } +} + +// The authorize URL must carry PKCE and state, and name the bound port. +func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { + flow := newTestOAuthFlow(t, testMeta("https://as.example/oauth2/token")) + + u, err := url.Parse(flow.AuthURL()) + require.NoError(t, err) + + q := u.Query() + assert.Equal(t, OAuthClientID, q.Get("client_id")) + assert.Equal(t, "code", q.Get("response_type")) + assert.Equal(t, "S256", q.Get("code_challenge_method")) + assert.NotEmpty(t, q.Get("code_challenge")) + assert.NotEmpty(t, q.Get("state")) + assert.Equal(t, "http://"+flow.localAddr()+"/", q.Get("redirect_uri")) + + // The verifier is the secret half and must never leave the process. + assert.NotContains(t, flow.AuthURL(), flow.oauthPKCE.verifier) +} + +// Regression: listener.Addr() resolves "localhost" to 127.0.0.1, and +// authorization servers compare redirect_uri as an exact string. A server that +// registered http://localhost:51164/ rejects the IP literal outright, before +// the user ever sees a login page. The +// hostname must come from the configured address, the port from the listener. +func TestOAuthFlow_RedirectURIKeepsConfiguredHostname(t *testing.T) { + flow, err := newBrowserFlowOAuthOn("localhost:0", testMeta("https://as.example/oauth2/token")) + require.NoError(t, err) + + t.Cleanup(func() { _ = flow.Close() }) + + assert.Contains(t, flow.redirectURI, "http://localhost:", + "redirect_uri must keep the configured hostname, not the resolved IP") + assert.NotContains(t, flow.redirectURI, "127.0.0.1") + + // And it must still name the port actually bound, not the requested :0. + _, boundPort, err := net.SplitHostPort(flow.localAddr()) + require.NoError(t, err) + assert.Contains(t, flow.redirectURI, boundPort) + assert.NotContains(t, flow.redirectURI, ":0/") +} + +func TestOAuthFlow_ExchangesCodeForToken(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"issued-token","token_type":"bearer","refresh_token":"r"}`) + + flow := newTestOAuthFlow(t, testMeta(srv.URL)) + + token, err := driveFlow(t, flow, "?code=abc&state="+flow.oauthPKCE.state) + + require.NoError(t, err) + assert.Equal(t, "issued-token", token, "Wait must return the exchanged access token, not the code") +} + +// The listener is open to any local process; state is the only thing +// distinguishing a callback this login started. +func TestOAuthFlow_RejectsStateMismatch(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"should-never-be-used"}`) + + flow := newTestOAuthFlow(t, testMeta(srv.URL)) + + token, err := driveFlow(t, flow, "?code=abc&state=not-the-right-state") + + require.Error(t, err) + assert.Empty(t, token) + assert.Contains(t, err.Error(), "state mismatch") + assert.NotErrorIs(t, err, ErrLoginInterrupted, + "a rejected callback is a failure, not the port-reclaim interrupt") +} + +// driveFlow fires the callback and returns Wait's result. +func driveFlow(t *testing.T, flow *BrowserFlow, query string) (string, error) { + t.Helper() + + go func() { + resp, err := http.Get(callbackURL(t, flow, query)) + if err == nil { + _ = resp.Body.Close() + } + }() + + return flow.Wait(context.Background()) +} + +// THE regression that matters: an empty `key` is the port-reclaim sentinel, so +// the OAuth branch must be evaluated first or every callback aborts the login. +func TestOAuthFlow_CodeCallbackIsNotSwallowedBySentinel(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"survived"}`) + + flow := newTestOAuthFlow(t, testMeta(srv.URL)) + + // No `key` parameter at all — exactly the shape the sentinel looks for. + token, err := driveFlow(t, flow, "?code=abc&state="+flow.oauthPKCE.state) + + require.NoError(t, err, "an OAuth callback must not be read as the interrupt sentinel") + assert.Equal(t, "survived", token) +} + +// In OAuth mode a `?key=` must never be accepted. Falling through to the legacy +// handler would let any local process inject a credential during the login +// window, bypassing the `state` check this flow exists to provide. +func TestOAuthFlow_RefusesInjectedKey(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"issued-token"}`) + flow := newTestOAuthFlow(t, testMeta(srv.URL)) + + // Wait starts the callback server, so both requests have to be made while + // it is running — hence the result channel rather than a plain call. + result := make(chan string, 1) + + go func() { + token, err := flow.Wait(context.Background()) + if err != nil { + t.Errorf("Wait: %v", err) + } + + result <- token + }() + + // The injected callback must be refused outright. + resp, err := http.Get(callbackURL(t, flow, "?key=injected-by-a-local-process")) + require.NoError(t, err) + + _ = resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + // ...and the genuine one must still complete the login. + genuine, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) + require.NoError(t, err) + + _ = genuine.Body.Close() + + select { + case token := <-result: + assert.Equal(t, "issued-token", token) + assert.NotEqual(t, "injected-by-a-local-process", token, + "an injected key must never become the stored credential") + case <-time.After(5 * time.Second): + t.Fatal("login did not complete; the injected key may have consumed it") + } +} + +// Gate off must not reach out at all — this protects hosts that serve a +// discovery document without supporting the flow. +func TestNewBrowserFlowContext_GateOffIssuesNoDiscovery(t *testing.T) { + t.Setenv(OAuthEnabledEnv, "") + + probed := false + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + probed = true + + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + flow, err := NewBrowserFlowContext(context.Background(), srv.URL, nil) + require.NoError(t, err) + + defer func() { _ = flow.Close() }() + + assert.False(t, probed, "no request may be made to the endpoint when the gate is off") + assert.Nil(t, flow.oauthMeta, "gate off must produce a legacy flow") + assert.Contains(t, flow.AuthURL(), "/account/developer-tools?cliRedirect=true") +} + +// Gate on but undiscoverable must fail loudly, not downgrade silently. +func TestNewBrowserFlowContext_GateOnWithoutDiscoveryFails(t *testing.T) { + t.Setenv(OAuthEnabledEnv, "true") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.NotFound(w, r) + })) + defer srv.Close() + + _, err := NewBrowserFlowContext(context.Background(), srv.URL, nil) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrOAuthNotSupported) +} + +// refreshServer is a token endpoint that records the form it received. +type refreshServer struct { + *httptest.Server + gotForm url.Values +} + +func newRefreshServer(t *testing.T, status int, body string) *refreshServer { + t.Helper() + + rs := &refreshServer{} + rs.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + r.Body = http.MaxBytesReader(w, r.Body, 1<<20) + _ = r.ParseForm() + rs.gotForm = r.PostForm + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + + t.Cleanup(rs.Close) + + return rs +} + +// A profile authenticated by the legacy hand-off has nothing to renew with. +// That is not a failure — it means "go and do an interactive login". +func TestRefreshAccessToken_NoStoredMaterial(t *testing.T) { + _, cleanup := setupTestEnvironment(t) + defer cleanup() + + ClearOAuthState() + + _, err := RefreshAccessToken(context.Background()) + + assert.ErrorIs(t, err, ErrNoRefreshToken) +} + +func TestRefreshAccessToken_Success(t *testing.T) { + _, cleanup := setupTestEnvironment(t) + defer cleanup() + + rs := newRefreshServer(t, http.StatusOK, `{"access_token":"renewed","token_type":"bearer"}`) + StoreOAuthState("stored-refresh", rs.URL) + + token, err := RefreshAccessToken(context.Background()) + + require.NoError(t, err) + assert.Equal(t, "renewed", token) + assert.Equal(t, "renewed", viperx.GetString(config.DataRobotAPIKey), + "the renewed token must land where the API layer reads it") + + assert.Equal(t, "refresh_token", rs.gotForm.Get("grant_type")) + assert.Equal(t, "stored-refresh", rs.gotForm.Get("refresh_token")) + assert.Equal(t, OAuthClientID, rs.gotForm.Get("client_id")) + + // Not rotated by this server, so the stored one must still be there. + assert.Equal(t, "stored-refresh", viperx.GetString(config.OAuthRefreshToken)) +} + +// Servers that rotate refresh tokens invalidate the old one on use, so a +// returned value has to replace what is stored or the NEXT renewal fails. +func TestRefreshAccessToken_StoresRotatedRefreshToken(t *testing.T) { + _, cleanup := setupTestEnvironment(t) + defer cleanup() + + rs := newRefreshServer(t, http.StatusOK, + `{"access_token":"renewed","refresh_token":"rotated"}`) + StoreOAuthState("stored-refresh", rs.URL) + + _, err := RefreshAccessToken(context.Background()) + + require.NoError(t, err) + assert.Equal(t, "rotated", viperx.GetString(config.OAuthRefreshToken), + "a rotated refresh token must replace the spent one") +} + +// A rejected refresh token is spent. Keeping it would make every subsequent +// command retry something that cannot work before falling back. +func TestRefreshAccessToken_RejectionClearsState(t *testing.T) { + _, cleanup := setupTestEnvironment(t) + defer cleanup() + + rs := newRefreshServer(t, http.StatusBadRequest, + `{"error":"invalid_grant","error_description":"token is expired"}`) + StoreOAuthState("stored-refresh", rs.URL) + + _, err := RefreshAccessToken(context.Background()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid_grant") + assert.Contains(t, err.Error(), "token is expired", "the description is the actionable half") + assert.Empty(t, viperx.GetString(config.OAuthRefreshToken), + "a spent refresh token must not be kept") + assert.Empty(t, viperx.GetString(config.OAuthTokenEndpoint)) +} + +// THE guard. Only a judged rejection means the token expired. Treating a 404, +// 429, 5xx or transport error as expiry would spend a refresh token — possibly +// a rotate-on-use one — against a server that never rejected anything, turning +// a transient outage into a forced re-login. +func TestTokenWasRejected(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"401 judged", &config.HTTPStatusError{StatusCode: 401}, true}, + {"403 judged", &config.HTTPStatusError{StatusCode: 403}, true}, + {"5xx unjudged", &config.HTTPStatusError{StatusCode: 503}, false}, + {"transport error unjudged", assert.AnError, false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, tokenWasRejected(tc.err)) + }) + } +} diff --git a/internal/config/constants.go b/internal/config/constants.go index e703cfdd0..55761621a 100644 --- a/internal/config/constants.go +++ b/internal/config/constants.go @@ -18,6 +18,17 @@ const ( DataRobotURL = "endpoint" DataRobotAPIKey = "token" + // OAuthRefreshToken and OAuthTokenEndpoint are written only by the OAuth2 + // login flow. Both are needed to renew an access token without a browser: + // the endpoint is stored rather than re-discovered so a renewal costs one + // request and does not depend on discovery still being reachable. + // + // Their presence is also what marks a profile as OAuth-authenticated, which + // is why the legacy flow clears them — a stale refresh token left beside a + // hand-off credential would be renewed against the wrong instance. + OAuthRefreshToken = "oauth-refresh-token" + OAuthTokenEndpoint = "oauth-token-endpoint" + APIConsumerTrackingEnabled = "api-consumer-tracking-enabled" // SkipAuthKey is the viper key behind the --skip-auth persistent flag. diff --git a/internal/config/write.go b/internal/config/write.go index e40a53334..8aaf7b4b4 100644 --- a/internal/config/write.go +++ b/internal/config/write.go @@ -42,6 +42,8 @@ var PersistableKeys = map[string]struct{}{ "pulumi_config_passphrase": {}, "ca-cert": {}, DefaultLLMID: {}, + OAuthRefreshToken: {}, + OAuthTokenEndpoint: {}, } // UpdateConfigFile writes only the allowlisted keys from viper back to the