From 2cc8b8700feab42837f7dda8ebfb14ccdf9ca7ec Mon Sep 17 00:00:00 2001 From: Erik Lattimore Date: Thu, 3 Sep 2026 06:50:18 -0400 Subject: [PATCH 1/4] feat(auth): OAuth2 authorization-code + PKCE login, opt-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Deployments that front their own authorization server — a self-hosted inference stack running Ory Hydra, say — have no /account/developer-tools handler, and the existing flow's ?key= hand-off puts the credential in a URL, where it lands in browser history. This adds a standard authorization-code flow with PKCE: the CLI holds the verifier, receives only a code on the loopback redirect, and exchanges it over POST, so the token arrives in a response body. It also gives the callback a real `state` to check, which the hand-off cannot provide. Off by default. `dr auth login` is unchanged unless asked, and issues no discovery request at all — deliberately, because some deployments serve an OIDC discovery document without a login service behind it, so the document's presence cannot be treated as evidence the flow works. Enable per-shell with DATAROBOT_OAUTH_ENABLED=true, or per-invocation with --oauth; --no-oauth forces the hand-off back on. Asked for but undiscoverable is an error, not a silent downgrade to a different kind of credential. Wait() still returns "the credential" in both modes, so callers need no changes. The OAuth branch in handleCallback is evaluated BEFORE the keyless-request check: an empty `key` is the port-reclaim interrupt sentinel, so an OAuth callback reaching it would abort the login rather than complete it. Callback failures travel on their own channel for the same reason — publishing an empty string would read as an interruption instead of an error. The redirect URI takes its hostname from the configured callback address and its port from the bound listener. listener.Addr() resolves localhost to 127.0.0.1, and servers match redirect_uri as an exact string, so the IP literal gets the request rejected before any login page appears. Verified against a live Hydra: discovery, PKCE authorize request, and Hydra accepting client/redirect_uri/scope and redirecting to its login provider. --- cmd/auth/login/cmd.go | 23 ++- internal/auth/browserflow.go | 218 ++++++++++++++++++++- internal/auth/oauth.go | 270 ++++++++++++++++++++++++++ internal/auth/oauth_test.go | 354 +++++++++++++++++++++++++++++++++++ 4 files changed, 859 insertions(+), 6 deletions(-) create mode 100644 internal/auth/oauth.go create mode 100644 internal/auth/oauth_test.go diff --git a/cmd/auth/login/cmd.go b/cmd/auth/login/cmd.go index 156cde747..cba8aa828 100644 --- a/cmd/auth/login/cmd.go +++ b/cmd/auth/login/cmd.go @@ -78,8 +78,18 @@ 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 + + if cmd.Flags().Changed("oauth") { + oauth, _ := cmd.Flags().GetBool("oauth") + oauthOverride = &oauth + } + key, err := auth.RunBrowserLoginWith(cmd.Context(), datarobotHost, auth.LoginOptions{ NoBrowser: noBrowser, + OAuth: oauthOverride, }) if err != nil { log.Error(err) @@ -119,7 +129,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 — such as a +self-hosted inference stack — 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 +145,10 @@ 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. Registered as + // one boolean so cobra gives us --oauth and --no-oauth for free, and read + // via Flags().Changed so "unset" stays distinguishable from "--no-oauth". + cmd.Flags().Bool("oauth", false, "log in with OAuth2 authorization-code + PKCE (default: DATAROBOT_OAUTH_ENABLED)") + return cmd } diff --git a/internal/auth/browserflow.go b/internal/auth/browserflow.go index b5eef1242..2700506ae 100644 --- a/internal/auth/browserflow.go +++ b/internal/auth/browserflow.go @@ -66,14 +66,59 @@ type BrowserFlow struct { keyCh chan string timeout time.Duration + // OAuth mode only (nil for the legacy ?key= hand-off). When set, the + // callback carries an authorization code that handleCallback exchanges for + // a token before publishing it on keyCh, so Wait's contract — "returns the + // credential" — is identical in both modes and callers need not care. + oauthMeta *OAuthMetadata + oauthPKCE *pkce + redirectURI string + + // errCh reports a failure that happened inside the callback handler, e.g. a + // state mismatch or a rejected token exchange. 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 + closeOnce sync.Once closeErr error } // NewBrowserFlow binds the callback listener and prepares the browser login for // datarobotHost. The caller must Close the returned flow. +// +// Prefer NewBrowserFlowContext where a context is available; this exists for +// callers that have none, and only matters when the OAuth flow is enabled, where +// it bounds the discovery probe on its own. 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 OAuth2 authorization-code + PKCE. +// +// oauthOverride comes from an explicit --oauth/--no-oauth flag and wins; nil +// defers to DATAROBOT_OAUTH_ENABLED, which is off by default. So unless someone +// opts in, this binds exactly the listener it always has and issues no discovery +// request at all — which is what keeps deployments that serve a discovery +// document without a login service behind it working as before. +// +// When OAuth IS requested and discovery does not produce a usable document this +// returns ErrOAuthNotSupported rather than falling back. Falling back would hand +// the user a different kind of 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 @@ -88,6 +133,7 @@ func newBrowserFlowOn(addr, datarobotHost string) (*BrowserFlow, error) { authURL: AuthCallbackURL(datarobotHost), listener: listener, keyCh: make(chan string, 1), + errCh: make(chan error, 1), timeout: DefaultLoginTimeout, } @@ -103,6 +149,78 @@ func newBrowserFlowOn(addr, datarobotHost string) (*BrowserFlow, error) { return flow, nil } +// newBrowserFlowOAuthOn binds the callback listener for an OAuth2 +// authorization-code + PKCE login against the server described by meta. +// +// The redirect URI is the loopback listener itself, which is what lets the code +// come back to this process; RFC 8252 blesses exactly this shape for native +// apps. The port is still CallbackAddr's, so the authorization server must have +// http://localhost:51164/ registered as a redirect URI 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 + } + + // The redirect URI must keep addr's HOSTNAME and take the listener's PORT. + // + // Both halves matter. listener.Addr() resolves "localhost" to 127.0.0.1, and + // authorization servers match redirect_uri as an exact string — Hydra + // registers http://localhost:51164/, so sending the IP literal instead gets + // the request rejected outright. The port has to come from the listener + // because tests bind :0 and the code must come back to the port actually + // bound. + 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 := &BrowserFlow{ + authURL: authURL, + listener: listener, + keyCh: make(chan string, 1), + errCh: make(chan error, 1), + timeout: DefaultLoginTimeout, + oauthMeta: meta, + oauthPKCE: p, + redirectURI: redirectURI, + } + + mux := http.NewServeMux() + mux.HandleFunc("/", flow.handleCallback) + + flow.server = &http.Server{ + Addr: addr, + Handler: mux, + ReadHeaderTimeout: 10 * time.Second, + } + + return flow, nil +} + // AuthURL is the DataRobot URL the user must visit to authorize the CLI. func (f *BrowserFlow) AuthURL() string { return f.authURL @@ -148,6 +266,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 +301,38 @@ 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. +// +// Two shapes arrive here. The legacy DataRobot web app sends the credential +// outright as `?key=`. An OAuth2 authorization server sends +// `?code=&state=`, which this exchanges for a token before +// publishing it, so Wait returns a usable credential either way. +// +// Order matters: the OAuth branch is checked FIRST. A keyless request is the +// port-reclaim interrupt sentinel (see Wait and listenReclaimingPort), so an +// OAuth callback falling through to that check would read as "another CLI wants +// this port" and abort the login instead of completing it. 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 + } + } + + apiKey := query.Get("key") w.Header().Set("Content-Type", "text/html; charset=utf-8") @@ -196,11 +349,66 @@ func (f *BrowserFlow) handleCallback(w http.ResponseWriter, r *http.Request) { } } +// handleOAuthCallback validates the redirect and performs the code exchange. +func (f *BrowserFlow) handleOAuthCallback(w http.ResponseWriter, r *http.Request, code, state string) { + // Constant-time is unnecessary — state is single-use, generated seconds ago, + // and an attacker 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 + } + + 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 value means +// "release the port" and would turn a real 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 +425,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 } diff --git a/internal/auth/oauth.go b/internal/auth/oauth.go new file mode 100644 index 000000000..12021e491 --- /dev/null +++ b/internal/auth/oauth.go @@ -0,0 +1,270 @@ +package auth + +// OAuth2 authorization-code + PKCE login, for DataRobot deployments that front +// their own OAuth2 authorization server (e.g. the inference stack's Hydra) +// rather than the SaaS `/account/developer-tools?cliRedirect=true` hand-off. +// +// Why this exists alongside the legacy flow rather than replacing it: the SaaS +// path returns the credential as a `?key=` query parameter, which means +// the token itself travels in a URL and lands in browser history. Here the CLI +// holds a PKCE verifier, receives only an authorization code on the loopback +// redirect, and exchanges it for the token over POST — so the token arrives in a +// response body. It also gives the callback a real `state` to check, which the +// legacy path has no way to provide. +// +// The legacy path remains the DEFAULT. This one activates only when explicitly +// asked for, because some deployments serve an OIDC discovery document without +// having a login service behind it: the document's presence is not evidence the +// flow is supported, so it cannot be used as an auto-detect signal. + +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/log" +) + +// OAuthEnabledEnv gates the discovery lookup. Unset (the default) means the CLI +// behaves exactly as it always has and issues no discovery request at all. +const OAuthEnabledEnv = "DATAROBOT_OAUTH_ENABLED" + +// OAuthClientID is the public OAuth2 client this CLI presents. Public means no +// client secret — there is nowhere safe to keep one on a user's machine — so the +// authorization server is expected to require PKCE for it. +const OAuthClientID = "datarobot-cli" + +// discoveryPath is the RFC 8414 / OpenID Connect Discovery location, resolved +// against the configured endpoint's scheme+host. +const discoveryPath = "/.well-known/openid-configuration" + +// oauthScopes are requested on every authorization. +// +// `offline_access` is requested deliberately: deployments may issue short-lived +// access tokens, and without a refresh token the user would have to re-run +// `dr auth login` every time one expired. An authorization server that does not +// grant it simply returns no refresh token, which degrades to the same +// re-login-on-expiry behavior the legacy flow already has. +var oauthScopes = []string{"openid", "offline_access", "email", "profile", "groups"} + +// discoveryTimeout bounds the probe. It runs before anything visible happens, so +// it must not be something a user waits on when the endpoint does not implement +// discovery at all. +const discoveryTimeout = 5 * time.Second + +// tokenExchangeTimeout bounds the code-for-token POST. This one runs inside the +// callback HTTP handler, with the user watching a browser tab, so it is short. +const tokenExchangeTimeout = 30 * time.Second + +// ErrOAuthNotSupported means the endpoint did not serve a usable discovery +// document. It is returned rather than silently falling back to the legacy flow: +// the user explicitly asked for OAuth, and quietly handing them a different kind +// of credential instead 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 comes from an explicit --oauth/--no-oauth flag and wins outright; +// nil means "not specified", in which case the environment decides. Default off. +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 for datarobotHost. +// +// It is strict on purpose. A single-page app that serves its shell for unknown +// paths answers 200 with HTML for this URL, which would otherwise look like a +// successful discovery; requiring parseable JSON carrying 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 carries one authorization attempt's PKCE verifier and CSRF state. +type pkce struct { + verifier string + challenge string + state string +} + +// newPKCE generates a fresh verifier/challenge/state triple. +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 +} + +// randomURLSafe returns n bytes of crypto-random data, base64url encoded. +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 +} + +// authorizeURL builds the browser-facing authorization request. +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 +} + +// tokenResponse is the subset of the token endpoint's reply we consume. +type tokenResponse struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + ExpiresIn int `json:"expires_in"` + 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 does not have, and the reason the token never +// appears in a URL: it comes back in this POST's response body. +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("client_id", OAuthClientID) + form.Set("code_verifier", p.verifier) + + ctx, cancel := context.WithTimeout(ctx, tokenExchangeTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, meta.TokenEndpoint, 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("exchanging authorization code at %s: %w", meta.TokenEndpoint, 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 before checking the status: OAuth2 error replies are JSON too, and + // their `error_description` is far more useful than a bare status code. + if err := json.Unmarshal(body, &tok); err != nil && resp.StatusCode == http.StatusOK { + return nil, fmt.Errorf("token endpoint returned unparseable JSON: %w", err) + } + + if resp.StatusCode != http.StatusOK { + if tok.Error != "" { + return nil, fmt.Errorf("token exchange rejected: %s: %s", tok.Error, tok.ErrorDescription) + } + + return nil, fmt.Errorf("token exchange failed with HTTP %d", resp.StatusCode) + } + + 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..c0534d05c --- /dev/null +++ b/internal/auth/oauth_test.go @@ -0,0 +1,354 @@ +// 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" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// 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 +} + +func TestOAuthRequested(t *testing.T) { + yes, no := true, false + + tests := []struct { + name string + env string + override *bool + want bool + }{ + {"unset is off — the whole point of the gate", "", nil, false}, + {"env true", "true", nil, true}, + {"env 1", "1", nil, true}, + {"env garbage is off, not an error", "banana", nil, false}, + {"flag beats unset env", "", &yes, true}, + {"flag --no-oauth beats env true", "true", &no, false}, + {"flag --oauth beats env absent", "", &yes, true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(OAuthEnabledEnv, tc.env) + assert.Equal(t, tc.want, OAuthRequested(tc.override)) + }) + } +} + +// A single-page app answers 200 with its HTML shell for unknown paths. Treating +// that as a successful discovery would send the CLI into an OAuth flow against +// an endpoint that cannot complete one. +func TestDiscoverOAuth_RejectsHTMLCatchAll(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + _, _ = w.Write([]byte("app shell")) + })) + defer srv.Close() + + _, err := DiscoverOAuth(context.Background(), srv.URL) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrOAuthNotSupported) +} + +func TestDiscoverOAuth_RejectsIncompleteDocument(t *testing.T) { + // Valid JSON, but missing token_endpoint — unusable. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"issuer":"http://x","authorization_endpoint":"http://x/auth"}`)) + })) + defer srv.Close() + + _, err := DiscoverOAuth(context.Background(), srv.URL) + + require.Error(t, err) + assert.ErrorIs(t, err, ErrOAuthNotSupported) +} + +func TestDiscoverOAuth_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, discoveryPath, r.URL.Path) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "issuer":"http://hydra", + "authorization_endpoint":"http://hydra/oauth2/auth", + "token_endpoint":"http://hydra/oauth2/token" + }`)) + })) + defer srv.Close() + + meta, err := DiscoverOAuth(context.Background(), srv.URL) + + require.NoError(t, err) + assert.Equal(t, "http://hydra/oauth2/auth", meta.AuthorizationEndpoint) + assert.Equal(t, "http://hydra/oauth2/token", meta.TokenEndpoint) +} + +// The authorize URL must carry PKCE and state, and the redirect URI must name +// the port actually bound or the code comes back to nothing. +func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: "http://hydra/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. Hydra +// registers http://localhost:51164/, so emitting the IP literal gets the +// authorize request rejected 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", &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: "http://hydra/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":"hydra-token","token_type":"bearer","refresh_token":"r"}`) + + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: srv.URL, + }) + + go func() { + resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) + if err == nil { + _ = resp.Body.Close() + } + }() + + token, err := flow.Wait(context.Background()) + + require.NoError(t, err) + assert.Equal(t, "hydra-token", token, "Wait must return the exchanged access token, not the code") +} + +// A callback this login did not start must be refused: the listener is open to +// any local process, and state is the only thing distinguishing them. +func TestOAuthFlow_RejectsStateMismatch(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"should-never-be-used"}`) + + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: srv.URL, + }) + + go func() { + resp, err := http.Get(callbackURL(t, flow, "?code=abc&state=not-the-right-state")) + if err == nil { + _ = resp.Body.Close() + } + }() + + token, err := flow.Wait(context.Background()) + + 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") +} + +func TestOAuthFlow_SurfacesTokenEndpointRejection(t *testing.T) { + srv := tokenServer(t, http.StatusBadRequest, + `{"error":"invalid_grant","error_description":"code already used"}`) + + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: srv.URL, + }) + + go func() { + resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) + if err == nil { + _ = resp.Body.Close() + } + }() + + _, err := flow.Wait(context.Background()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid_grant") + assert.Contains(t, err.Error(), "code already used", "the description is the actionable half") +} + +// The authorization server can report failure on the redirect itself, e.g. the +// user declining consent. That must not wait out the five-minute timeout. +func TestOAuthFlow_SurfacesAuthorizationError(t *testing.T) { + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: "http://hydra/oauth2/token", + }) + + go func() { + resp, err := http.Get(callbackURL(t, flow, "?error=access_denied&error_description=user+declined")) + if err == nil { + _ = resp.Body.Close() + } + }() + + _, err := flow.Wait(context.Background()) + + require.Error(t, err) + assert.Contains(t, err.Error(), "access_denied") + assert.NotErrorIs(t, err, ErrLoginInterrupted) +} + +// THE regression that matters. An empty `key` is the port-reclaim sentinel, so +// the OAuth branch has to be evaluated first — otherwise every OAuth callback +// reads as "another CLI wants this port" and aborts the login. +func TestOAuthFlow_CodeCallbackIsNotSwallowedBySentinel(t *testing.T) { + srv := tokenServer(t, http.StatusOK, `{"access_token":"survived"}`) + + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: srv.URL, + }) + + // No `key` parameter at all — exactly the shape the sentinel looks for. + go func() { + resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) + if err == nil { + _ = resp.Body.Close() + } + }() + + token, err := flow.Wait(context.Background()) + + require.NoError(t, err, "an OAuth callback must not be read as the interrupt sentinel") + assert.Equal(t, "survived", token) +} + +// And the sentinel itself must still work in OAuth mode, or a stale process can +// never be asked to release the port. +func TestOAuthFlow_KeylessCallbackStillInterrupts(t *testing.T) { + flow := newTestOAuthFlow(t, &OAuthMetadata{ + AuthorizationEndpoint: "http://hydra/oauth2/auth", + TokenEndpoint: "http://hydra/oauth2/token", + }) + + go func() { + resp, err := http.Get(callbackURL(t, flow, "")) + if err == nil { + _ = resp.Body.Close() + } + }() + + _, err := flow.Wait(context.Background()) + + assert.ErrorIs(t, err, ErrLoginInterrupted) +} + +// With the gate off, the constructor must not reach out at all — this is what +// protects deployments that serve a discovery document with no login service. +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 discovery unusable must fail loudly. A silent downgrade to the +// legacy flow would look like success while producing a different credential. +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) +} From 227bb224a9e885293d1db8a75a82ff23d2abccaa Mon Sep 17 00:00:00 2001 From: Erik Lattimore Date: Thu, 3 Sep 2026 12:21:05 -0400 Subject: [PATCH 2/4] feat(auth): renew the access token with a refresh token MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The OAuth flow already asked for offline_access but discarded the refresh token, so a deployment issuing short-lived access tokens sent the user through a browser every time one expired — which defeats the point of asking. The refresh token and token endpoint are now persisted with the profile, and a rejected access token is renewed in the background before falling back to an interactive login. `dr auth logout` clears both: dropping the access token while leaving a working refresh token on disk does not log anyone out. Only a JUDGED rejection triggers a renewal. A 404, 429, 5xx or transport error says nothing about the token, and spending a rotate-on-use refresh token against a server that never rejected anything would turn a transient outage into a forced re-login. The renewed token is then verified rather than trusted, so a server handing back a credential it will not accept cannot loop us silently. A rotated refresh token replaces the stored one, since servers that rotate invalidate the old value on use. A rejected refresh clears the stored material so later commands go straight to a browser instead of retrying something spent. Also drops vendor and deployment names from comments, help text and test fixtures: this describes a standard OAuth2 authorization server, and naming a particular one in a public repo is both noise and leakage. --- cmd/auth/login/cmd.go | 10 +- cmd/auth/logout/cmd.go | 5 + internal/auth/auth.go | 70 +++++++++++++ internal/auth/browserflow.go | 44 +++++++- internal/auth/oauth.go | 123 ++++++++++++++++++++++- internal/auth/oauth_test.go | 44 ++++---- internal/auth/refresh_test.go | 184 ++++++++++++++++++++++++++++++++++ internal/config/constants.go | 11 ++ internal/config/write.go | 2 + 9 files changed, 459 insertions(+), 34 deletions(-) create mode 100644 internal/auth/refresh_test.go diff --git a/cmd/auth/login/cmd.go b/cmd/auth/login/cmd.go index cba8aa828..fe1675fb8 100644 --- a/cmd/auth/login/cmd.go +++ b/cmd/auth/login/cmd.go @@ -131,11 +131,11 @@ This command will: 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. -Deployments that front their own OAuth2 authorization server — such as a -self-hosted inference stack — 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.`, +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, 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..57d030f62 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -371,6 +371,20 @@ func EnsureAuthenticated(ctx context.Context) bool { //nolint: cyclop return true } + // A stored token that was actually REJECTED may be renewable without a + // browser. Try that before the interactive flow: an authorization server + // with short access-token lifetimes would otherwise send the user through + // Okta every few minutes, which is the whole reason `offline_access` is + // requested at login. + // + // Only on a judged rejection. An unjudged failure (404, 5xx, a network + // blip) is not evidence the token expired, and burning a + // rotation-on-use refresh token against a server that never rejected + // anything would turn a transient 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 +662,59 @@ func GetBaseURLOrAsk() string { return datarobotHost } + +// renewStoredToken tries to swap an expired access token for a fresh one using +// the stored refresh token, and reports whether the profile ends up usable. +// +// false means "carry on to the interactive login" — including the ordinary case +// of a profile that has 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: the renewed credential has to satisfy the same + // check the stored one just failed, or a server handing back a token it + // will not accept would 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 verification failure was the server +// judging the credential (401/403) rather than being unable to answer. +// +// The distinction is the same one fprintServerStatus makes: a 404, 429, 5xx or +// transport error says nothing about the token, and treating it as expiry would +// spend a refresh token — possibly a rotate-on-use one — on a server that never +// rejected anything. +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 2700506ae..677e1a648 100644 --- a/internal/auth/browserflow.go +++ b/internal/auth/browserflow.go @@ -80,6 +80,10 @@ type BrowserFlow struct { // 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 } @@ -170,11 +174,10 @@ func newBrowserFlowOAuthOn(addr string, meta *OAuthMetadata) (*BrowserFlow, erro // The redirect URI must keep addr's HOSTNAME and take the listener's PORT. // // Both halves matter. listener.Addr() resolves "localhost" to 127.0.0.1, and - // authorization servers match redirect_uri as an exact string — Hydra - // registers http://localhost:51164/, so sending the IP literal instead gets - // the request rejected outright. The port has to come from the listener - // because tests bind :0 and the code must come back to the port actually - // bound. + // authorization servers match redirect_uri as an exact string — a server + // that registered http://localhost:51164/ rejects the IP literal outright. + // The port has to come from the listener because tests bind :0 and the code + // must come back to the port actually bound. redirectHost, _, err := net.SplitHostPort(addr) if err != nil { listener.Close() @@ -226,6 +229,22 @@ func (f *BrowserFlow) AuthURL() string { return f.authURL } +// RefreshToken is the refresh token from the OAuth exchange, or "" when the +// login took the legacy path or the server issued none. Only meaningful after +// Wait has returned successfully. +func (f *BrowserFlow) RefreshToken() string { + return f.refreshToken +} + +// TokenEndpoint is where a renewal should be 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 { @@ -367,6 +386,11 @@ func (f *BrowserFlow) handleOAuthCallback(w http.ResponseWriter, r *http.Request return } + // Kept for the caller to persist after Wait returns. Without a refresh + // token the credential simply expires and the user logs in again — the + // same behavior as the legacy flow. + f.refreshToken = tok.RefreshToken + if tok.RefreshToken == "" { log.Debug("Authorization server issued no refresh token; the credential expires without renewal") } @@ -499,6 +523,16 @@ func runLoginWithFlow(ctx context.Context, flow *BrowserFlow, opts LoginOptions) return "", err } + // Record what a later renewal needs — or clear any stale material when this + // login took the legacy path, so a refresh token from a previous OAuth + // login is never renewed against a different instance. Persisting is the + // caller's job; both call sites write the config immediately 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 index 12021e491..2b3a484c7 100644 --- a/internal/auth/oauth.go +++ b/internal/auth/oauth.go @@ -1,8 +1,8 @@ package auth // OAuth2 authorization-code + PKCE login, for DataRobot deployments that front -// their own OAuth2 authorization server (e.g. the inference stack's Hydra) -// rather than the SaaS `/account/developer-tools?cliRedirect=true` hand-off. +// their own OAuth2 authorization server rather than the SaaS +// `/account/developer-tools?cliRedirect=true` hand-off. // // Why this exists alongside the legacy flow rather than replacing it: the SaaS // path returns the credential as a `?key=` query parameter, which means @@ -32,6 +32,8 @@ import ( "strings" "time" + "github.com/datarobot/cli/internal/config" + "github.com/datarobot/cli/internal/config/viperx" "github.com/datarobot/cli/internal/log" ) @@ -268,3 +270,120 @@ func exchangeCode(ctx context.Context, meta *OAuthMetadata, p *pkce, redirectURI return &tok, nil } + +// ErrNoRefreshToken means this profile has nothing to renew with — either it +// was authenticated by the legacy hand-off, or the authorization server did not +// issue a refresh token. +var ErrNoRefreshToken = errors.New("no refresh token stored for this profile") + +// RefreshAccessToken renews the stored access token without a browser. +// +// This is the whole point of asking for `offline_access`: an authorization +// server that issues short-lived access tokens would otherwise force a full +// interactive login every time one expired. +// +// It writes the new credential into viper but does NOT persist it — the caller +// owns writing drconfig.yaml, matching how the login flow behaves. Returns +// ErrNoRefreshToken when the profile has nothing to renew with, which callers +// should treat as "fall back to interactive login", not as a 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) + + // Servers that rotate refresh tokens invalidate the old one on use, so a + // rotated value MUST replace what is stored or the next renewal fails. + // Servers that do not rotate simply omit it, and the stored one stays good. + if tok.RefreshToken != "" { + viperx.Set(config.OAuthRefreshToken, tok.RefreshToken) + } + + return tok.AccessToken, nil +} + +// StoreOAuthState records what a successful OAuth login needs for later +// renewal. 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 later be renewed against whatever +// instance issued it, which is not necessarily the one now configured. +func ClearOAuthState() { + viperx.Set(config.OAuthRefreshToken, "") + viperx.Set(config.OAuthTokenEndpoint, "") +} + +// postRefresh performs the refresh_token grant and returns the parsed response. +// +// A non-200 clears the stored material before returning: a rejected refresh +// token is spent — revoked, expired, or already rotated — 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) + 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 refresh 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("refreshing access token at %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 refresh response: %w", err) + } + + var tok tokenResponse + + // Decode before checking the status: OAuth2 error replies are JSON too, and + // their error_description is the actionable half. + unmarshalErr := json.Unmarshal(body, &tok) + + if resp.StatusCode != http.StatusOK { + ClearOAuthState() + + if tok.Error != "" { + return nil, fmt.Errorf("refresh rejected: %s: %s", tok.Error, tok.ErrorDescription) + } + + return nil, fmt.Errorf("refresh failed with HTTP %d", resp.StatusCode) + } + + if unmarshalErr != nil { + return nil, fmt.Errorf("refresh endpoint returned unparseable JSON: %w", unmarshalErr) + } + + if tok.AccessToken == "" { + return nil, errors.New("refresh endpoint returned no access_token") + } + + return &tok, nil +} diff --git a/internal/auth/oauth_test.go b/internal/auth/oauth_test.go index c0534d05c..3acd39831 100644 --- a/internal/auth/oauth_test.go +++ b/internal/auth/oauth_test.go @@ -114,9 +114,9 @@ func TestDiscoverOAuth_Success(t *testing.T) { assert.Equal(t, discoveryPath, r.URL.Path) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ - "issuer":"http://hydra", - "authorization_endpoint":"http://hydra/oauth2/auth", - "token_endpoint":"http://hydra/oauth2/token" + "issuer":"https://as.example", + "authorization_endpoint":"https://as.example/oauth2/auth", + "token_endpoint":"https://as.example/oauth2/token" }`)) })) defer srv.Close() @@ -124,16 +124,16 @@ func TestDiscoverOAuth_Success(t *testing.T) { meta, err := DiscoverOAuth(context.Background(), srv.URL) require.NoError(t, err) - assert.Equal(t, "http://hydra/oauth2/auth", meta.AuthorizationEndpoint) - assert.Equal(t, "http://hydra/oauth2/token", meta.TokenEndpoint) + assert.Equal(t, "https://as.example/oauth2/auth", meta.AuthorizationEndpoint) + assert.Equal(t, "https://as.example/oauth2/token", meta.TokenEndpoint) } // The authorize URL must carry PKCE and state, and the redirect URI must name // the port actually bound or the code comes back to nothing. func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", - TokenEndpoint: "http://hydra/oauth2/token", + AuthorizationEndpoint: "https://as.example/oauth2/auth", + TokenEndpoint: "https://as.example/oauth2/token", }) u, err := url.Parse(flow.AuthURL()) @@ -152,14 +152,14 @@ func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { } // Regression: listener.Addr() resolves "localhost" to 127.0.0.1, and -// authorization servers compare redirect_uri as an exact string. Hydra -// registers http://localhost:51164/, so emitting the IP literal gets the -// authorize request rejected before the user ever sees a login page. The +// 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", &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", - TokenEndpoint: "http://hydra/oauth2/token", + AuthorizationEndpoint: "https://as.example/oauth2/auth", + TokenEndpoint: "https://as.example/oauth2/token", }) require.NoError(t, err) @@ -177,10 +177,10 @@ func TestOAuthFlow_RedirectURIKeepsConfiguredHostname(t *testing.T) { } func TestOAuthFlow_ExchangesCodeForToken(t *testing.T) { - srv := tokenServer(t, http.StatusOK, `{"access_token":"hydra-token","token_type":"bearer","refresh_token":"r"}`) + srv := tokenServer(t, http.StatusOK, `{"access_token":"issued-token","token_type":"bearer","refresh_token":"r"}`) flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", + AuthorizationEndpoint: "https://as.example/oauth2/auth", TokenEndpoint: srv.URL, }) @@ -194,7 +194,7 @@ func TestOAuthFlow_ExchangesCodeForToken(t *testing.T) { token, err := flow.Wait(context.Background()) require.NoError(t, err) - assert.Equal(t, "hydra-token", token, "Wait must return the exchanged access token, not the code") + assert.Equal(t, "issued-token", token, "Wait must return the exchanged access token, not the code") } // A callback this login did not start must be refused: the listener is open to @@ -203,7 +203,7 @@ func TestOAuthFlow_RejectsStateMismatch(t *testing.T) { srv := tokenServer(t, http.StatusOK, `{"access_token":"should-never-be-used"}`) flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", + AuthorizationEndpoint: "https://as.example/oauth2/auth", TokenEndpoint: srv.URL, }) @@ -228,7 +228,7 @@ func TestOAuthFlow_SurfacesTokenEndpointRejection(t *testing.T) { `{"error":"invalid_grant","error_description":"code already used"}`) flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", + AuthorizationEndpoint: "https://as.example/oauth2/auth", TokenEndpoint: srv.URL, }) @@ -250,8 +250,8 @@ func TestOAuthFlow_SurfacesTokenEndpointRejection(t *testing.T) { // user declining consent. That must not wait out the five-minute timeout. func TestOAuthFlow_SurfacesAuthorizationError(t *testing.T) { flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", - TokenEndpoint: "http://hydra/oauth2/token", + AuthorizationEndpoint: "https://as.example/oauth2/auth", + TokenEndpoint: "https://as.example/oauth2/token", }) go func() { @@ -275,7 +275,7 @@ func TestOAuthFlow_CodeCallbackIsNotSwallowedBySentinel(t *testing.T) { srv := tokenServer(t, http.StatusOK, `{"access_token":"survived"}`) flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", + AuthorizationEndpoint: "https://as.example/oauth2/auth", TokenEndpoint: srv.URL, }) @@ -297,8 +297,8 @@ func TestOAuthFlow_CodeCallbackIsNotSwallowedBySentinel(t *testing.T) { // never be asked to release the port. func TestOAuthFlow_KeylessCallbackStillInterrupts(t *testing.T) { flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "http://hydra/oauth2/auth", - TokenEndpoint: "http://hydra/oauth2/token", + AuthorizationEndpoint: "https://as.example/oauth2/auth", + TokenEndpoint: "https://as.example/oauth2/token", }) go func() { diff --git a/internal/auth/refresh_test.go b/internal/auth/refresh_test.go new file mode 100644 index 000000000..5aa01dce7 --- /dev/null +++ b/internal/auth/refresh_test.go @@ -0,0 +1,184 @@ +// 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/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/datarobot/cli/internal/config" + "github.com/datarobot/cli/internal/config/viperx" +) + +// refreshServer records the form it was sent and answers with the given status +// and body. +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) { + if err := r.ParseForm(); err != nil { + // Not require/assert: this runs on the server goroutine, where a + // failed assertion cannot fail the test cleanly. Surface it as a + // response the test will notice instead. + w.WriteHeader(http.StatusBadRequest) + + return + } + + 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 is not a rejection", nil, false}, + {"401 is a rejection", &config.HTTPStatusError{StatusCode: 401}, true}, + {"403 is a rejection", &config.HTTPStatusError{StatusCode: 403}, true}, + {"404 is unjudged", &config.HTTPStatusError{StatusCode: 404}, false}, + {"429 is unjudged", &config.HTTPStatusError{StatusCode: 429}, false}, + {"500 is unjudged", &config.HTTPStatusError{StatusCode: 500}, false}, + {"503 is unjudged", &config.HTTPStatusError{StatusCode: 503}, false}, + {"a transport error is unjudged", assert.AnError, false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, tokenWasRejected(tc.err)) + }) + } +} + +// Clearing the access token alone would leave a working refresh token on disk. +func TestClearOAuthState(t *testing.T) { + _, cleanup := setupTestEnvironment(t) + defer cleanup() + + StoreOAuthState("stored-refresh", "https://as.example/oauth2/token") + require.NotEmpty(t, viperx.GetString(config.OAuthRefreshToken)) + + ClearOAuthState() + + assert.Empty(t, viperx.GetString(config.OAuthRefreshToken)) + assert.Empty(t, viperx.GetString(config.OAuthTokenEndpoint)) +} + +// Both keys must be on the persist allowlist, or a renewal survives only until +// the process exits and every new command starts with a browser. +func TestOAuthKeysArePersistable(t *testing.T) { + for _, key := range []string{config.OAuthRefreshToken, config.OAuthTokenEndpoint} { + _, ok := config.PersistableKeys[key] + assert.True(t, ok, "%s must be persistable or renewal state never reaches disk", key) + } +} 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 From 395694cb94f07fe99390ea0e81b1549ce4fa658c Mon Sep 17 00:00:00 2001 From: Erik Lattimore Date: Thu, 3 Sep 2026 12:47:55 -0400 Subject: [PATCH 3/4] =?UTF-8?q?refactor(auth):=20shrink=20the=20change=20?= =?UTF-8?q?=E2=80=94=20dedupe,=20tighten=20comments,=20cut=20thin=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same behavior, ~270 fewer lines, to keep the diff reviewable. Real duplication removed rather than lines golfed: - One postToken for both grants; exchangeCode and postRefresh differed only in the form they send and what rejection means to the caller. - One newFlow builds the flow and its callback server; the two constructors differed only in the URL and the OAuth fields. - Dropped tokenResponse.TokenType and .ExpiresIn, declared and never read. - Merged the refresh tests into oauth_test.go and shared one driveFlow helper, which removes the goroutine boilerplate repeated in five tests. Tests cut as unlikely to fail or covered elsewhere: - OAuthKeysArePersistable and ClearOAuthState asserted a map literal and a two-line setter. - SurfacesFailures: the errCh-not-sentinel property is asserted by RejectsStateMismatch; error_description propagation is formatting. - KeylessCallbackStillInterrupts: paired with CodeCallbackIsNotSwallowedBySentinel, which covers the reordering from the direction that regressed. - OAuthRequested: its load-bearing case is asserted more meaningfully by GateOffIssuesNoDiscovery, which checks that no network call is made rather than a boolean. - Trimmed the status-code table to one case per branch. Kept every test that guards a bug actually hit or a property that would fail silently: the sentinel ordering, the redirect-URI hostname, strict discovery, and the judged-vs-unjudged refresh guard. --- internal/auth/auth.go | 38 ++-- internal/auth/browserflow.go | 138 +++++--------- internal/auth/oauth.go | 207 +++++++------------- internal/auth/oauth_test.go | 349 +++++++++++++++++----------------- internal/auth/refresh_test.go | 184 ------------------ 5 files changed, 310 insertions(+), 606 deletions(-) delete mode 100644 internal/auth/refresh_test.go diff --git a/internal/auth/auth.go b/internal/auth/auth.go index 57d030f62..acd57557c 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -371,16 +371,11 @@ func EnsureAuthenticated(ctx context.Context) bool { //nolint: cyclop return true } - // A stored token that was actually REJECTED may be renewable without a - // browser. Try that before the interactive flow: an authorization server - // with short access-token lifetimes would otherwise send the user through - // Okta every few minutes, which is the whole reason `offline_access` is - // requested at login. - // - // Only on a judged rejection. An unjudged failure (404, 5xx, a network - // blip) is not evidence the token expired, and burning a - // rotation-on-use refresh token against a server that never rejected - // anything would turn a transient outage into a forced re-login. + // 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 } @@ -663,11 +658,9 @@ func GetBaseURLOrAsk() string { return datarobotHost } -// renewStoredToken tries to swap an expired access token for a fresh one using -// the stored refresh token, and reports whether the profile ends up usable. -// -// false means "carry on to the interactive login" — including the ordinary case -// of a profile that has nothing to renew with. +// 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) { @@ -677,9 +670,8 @@ func renewStoredToken(ctx context.Context) bool { return false } - // Verify rather than trust: the renewed credential has to satisfy the same - // check the stored one just failed, or a server handing back a token it - // will not accept would loop us silently. + // 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() @@ -698,13 +690,9 @@ func renewStoredToken(ctx context.Context) bool { return true } -// tokenWasRejected reports whether a verification failure was the server -// judging the credential (401/403) rather than being unable to answer. -// -// The distinction is the same one fprintServerStatus makes: a 404, 429, 5xx or -// transport error says nothing about the token, and treating it as expiry would -// spend a refresh token — possibly a rotate-on-use one — on a server that never -// rejected anything. +// 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 diff --git a/internal/auth/browserflow.go b/internal/auth/browserflow.go index 677e1a648..73e0a38d4 100644 --- a/internal/auth/browserflow.go +++ b/internal/auth/browserflow.go @@ -67,17 +67,15 @@ type BrowserFlow struct { timeout time.Duration // OAuth mode only (nil for the legacy ?key= hand-off). When set, the - // callback carries an authorization code that handleCallback exchanges for - // a token before publishing it on keyCh, so Wait's contract — "returns the - // credential" — is identical in both modes and callers need not care. + // 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 a failure that happened inside the callback handler, e.g. a - // state mismatch or a rejected token exchange. 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 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 @@ -89,27 +87,20 @@ type BrowserFlow struct { } // NewBrowserFlow binds the callback listener and prepares the browser login for -// datarobotHost. The caller must Close the returned flow. -// -// Prefer NewBrowserFlowContext where a context is available; this exists for -// callers that have none, and only matters when the OAuth flow is enabled, where -// it bounds the discovery probe on its own. +// datarobotHost. The caller must Close the returned flow. Prefer +// NewBrowserFlowContext where a context is available. func NewBrowserFlow(datarobotHost string) (*BrowserFlow, error) { return NewBrowserFlowContext(context.Background(), datarobotHost, nil) } -// NewBrowserFlowContext binds the callback listener, choosing between the legacy -// `?key=` hand-off and OAuth2 authorization-code + PKCE. -// -// oauthOverride comes from an explicit --oauth/--no-oauth flag and wins; nil -// defers to DATAROBOT_OAUTH_ENABLED, which is off by default. So unless someone -// opts in, this binds exactly the listener it always has and issues no discovery -// request at all — which is what keeps deployments that serve a discovery -// document without a login service behind it working as before. +// NewBrowserFlowContext binds the callback listener, choosing between the +// legacy `?key=` hand-off and PKCE. // -// When OAuth IS requested and discovery does not produce a usable document this -// returns ErrOAuthNotSupported rather than falling back. Falling back would hand -// the user a different kind of credential while looking like success. +// 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) @@ -133,8 +124,14 @@ 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), @@ -150,16 +147,12 @@ func newBrowserFlowOn(addr, datarobotHost string) (*BrowserFlow, error) { ReadHeaderTimeout: 10 * time.Second, } - return flow, nil + return flow } -// newBrowserFlowOAuthOn binds the callback listener for an OAuth2 -// authorization-code + PKCE login against the server described by meta. -// -// The redirect URI is the loopback listener itself, which is what lets the code -// come back to this process; RFC 8252 blesses exactly this shape for native -// apps. The port is still CallbackAddr's, so the authorization server must have -// http://localhost:51164/ registered as a redirect URI for OAuthClientID. +// 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 { @@ -171,13 +164,10 @@ func newBrowserFlowOAuthOn(addr string, meta *OAuthMetadata) (*BrowserFlow, erro return nil, err } - // The redirect URI must keep addr's HOSTNAME and take the listener's PORT. - // - // Both halves matter. listener.Addr() resolves "localhost" to 127.0.0.1, and - // authorization servers match redirect_uri as an exact string — a server - // that registered http://localhost:51164/ rejects the IP literal outright. - // The port has to come from the listener because tests bind :0 and the code - // must come back to the port actually bound. + // 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() @@ -201,25 +191,10 @@ func newBrowserFlowOAuthOn(addr string, meta *OAuthMetadata) (*BrowserFlow, erro return nil, err } - flow := &BrowserFlow{ - authURL: authURL, - listener: listener, - keyCh: make(chan string, 1), - errCh: make(chan error, 1), - timeout: DefaultLoginTimeout, - oauthMeta: meta, - oauthPKCE: p, - redirectURI: redirectURI, - } - - mux := http.NewServeMux() - mux.HandleFunc("/", flow.handleCallback) - - flow.server = &http.Server{ - Addr: addr, - Handler: mux, - ReadHeaderTimeout: 10 * time.Second, - } + flow := newFlow(addr, listener, authURL) + flow.oauthMeta = meta + flow.oauthPKCE = p + flow.redirectURI = redirectURI return flow, nil } @@ -229,14 +204,13 @@ func (f *BrowserFlow) AuthURL() string { return f.authURL } -// RefreshToken is the refresh token from the OAuth exchange, or "" when the -// login took the legacy path or the server issued none. Only meaningful after -// Wait has returned successfully. +// 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 should be sent, or "" in legacy mode. +// TokenEndpoint is where a renewal is sent, or "" in legacy mode. func (f *BrowserFlow) TokenEndpoint() string { if f.oauthMeta == nil { return "" @@ -320,17 +294,13 @@ func (f *BrowserFlow) Close() error { return f.closeErr } -// handleCallback receives the redirect that ends the browser login. +// 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. // -// Two shapes arrive here. The legacy DataRobot web app sends the credential -// outright as `?key=`. An OAuth2 authorization server sends -// `?code=&state=`, which this exchanges for a token before -// publishing it, so Wait returns a usable credential either way. -// -// Order matters: the OAuth branch is checked FIRST. A keyless request is the -// port-reclaim interrupt sentinel (see Wait and listenReclaimingPort), so an -// OAuth callback falling through to that check would read as "another CLI wants -// this port" and abort the login instead of completing it. +// 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) { query := r.URL.Query() @@ -368,11 +338,11 @@ func (f *BrowserFlow) handleCallback(w http.ResponseWriter, r *http.Request) { } } -// handleOAuthCallback validates the redirect and performs the code exchange. +// 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, generated seconds ago, - // and an attacker 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. + // 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")) @@ -386,9 +356,7 @@ func (f *BrowserFlow) handleOAuthCallback(w http.ResponseWriter, r *http.Request return } - // Kept for the caller to persist after Wait returns. Without a refresh - // token the credential simply expires and the user logs in again — the - // same behavior as the legacy flow. + // For the caller to persist after Wait. f.refreshToken = tok.RefreshToken if tok.RefreshToken == "" { @@ -409,9 +377,8 @@ func (f *BrowserFlow) handleOAuthCallback(w http.ResponseWriter, r *http.Request } // failCallback tells the browser the login failed and hands the reason to Wait. -// -// It deliberately does NOT publish an empty string on keyCh: that value means -// "release the port" and would turn a real failure into a silent interruption. +// 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) @@ -523,10 +490,9 @@ func runLoginWithFlow(ctx context.Context, flow *BrowserFlow, opts LoginOptions) return "", err } - // Record what a later renewal needs — or clear any stale material when this - // login took the legacy path, so a refresh token from a previous OAuth - // login is never renewed against a different instance. Persisting is the - // caller's job; both call sites write the config immediately after. + // 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 { diff --git a/internal/auth/oauth.go b/internal/auth/oauth.go index 2b3a484c7..009505b68 100644 --- a/internal/auth/oauth.go +++ b/internal/auth/oauth.go @@ -1,21 +1,10 @@ package auth -// OAuth2 authorization-code + PKCE login, for DataRobot deployments that front -// their own OAuth2 authorization server rather than the SaaS -// `/account/developer-tools?cliRedirect=true` hand-off. -// -// Why this exists alongside the legacy flow rather than replacing it: the SaaS -// path returns the credential as a `?key=` query parameter, which means -// the token itself travels in a URL and lands in browser history. Here the CLI -// holds a PKCE verifier, receives only an authorization code on the loopback -// redirect, and exchanges it for the token over POST — so the token arrives in a -// response body. It also gives the callback a real `state` to check, which the -// legacy path has no way to provide. -// -// The legacy path remains the DEFAULT. This one activates only when explicitly -// asked for, because some deployments serve an OIDC discovery document without -// having a login service behind it: the document's presence is not evidence the -// flow is supported, so it cannot be used as an auto-detect signal. +// 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" @@ -37,41 +26,29 @@ import ( "github.com/datarobot/cli/internal/log" ) -// OAuthEnabledEnv gates the discovery lookup. Unset (the default) means the CLI -// behaves exactly as it always has and issues no discovery request at all. +// 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 OAuth2 client this CLI presents. Public means no -// client secret — there is nowhere safe to keep one on a user's machine — so the -// authorization server is expected to require PKCE for it. +// 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" -// discoveryPath is the RFC 8414 / OpenID Connect Discovery location, resolved -// against the configured endpoint's scheme+host. const discoveryPath = "/.well-known/openid-configuration" -// oauthScopes are requested on every authorization. -// -// `offline_access` is requested deliberately: deployments may issue short-lived -// access tokens, and without a refresh token the user would have to re-run -// `dr auth login` every time one expired. An authorization server that does not -// grant it simply returns no refresh token, which degrades to the same -// re-login-on-expiry behavior the legacy flow already has. +// `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"} -// discoveryTimeout bounds the probe. It runs before anything visible happens, so -// it must not be something a user waits on when the endpoint does not implement -// discovery at all. +// The probe runs before anything visible happens, so it must stay short for +// hosts that do not implement discovery. const discoveryTimeout = 5 * time.Second -// tokenExchangeTimeout bounds the code-for-token POST. This one runs inside the -// callback HTTP handler, with the user watching a browser tab, so it is short. const tokenExchangeTimeout = 30 * time.Second -// ErrOAuthNotSupported means the endpoint did not serve a usable discovery -// document. It is returned rather than silently falling back to the legacy flow: -// the user explicitly asked for OAuth, and quietly handing them a different kind -// of credential instead would look like success. +// 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. @@ -81,10 +58,8 @@ type OAuthMetadata struct { TokenEndpoint string `json:"token_endpoint"` } -// OAuthRequested reports whether the caller asked for the OAuth flow. -// -// override comes from an explicit --oauth/--no-oauth flag and wins outright; -// nil means "not specified", in which case the environment decides. Default off. +// 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 @@ -98,12 +73,10 @@ func OAuthRequested(override *bool) bool { } } -// DiscoverOAuth fetches and validates the discovery document for datarobotHost. -// -// It is strict on purpose. A single-page app that serves its shell for unknown -// paths answers 200 with HTML for this URL, which would otherwise look like a -// successful discovery; requiring parseable JSON carrying both endpoints is what -// separates a real authorization server from a catch-all route. +// 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, "/") @@ -147,14 +120,13 @@ func DiscoverOAuth(ctx context.Context, datarobotHost string) (*OAuthMetadata, e return &meta, nil } -// pkce carries one authorization attempt's PKCE verifier and CSRF state. +// pkce is one attempt's verifier and CSRF state. type pkce struct { verifier string challenge string state string } -// newPKCE generates a fresh verifier/challenge/state triple. func newPKCE() (*pkce, error) { verifier, err := randomURLSafe(32) if err != nil { @@ -175,7 +147,6 @@ func newPKCE() (*pkce, error) { }, nil } -// randomURLSafe returns n bytes of crypto-random data, base64url encoded. func randomURLSafe(n int) (string, error) { buf := make([]byte, n) if _, err := rand.Read(buf); err != nil { @@ -185,7 +156,6 @@ func randomURLSafe(n int) (string, error) { return base64.RawURLEncoding.EncodeToString(buf), nil } -// authorizeURL builds the browser-facing authorization request. func authorizeURL(meta *OAuthMetadata, p *pkce, redirectURI string) (string, error) { u, err := url.Parse(meta.AuthorizationEndpoint) if err != nil { @@ -205,87 +175,37 @@ func authorizeURL(meta *OAuthMetadata, p *pkce, redirectURI string) (string, err return u.String(), nil } -// tokenResponse is the subset of the token endpoint's reply we consume. type tokenResponse struct { AccessToken string `json:"access_token"` - TokenType string `json:"token_type"` - ExpiresIn int `json:"expires_in"` 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 does not have, and the reason the token never -// appears in a URL: it comes back in this POST's response body. +// 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("client_id", OAuthClientID) form.Set("code_verifier", p.verifier) - ctx, cancel := context.WithTimeout(ctx, tokenExchangeTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, meta.TokenEndpoint, strings.NewReader(form.Encode())) + tok, err := postToken(ctx, meta.TokenEndpoint, form) 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("exchanging authorization code at %s: %w", meta.TokenEndpoint, 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 before checking the status: OAuth2 error replies are JSON too, and - // their `error_description` is far more useful than a bare status code. - if err := json.Unmarshal(body, &tok); err != nil && resp.StatusCode == http.StatusOK { - return nil, fmt.Errorf("token endpoint returned unparseable JSON: %w", err) + return nil, fmt.Errorf("token exchange failed: %w", err) } - if resp.StatusCode != http.StatusOK { - if tok.Error != "" { - return nil, fmt.Errorf("token exchange rejected: %s: %s", tok.Error, tok.ErrorDescription) - } - - return nil, fmt.Errorf("token exchange failed with HTTP %d", resp.StatusCode) - } - - if tok.AccessToken == "" { - return nil, errors.New("token endpoint returned no access_token") - } - - return &tok, nil + return tok, nil } -// ErrNoRefreshToken means this profile has nothing to renew with — either it -// was authenticated by the legacy hand-off, or the authorization server did not -// issue a refresh token. +// 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. -// -// This is the whole point of asking for `offline_access`: an authorization -// server that issues short-lived access tokens would otherwise force a full -// interactive login every time one expired. -// -// It writes the new credential into viper but does NOT persist it — the caller -// owns writing drconfig.yaml, matching how the login flow behaves. Returns -// ErrNoRefreshToken when the profile has nothing to renew with, which callers -// should treat as "fall back to interactive login", not as a failure. +// 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) @@ -301,9 +221,8 @@ func RefreshAccessToken(ctx context.Context) (string, error) { viperx.Set(config.DataRobotAPIKey, tok.AccessToken) - // Servers that rotate refresh tokens invalidate the old one on use, so a - // rotated value MUST replace what is stored or the next renewal fails. - // Servers that do not rotate simply omit it, and the stored one stays good. + // 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) } @@ -311,32 +230,46 @@ func RefreshAccessToken(ctx context.Context) (string, error) { return tok.AccessToken, nil } -// StoreOAuthState records what a successful OAuth login needs for later -// renewal. Persisting is the caller's job. +// 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 later be renewed against whatever -// instance issued it, which is not necessarily the one now configured. +// 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 and returns the parsed response. +// postRefresh performs the refresh_token grant. // -// A non-200 clears the stored material before returning: a rejected refresh -// token is spent — revoked, expired, or already rotated — so keeping it would -// make every later command retry something that cannot work. +// 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) @@ -344,7 +277,7 @@ func postRefresh(ctx context.Context, endpoint, refresh string) (*tokenResponse, req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode())) if err != nil { - return nil, fmt.Errorf("building refresh request: %w", err) + return nil, fmt.Errorf("building token request: %w", err) } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") @@ -352,37 +285,35 @@ func postRefresh(ctx context.Context, endpoint, refresh string) (*tokenResponse, resp, err := http.DefaultClient.Do(req) if err != nil { - return nil, fmt.Errorf("refreshing access token at %s: %w", endpoint, err) + 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 refresh response: %w", err) + return nil, fmt.Errorf("reading token response: %w", err) } var tok tokenResponse - // Decode before checking the status: OAuth2 error replies are JSON too, and - // their error_description is the actionable half. + // Decode first: error replies are JSON too, and error_description is the + // actionable half. unmarshalErr := json.Unmarshal(body, &tok) if resp.StatusCode != http.StatusOK { - ClearOAuthState() - if tok.Error != "" { - return nil, fmt.Errorf("refresh rejected: %s: %s", tok.Error, tok.ErrorDescription) + return &tok, fmt.Errorf("%s: %s", tok.Error, tok.ErrorDescription) } - return nil, fmt.Errorf("refresh failed with HTTP %d", resp.StatusCode) + return &tok, fmt.Errorf("HTTP %d", resp.StatusCode) } if unmarshalErr != nil { - return nil, fmt.Errorf("refresh endpoint returned unparseable JSON: %w", unmarshalErr) + return nil, fmt.Errorf("token endpoint returned unparseable JSON: %w", unmarshalErr) } if tok.AccessToken == "" { - return nil, errors.New("refresh endpoint returned no access_token") + 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 index 3acd39831..4944b6087 100644 --- a/internal/auth/oauth_test.go +++ b/internal/auth/oauth_test.go @@ -24,8 +24,19 @@ import ( "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() @@ -53,88 +64,46 @@ func tokenServer(t *testing.T, status int, body string) *httptest.Server { return srv } -func TestOAuthRequested(t *testing.T) { - yes, no := true, false +// 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 string - env string - override *bool - want bool + name, contentType, body string + wantOK bool }{ - {"unset is off — the whole point of the gate", "", nil, false}, - {"env true", "true", nil, true}, - {"env 1", "1", nil, true}, - {"env garbage is off, not an error", "banana", nil, false}, - {"flag beats unset env", "", &yes, true}, - {"flag --no-oauth beats env true", "true", &no, false}, - {"flag --oauth beats env absent", "", &yes, true}, + {"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) { - t.Setenv(OAuthEnabledEnv, tc.env) - assert.Equal(t, tc.want, OAuthRequested(tc.override)) - }) - } -} - -// A single-page app answers 200 with its HTML shell for unknown paths. Treating -// that as a successful discovery would send the CLI into an OAuth flow against -// an endpoint that cannot complete one. -func TestDiscoverOAuth_RejectsHTMLCatchAll(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html") - _, _ = w.Write([]byte("app shell")) - })) - defer srv.Close() - - _, err := DiscoverOAuth(context.Background(), srv.URL) - - require.Error(t, err) - assert.ErrorIs(t, err, ErrOAuthNotSupported) -} + 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() -func TestDiscoverOAuth_RejectsIncompleteDocument(t *testing.T) { - // Valid JSON, but missing token_endpoint — unusable. - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"issuer":"http://x","authorization_endpoint":"http://x/auth"}`)) - })) - defer srv.Close() + meta, err := DiscoverOAuth(context.Background(), srv.URL) - _, err := DiscoverOAuth(context.Background(), srv.URL) + if tc.wantOK { + require.NoError(t, err) + assert.Equal(t, "https://as.example/oauth2/token", meta.TokenEndpoint) - require.Error(t, err) - assert.ErrorIs(t, err, ErrOAuthNotSupported) -} + return + } -func TestDiscoverOAuth_Success(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - assert.Equal(t, discoveryPath, r.URL.Path) - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{ - "issuer":"https://as.example", - "authorization_endpoint":"https://as.example/oauth2/auth", - "token_endpoint":"https://as.example/oauth2/token" - }`)) - })) - defer srv.Close() - - meta, err := DiscoverOAuth(context.Background(), srv.URL) - - require.NoError(t, err) - assert.Equal(t, "https://as.example/oauth2/auth", meta.AuthorizationEndpoint) - assert.Equal(t, "https://as.example/oauth2/token", meta.TokenEndpoint) + require.Error(t, err) + assert.ErrorIs(t, err, ErrOAuthNotSupported) + }) + } } -// The authorize URL must carry PKCE and state, and the redirect URI must name -// the port actually bound or the code comes back to nothing. +// The authorize URL must carry PKCE and state, and name the bound port. func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { - flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: "https://as.example/oauth2/token", - }) + flow := newTestOAuthFlow(t, testMeta("https://as.example/oauth2/token")) u, err := url.Parse(flow.AuthURL()) require.NoError(t, err) @@ -157,10 +126,7 @@ func TestOAuthFlow_AuthURLCarriesPKCEAndState(t *testing.T) { // 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", &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: "https://as.example/oauth2/token", - }) + flow, err := newBrowserFlowOAuthOn("localhost:0", testMeta("https://as.example/oauth2/token")) require.NoError(t, err) t.Cleanup(func() { _ = flow.Close() }) @@ -179,42 +145,22 @@ func TestOAuthFlow_RedirectURIKeepsConfiguredHostname(t *testing.T) { func TestOAuthFlow_ExchangesCodeForToken(t *testing.T) { srv := tokenServer(t, http.StatusOK, `{"access_token":"issued-token","token_type":"bearer","refresh_token":"r"}`) - flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: srv.URL, - }) + flow := newTestOAuthFlow(t, testMeta(srv.URL)) - go func() { - resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) - if err == nil { - _ = resp.Body.Close() - } - }() - - token, err := flow.Wait(context.Background()) + 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") } -// A callback this login did not start must be refused: the listener is open to -// any local process, and state is the only thing distinguishing them. +// 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, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: srv.URL, - }) + flow := newTestOAuthFlow(t, testMeta(srv.URL)) - go func() { - resp, err := http.Get(callbackURL(t, flow, "?code=abc&state=not-the-right-state")) - if err == nil { - _ = resp.Body.Close() - } - }() - - token, err := flow.Wait(context.Background()) + token, err := driveFlow(t, flow, "?code=abc&state=not-the-right-state") require.Error(t, err) assert.Empty(t, token) @@ -223,98 +169,36 @@ func TestOAuthFlow_RejectsStateMismatch(t *testing.T) { "a rejected callback is a failure, not the port-reclaim interrupt") } -func TestOAuthFlow_SurfacesTokenEndpointRejection(t *testing.T) { - srv := tokenServer(t, http.StatusBadRequest, - `{"error":"invalid_grant","error_description":"code already used"}`) - - flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: srv.URL, - }) - - go func() { - resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) - if err == nil { - _ = resp.Body.Close() - } - }() - - _, err := flow.Wait(context.Background()) - - require.Error(t, err) - assert.Contains(t, err.Error(), "invalid_grant") - assert.Contains(t, err.Error(), "code already used", "the description is the actionable half") -} - -// The authorization server can report failure on the redirect itself, e.g. the -// user declining consent. That must not wait out the five-minute timeout. -func TestOAuthFlow_SurfacesAuthorizationError(t *testing.T) { - flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: "https://as.example/oauth2/token", - }) +// 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, "?error=access_denied&error_description=user+declined")) + resp, err := http.Get(callbackURL(t, flow, query)) if err == nil { _ = resp.Body.Close() } }() - _, err := flow.Wait(context.Background()) - - require.Error(t, err) - assert.Contains(t, err.Error(), "access_denied") - assert.NotErrorIs(t, err, ErrLoginInterrupted) + return flow.Wait(context.Background()) } -// THE regression that matters. An empty `key` is the port-reclaim sentinel, so -// the OAuth branch has to be evaluated first — otherwise every OAuth callback -// reads as "another CLI wants this port" and aborts the login. +// 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, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: srv.URL, - }) + flow := newTestOAuthFlow(t, testMeta(srv.URL)) // No `key` parameter at all — exactly the shape the sentinel looks for. - go func() { - resp, err := http.Get(callbackURL(t, flow, "?code=abc&state="+flow.oauthPKCE.state)) - if err == nil { - _ = resp.Body.Close() - } - }() - - token, err := flow.Wait(context.Background()) + 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) } -// And the sentinel itself must still work in OAuth mode, or a stale process can -// never be asked to release the port. -func TestOAuthFlow_KeylessCallbackStillInterrupts(t *testing.T) { - flow := newTestOAuthFlow(t, &OAuthMetadata{ - AuthorizationEndpoint: "https://as.example/oauth2/auth", - TokenEndpoint: "https://as.example/oauth2/token", - }) - - go func() { - resp, err := http.Get(callbackURL(t, flow, "")) - if err == nil { - _ = resp.Body.Close() - } - }() - - _, err := flow.Wait(context.Background()) - - assert.ErrorIs(t, err, ErrLoginInterrupted) -} - -// With the gate off, the constructor must not reach out at all — this is what -// protects deployments that serve a discovery document with no login service. +// 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, "") @@ -337,8 +221,7 @@ func TestNewBrowserFlowContext_GateOffIssuesNoDiscovery(t *testing.T) { assert.Contains(t, flow.AuthURL(), "/account/developer-tools?cliRedirect=true") } -// Gate on but discovery unusable must fail loudly. A silent downgrade to the -// legacy flow would look like success while producing a different credential. +// Gate on but undiscoverable must fail loudly, not downgrade silently. func TestNewBrowserFlowContext_GateOnWithoutDiscoveryFails(t *testing.T) { t.Setenv(OAuthEnabledEnv, "true") @@ -352,3 +235,123 @@ func TestNewBrowserFlowContext_GateOnWithoutDiscoveryFails(t *testing.T) { 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.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/auth/refresh_test.go b/internal/auth/refresh_test.go deleted file mode 100644 index 5aa01dce7..000000000 --- a/internal/auth/refresh_test.go +++ /dev/null @@ -1,184 +0,0 @@ -// 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/http" - "net/http/httptest" - "net/url" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/datarobot/cli/internal/config" - "github.com/datarobot/cli/internal/config/viperx" -) - -// refreshServer records the form it was sent and answers with the given status -// and body. -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) { - if err := r.ParseForm(); err != nil { - // Not require/assert: this runs on the server goroutine, where a - // failed assertion cannot fail the test cleanly. Surface it as a - // response the test will notice instead. - w.WriteHeader(http.StatusBadRequest) - - return - } - - 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 is not a rejection", nil, false}, - {"401 is a rejection", &config.HTTPStatusError{StatusCode: 401}, true}, - {"403 is a rejection", &config.HTTPStatusError{StatusCode: 403}, true}, - {"404 is unjudged", &config.HTTPStatusError{StatusCode: 404}, false}, - {"429 is unjudged", &config.HTTPStatusError{StatusCode: 429}, false}, - {"500 is unjudged", &config.HTTPStatusError{StatusCode: 500}, false}, - {"503 is unjudged", &config.HTTPStatusError{StatusCode: 503}, false}, - {"a transport error is unjudged", assert.AnError, false}, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - assert.Equal(t, tc.want, tokenWasRejected(tc.err)) - }) - } -} - -// Clearing the access token alone would leave a working refresh token on disk. -func TestClearOAuthState(t *testing.T) { - _, cleanup := setupTestEnvironment(t) - defer cleanup() - - StoreOAuthState("stored-refresh", "https://as.example/oauth2/token") - require.NotEmpty(t, viperx.GetString(config.OAuthRefreshToken)) - - ClearOAuthState() - - assert.Empty(t, viperx.GetString(config.OAuthRefreshToken)) - assert.Empty(t, viperx.GetString(config.OAuthTokenEndpoint)) -} - -// Both keys must be on the persist allowlist, or a renewal survives only until -// the process exits and every new command starts with a browser. -func TestOAuthKeysArePersistable(t *testing.T) { - for _, key := range []string{config.OAuthRefreshToken, config.OAuthTokenEndpoint} { - _, ok := config.PersistableKeys[key] - assert.True(t, ok, "%s must be persistable or renewal state never reaches disk", key) - } -} From 87f0868c7be58b71f715e6f12fd1cab70debae96 Mon Sep 17 00:00:00 2001 From: Erik Lattimore Date: Thu, 3 Sep 2026 13:05:28 -0400 Subject: [PATCH 4/4] fix(auth): refuse injected keys in OAuth mode; register --no-oauth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four issues from CI and review on #892. In OAuth mode handleCallback intercepted `code` and `error` and then fell through to the legacy `key` handler, so any local process could hit the loopback listener with ?key= during the login window and have that value stored as the credential — bypassing the `state` check the flow exists to provide. A key cannot authenticate this flow, so it is now refused with a 400 and published on neither channel: the login keeps waiting for the real callback, making it neither an injection nor a way to cancel someone else's login. Regression test covers both halves — the injected callback refused, the genuine one still completing. --no-oauth did not exist. The help text and a comment claimed cobra synthesises --no- variants from a bool flag; pflag does not, so the documented way to force the legacy hand-off failed as an unknown flag. Registered explicitly and marked mutually exclusive with --oauth. CI: internal/auth/oauth.go was missing the Apache header (Copyrights), and the test stub called ParseForm without bounding the body (gosec G120). Both fixed, and verified with the repo's own `task lint` and `task copyright` rather than a hand-picked package subset — linting one GOOS over selected packages is what let the gosec finding through. --- cmd/auth/login/cmd.go | 17 +++++++++---- internal/auth/browserflow.go | 15 ++++++++++++ internal/auth/oauth.go | 14 +++++++++++ internal/auth/oauth_test.go | 46 ++++++++++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 4 deletions(-) diff --git a/cmd/auth/login/cmd.go b/cmd/auth/login/cmd.go index fe1675fb8..59b3de56e 100644 --- a/cmd/auth/login/cmd.go +++ b/cmd/auth/login/cmd.go @@ -82,9 +82,13 @@ func RunE(cmd *cobra.Command, args []string) error { //nolint: cyclop // DATAROBOT_OAUTH_ENABLED decides, and its default is off. var oauthOverride *bool - if cmd.Flags().Changed("oauth") { + 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{ @@ -145,10 +149,15 @@ URL. Set DATAROBOT_OAUTH_ENABLED=true to make that the default for a shell; // 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. Registered as - // one boolean so cobra gives us --oauth and --no-oauth for free, and read - // via Flags().Changed so "unset" stays distinguishable from "--no-oauth". + // 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/internal/auth/browserflow.go b/internal/auth/browserflow.go index 73e0a38d4..f62487482 100644 --- a/internal/auth/browserflow.go +++ b/internal/auth/browserflow.go @@ -319,6 +319,21 @@ func (f *BrowserFlow) handleCallback(w http.ResponseWriter, r *http.Request) { 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") diff --git a/internal/auth/oauth.go b/internal/auth/oauth.go index 009505b68..80d396d76 100644 --- a/internal/auth/oauth.go +++ b/internal/auth/oauth.go @@ -1,3 +1,17 @@ +// 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 diff --git a/internal/auth/oauth_test.go b/internal/auth/oauth_test.go index 4944b6087..e3a1fec24 100644 --- a/internal/auth/oauth_test.go +++ b/internal/auth/oauth_test.go @@ -21,6 +21,7 @@ import ( "net/http/httptest" "net/url" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -197,6 +198,50 @@ func TestOAuthFlow_CodeCallbackIsNotSwallowedBySentinel(t *testing.T) { 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) { @@ -247,6 +292,7 @@ func newRefreshServer(t *testing.T, status int, body string) *refreshServer { 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