From 0f60b9c550c037378b8c90bcdb539fe353317e2f Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Thu, 23 Jul 2026 10:12:40 +0530 Subject: [PATCH 1/2] fix: harden OAuth device and PKCE flows --- admin/auth_code_test.go | 47 +- admin/database/postgres/migrations/0096.sql | 5 + admin/server/auth/device_code.go | 56 ++- admin/server/auth/device_code_test.go | 320 ++++++++++++++ admin/server/auth/handlers.go | 3 + admin/server/auth/mcp_oauth.go | 93 +++- admin/server/auth/oauth_registration_test.go | 278 ++++++++++++ admin/server/auth/pkce.go | 126 ++++-- admin/server/auth/pkce_flow_test.go | 333 +++++++++++++++ admin/server/auth/refresh_token_test.go | 418 ++++++++++++++++++ cli/pkg/deviceauth/authenticator.go | 154 ++++--- cli/pkg/deviceauth/authenticator_test.go | 424 +++++++++++++++++++ cli/pkg/pkce/authenticator.go | 23 +- cli/pkg/pkce/authenticator_test.go | 206 ++++++++- 14 files changed, 2351 insertions(+), 135 deletions(-) create mode 100644 admin/database/postgres/migrations/0096.sql create mode 100644 admin/server/auth/device_code_test.go create mode 100644 admin/server/auth/oauth_registration_test.go create mode 100644 admin/server/auth/pkce_flow_test.go create mode 100644 admin/server/auth/refresh_token_test.go create mode 100644 cli/pkg/deviceauth/authenticator_test.go diff --git a/admin/auth_code_test.go b/admin/auth_code_test.go index 1d9d204558c1..dbd69c63dfc3 100644 --- a/admin/auth_code_test.go +++ b/admin/auth_code_test.go @@ -1,44 +1,31 @@ package admin import ( + "encoding/base64" + "strings" "testing" + "time" ) func TestGenerateDeviceAndUserCode(t *testing.T) { + // Check the externally visible code format and lifetime. Collision handling is + // covered deterministically at the database-backed issuance boundary. + started := time.Now() authCode, err := generateDeviceAndUserCode() if err != nil { - t.Error(err) + t.Fatal(err) } - if len(authCode.DeviceCode) != 32 { - t.Errorf("device code length is incorrect; got %d, want 32", len(authCode.DeviceCode)) + decoded, err := base64.StdEncoding.DecodeString(authCode.DeviceCode) + if err != nil { + t.Fatalf("device code is not valid base64: %v", err) } - if len(authCode.UserCode) != 8 { - t.Errorf("user code length is incorrect; got %d, want 8", len(authCode.UserCode)) + if len(decoded) != 24 { + t.Errorf("decoded device code length = %d, want 24", len(decoded)) } -} - -func TestGenerateDeviceAndUserCodeCollision(t *testing.T) { - deviceCodeSet := make(map[string]bool) - userCodeSet := make(map[string]bool) - - for i := 0; i < 10000; i++ { - authCode, err := generateDeviceAndUserCode() - if err != nil { - t.Error(err) - } - if len(authCode.DeviceCode) != 32 { - t.Errorf("device code length is incorrect; got %d, want 32", len(authCode.DeviceCode)) - } - if len(authCode.UserCode) != 8 { - t.Errorf("user code length is incorrect; got %d, want 8", len(authCode.UserCode)) - } - if deviceCodeSet[authCode.DeviceCode] { - t.Errorf("device code collision: %s", authCode.DeviceCode) - } - if userCodeSet[authCode.UserCode] { - t.Errorf("user code collision: %s", authCode.UserCode) - } - deviceCodeSet[authCode.DeviceCode] = true - userCodeSet[authCode.UserCode] = true + if len(authCode.UserCode) != 8 || strings.Trim(authCode.UserCode, "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ") != "" { + t.Errorf("user code %q is not eight uppercase base36 characters", authCode.UserCode) + } + if authCode.Expiry.Before(started.Add(DeviceAuthCodeTTL)) || authCode.Expiry.After(time.Now().Add(DeviceAuthCodeTTL)) { + t.Errorf("expiry %s is outside the issuance window", authCode.Expiry) } } diff --git a/admin/database/postgres/migrations/0096.sql b/admin/database/postgres/migrations/0096.sql new file mode 100644 index 000000000000..67a92e823eaf --- /dev/null +++ b/admin/database/postgres/migrations/0096.sql @@ -0,0 +1,5 @@ +-- User codes are lookup keys while a device request is pending. Keeping them +-- unique prevents a confirmation from approving an arbitrary colliding request. +CREATE UNIQUE INDEX device_auth_codes_pending_user_code_idx +ON device_auth_codes (user_code) +WHERE approval_state = 0; diff --git a/admin/server/auth/device_code.go b/admin/server/auth/device_code.go index 178af80d9214..b141a3aaf2df 100644 --- a/admin/server/auth/device_code.go +++ b/admin/server/auth/device_code.go @@ -15,6 +15,11 @@ import ( "github.com/rilldata/rill/admin/pkg/oauth" ) +const ( + deviceCodePollingInterval = 5 * time.Second + deviceCodeIssueAttempts = 3 +) + // DeviceCodeResponse encapsulates the response for obtaining a device code. type DeviceCodeResponse struct { DeviceCode string `json:"device_code"` @@ -58,16 +63,26 @@ func (a *Authenticator) handleDeviceCodeRequest(w http.ResponseWriter, r *http.R http.Error(w, "client_id is required", http.StatusBadRequest) return } - scopes := strings.Split(values.Get("scope"), " ") - if len(scopes) == 0 { + scope := strings.TrimSpace(values.Get("scope")) + if scope == "" { http.Error(w, "scope is required", http.StatusBadRequest) return } + scopes := strings.Fields(scope) if len(scopes) > 1 || scopes[0] != "full_account" { http.Error(w, "invalid scope", http.StatusBadRequest) return } - authCode, err := a.admin.IssueDeviceAuthCode(r.Context(), clientID) + + // A user code must identify one pending request. The database enforces that + // invariant, and retrying here turns a rare random-code collision into a new code. + var authCode *database.DeviceAuthCode + for range deviceCodeIssueAttempts { + authCode, err = a.admin.IssueDeviceAuthCode(r.Context(), clientID) + if err == nil || !errors.Is(err, database.ErrNotUnique) { + break + } + } if err != nil { internalServerError(w, fmt.Errorf("failed to issue auth code: %w", err)) return @@ -89,7 +104,9 @@ func (a *Authenticator) handleDeviceCodeRequest(w http.ResponseWriter, r *http.R VerificationURI: a.admin.URLs.AuthVerifyDeviceUI(nil), VerificationCompleteURI: a.admin.URLs.AuthVerifyDeviceUI(qry), ExpiresIn: int(admin.DeviceAuthCodeTTL.Seconds()), - PollingInterval: 5, + // This is the client's minimum polling cadence. Server-side slow_down is + // not enforced until poll history is persisted with the device code. + PollingInterval: int(deviceCodePollingInterval.Seconds()), } respBytes, err := json.Marshal(resp) @@ -174,7 +191,17 @@ func (a *Authenticator) getAccessTokenForDeviceCode(w http.ResponseWriter, r *ht } responseVersion := values.Get("token_response_version") - authCode, err := a.admin.DB.FindDeviceAuthCodeByDeviceCode(r.Context(), deviceCode) + // Token creation and one-time code consumption must commit together. Besides + // keeping failures retryable, the transaction ensures concurrent polls cannot + // persist more than one token for the same approved code. + txCtx, tx, err := a.admin.DB.NewTx(r.Context(), false) + if err != nil { + internalServerError(w, fmt.Errorf("failed to start device code transaction: %w", err)) + return + } + defer func() { _ = tx.Rollback() }() + + authCode, err := a.admin.DB.FindDeviceAuthCodeByDeviceCode(txCtx, deviceCode) if err != nil { if errors.Is(err, database.ErrNotFound) { http.Error(w, fmt.Sprintf("no such device code: %s found", deviceCode), http.StatusBadRequest) @@ -189,16 +216,22 @@ func (a *Authenticator) getAccessTokenForDeviceCode(w http.ResponseWriter, r *ht return } - if authCode.Expiry.Before(time.Now()) { + // Expiry is checked before approval state so an expired code can never issue a + // token, even if it was approved just before its deadline. + if !authCode.Expiry.After(time.Now()) { http.Error(w, "expired_token", http.StatusUnauthorized) return } if authCode.ApprovalState == database.DeviceAuthCodeStateRejected { - err = a.admin.DB.DeleteDeviceAuthCode(r.Context(), deviceCode) + err = a.admin.DB.DeleteDeviceAuthCode(txCtx, deviceCode) if err != nil { internalServerError(w, fmt.Errorf("failed to clean up rejected code: %s, %w", deviceCode, err)) return } + if err := tx.Commit(); err != nil { + internalServerError(w, fmt.Errorf("failed to commit rejected code cleanup: %w", err)) + return + } http.Error(w, "rejected", http.StatusUnauthorized) return } @@ -210,19 +243,22 @@ func (a *Authenticator) getAccessTokenForDeviceCode(w http.ResponseWriter, r *ht internalServerError(w, fmt.Errorf("inconsistent state, %w", errors.New("server error"))) return } - // TODO handle too many requests - authToken, err := a.admin.IssueUserAuthToken(r.Context(), *authCode.UserID, authCode.ClientID, "", nil, nil, false) + authToken, err := a.admin.IssueUserAuthToken(txCtx, *authCode.UserID, authCode.ClientID, "", nil, nil, false) if err != nil { internalServerError(w, fmt.Errorf("failed to issue access token, %w", err)) return } - err = a.admin.DB.DeleteDeviceAuthCode(r.Context(), deviceCode) + err = a.admin.DB.DeleteDeviceAuthCode(txCtx, deviceCode) if err != nil { internalServerError(w, fmt.Errorf("failed to clean up approved code: %s, %w", deviceCode, err)) return } + if err := tx.Commit(); err != nil { + internalServerError(w, fmt.Errorf("failed to commit approved code exchange: %w", err)) + return + } var respBytes []byte if responseVersion == "standard" { diff --git a/admin/server/auth/device_code_test.go b/admin/server/auth/device_code_test.go new file mode 100644 index 000000000000..444b99d5d242 --- /dev/null +++ b/admin/server/auth/device_code_test.go @@ -0,0 +1,320 @@ +package auth_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rilldata/rill/admin" + "github.com/rilldata/rill/admin/database" + serverauth "github.com/rilldata/rill/admin/server/auth" + "github.com/rilldata/rill/admin/testadmin" + "github.com/stretchr/testify/require" +) + +const oauth004DeviceGrantType = "urn:ietf:params:oauth:grant-type:device_code" + +var ( + oauth004Sequence atomic.Uint64 + errInjectedTokenIssue = errors.New("injected token issuance failure") + errInjectedCodeDelete = errors.New("injected device code deletion failure") +) + +type oauth004FormResult struct { + status int + body string + err error +} + +type oauth004CollisionDB struct { + database.DB + calls atomic.Int32 +} + +func (d *oauth004CollisionDB) InsertDeviceAuthCode(ctx context.Context, deviceCode, userCode, clientID string, expiresOn time.Time) (*database.DeviceAuthCode, error) { + if d.calls.Add(1) == 1 { + return nil, database.NewNotUniqueError("injected user code collision") + } + return d.DB.InsertDeviceAuthCode(ctx, deviceCode, userCode, clientID, expiresOn) +} + +type oauth004IssueFailureDB struct { + database.DB +} + +func (d *oauth004IssueFailureDB) InsertUserAuthToken(context.Context, *database.InsertUserAuthTokenOptions) (*database.UserAuthToken, error) { + return nil, errInjectedTokenIssue +} + +type oauth004DeleteFailureDB struct { + database.DB +} + +func (d *oauth004DeleteFailureDB) DeleteDeviceAuthCode(context.Context, string) error { + return errInjectedCodeDelete +} + +func TestOAUTH004DeviceCodeFlow(t *testing.T) { + // Exercise the full real-Postgres device-code lifecycle, including collision, concurrency, and rollback boundaries. + fix := testadmin.New(t) + deviceAuthorizationURL := fix.ExternalURL() + "/auth/oauth/device_authorization" + tokenURL := fix.ExternalURL() + "/auth/oauth/token" + + t.Run("requires the one supported scope and advertises polling", func(t *testing.T) { + res := oauth004PostForm(t.Context(), deviceAuthorizationURL, url.Values{ + "client_id": {database.AuthClientIDRillCLI}, + }) + require.NoError(t, res.err) + require.Equal(t, http.StatusBadRequest, res.status) + require.Equal(t, "scope is required\n", res.body) + + for _, scope := range []string{"openid", "full_account openid"} { + res = oauth004PostForm(t.Context(), deviceAuthorizationURL, url.Values{ + "client_id": {database.AuthClientIDRillCLI}, + "scope": {scope}, + }) + require.NoError(t, res.err) + require.Equal(t, http.StatusBadRequest, res.status) + require.Equal(t, "invalid scope\n", res.body) + } + + res = oauth004PostForm(t.Context(), deviceAuthorizationURL, url.Values{ + "client_id": {database.AuthClientIDRillCLI}, + "scope": {" full_account "}, + }) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + + var payload serverauth.DeviceCodeResponse + require.NoError(t, json.Unmarshal([]byte(res.body), &payload)) + require.NotEmpty(t, payload.DeviceCode) + require.Len(t, strings.ReplaceAll(payload.UserCode, "-", ""), 8) + require.Equal(t, int(admin.DeviceAuthCodeTTL.Seconds()), payload.ExpiresIn) + // The advertised interval is the current polling contract. A slow_down + // response is not supported because the server stores no poll history. + require.Equal(t, 5, payload.PollingInterval) + + stored, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), payload.DeviceCode) + require.NoError(t, err) + require.Equal(t, database.DeviceAuthCodeStatePending, stored.ApprovalState) + }) + + t.Run("keeps pending user codes unique and retries a collision", func(t *testing.T) { + // A confirmation code must select exactly one pending device request. The + // real database constraint makes a duplicate deterministic for this test. + expiresOn := time.Now().Add(time.Minute) + _, err := fix.Admin.DB.InsertDeviceAuthCode(t.Context(), oauth004ID("collision-a"), "COLLIDE1", database.AuthClientIDRillCLI, expiresOn) + require.NoError(t, err) + _, err = fix.Admin.DB.InsertDeviceAuthCode(t.Context(), oauth004ID("collision-b"), "COLLIDE1", database.AuthClientIDRillCLI, expiresOn) + require.ErrorIs(t, err, database.ErrNotUnique) + + originalDB := fix.Admin.DB + collisionDB := &oauth004CollisionDB{DB: originalDB} + oauth004WithDB(t, fix, collisionDB, func() { + res := oauth004PostForm(t.Context(), deviceAuthorizationURL, url.Values{ + "client_id": {database.AuthClientIDRillCLI}, + "scope": {"full_account"}, + }) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + }) + require.EqualValues(t, 2, collisionDB.calls.Load()) + }) + + t.Run("reports pending without consuming the code", func(t *testing.T) { + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStatePending, nil, time.Now().Add(time.Minute)) + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusUnauthorized, res.status) + require.Equal(t, "authorization_pending\n", res.body) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.NoError(t, err) + }) + + t.Run("rejects and consumes a denied code", func(t *testing.T) { + user, _ := fix.NewUser(t) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStateRejected, &user.ID, time.Now().Add(time.Minute)) + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusUnauthorized, res.status) + require.Equal(t, "rejected\n", res.body) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.ErrorIs(t, err, database.ErrNotFound) + }) + + t.Run("never issues from an expired approved code", func(t *testing.T) { + user, _ := fix.NewUser(t) + before := oauth004TokenCount(t, fix.Admin.DB, user.ID) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStateApproved, &user.ID, time.Now().Add(-time.Minute)) + + // Expiry wins over approval: a stale approval must not mint a token. + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusUnauthorized, res.status) + require.Equal(t, "expired_token\n", res.body) + require.Equal(t, before, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.NoError(t, err) + }) + + t.Run("fails closed on a malformed persisted state", func(t *testing.T) { + user, _ := fix.NewUser(t) + before := oauth004TokenCount(t, fix.Admin.DB, user.ID) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeState(99), &user.ID, time.Now().Add(time.Minute)) + + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status) + require.Contains(t, res.body, "inconsistent state") + require.Equal(t, before, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.NoError(t, err) + }) + + t.Run("simultaneous approved polls persist one token", func(t *testing.T) { + user, _ := fix.NewUser(t) + before := oauth004TokenCount(t, fix.Admin.DB, user.ID) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStateApproved, &user.ID, time.Now().Add(time.Minute)) + + const polls = 8 + start := make(chan struct{}) + results := make(chan oauth004FormResult, polls) + var group sync.WaitGroup + for range polls { + group.Add(1) + go func() { + defer group.Done() + <-start + results <- oauth004Poll(context.Background(), tokenURL, code.DeviceCode) + }() + } + close(start) + group.Wait() + close(results) + + successes := 0 + for res := range results { + require.NoError(t, res.err) + if res.status == http.StatusOK { + successes++ + } + } + require.Equal(t, 1, successes) + require.Equal(t, before+1, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.ErrorIs(t, err, database.ErrNotFound) + }) + + t.Run("issuance failure leaves the approved code retryable", func(t *testing.T) { + user, _ := fix.NewUser(t) + before := oauth004TokenCount(t, fix.Admin.DB, user.ID) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStateApproved, &user.ID, time.Now().Add(time.Minute)) + + // Issuance and deletion are one transaction: either both persist or neither does. + oauth004WithDB(t, fix, &oauth004IssueFailureDB{DB: fix.Admin.DB}, func() { + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status) + require.Contains(t, res.body, errInjectedTokenIssue.Error()) + }) + require.Equal(t, before, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.NoError(t, err) + + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + require.Equal(t, before+1, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + }) + + t.Run("deletion failure rolls token issuance back", func(t *testing.T) { + user, _ := fix.NewUser(t) + before := oauth004TokenCount(t, fix.Admin.DB, user.ID) + code := oauth004StoreCode(t, fix.Admin.DB, database.DeviceAuthCodeStateApproved, &user.ID, time.Now().Add(time.Minute)) + + oauth004WithDB(t, fix, &oauth004DeleteFailureDB{DB: fix.Admin.DB}, func() { + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status) + require.Contains(t, res.body, errInjectedCodeDelete.Error()) + }) + require.Equal(t, before, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + _, err := fix.Admin.DB.FindDeviceAuthCodeByDeviceCode(t.Context(), code.DeviceCode) + require.NoError(t, err) + + res := oauth004Poll(t.Context(), tokenURL, code.DeviceCode) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + require.Equal(t, before+1, oauth004TokenCount(t, fix.Admin.DB, user.ID)) + }) +} + +func oauth004PostForm(ctx context.Context, endpoint string, values url.Values) oauth004FormResult { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(values.Encode())) + if err != nil { + return oauth004FormResult{err: err} + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return oauth004FormResult{err: err} + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + return oauth004FormResult{status: resp.StatusCode, body: string(body), err: err} +} + +func oauth004Poll(ctx context.Context, endpoint, deviceCode string) oauth004FormResult { + return oauth004PostForm(ctx, endpoint, url.Values{ + "client_id": {database.AuthClientIDRillCLI}, + "device_code": {deviceCode}, + "grant_type": {oauth004DeviceGrantType}, + "token_response_version": {"standard"}, + }) +} + +func oauth004StoreCode(t *testing.T, db database.DB, state database.DeviceAuthCodeState, userID *string, expiry time.Time) *database.DeviceAuthCode { + t.Helper() + sequence := oauth004Sequence.Add(1) + code, err := db.InsertDeviceAuthCode( + t.Context(), + fmt.Sprintf("oauth004-device-%d", sequence), + fmt.Sprintf("T%07d", sequence), + database.AuthClientIDRillCLI, + expiry, + ) + require.NoError(t, err) + if state != database.DeviceAuthCodeStatePending { + require.NotNil(t, userID) + require.NoError(t, db.UpdateDeviceAuthCode(t.Context(), code.ID, *userID, state)) + } + return code +} + +func oauth004TokenCount(t *testing.T, db database.DB, userID string) int { + t.Helper() + tokens, err := db.FindUserAuthTokens(t.Context(), userID, "", 100, nil) + require.NoError(t, err) + return len(tokens) +} + +func oauth004WithDB(t *testing.T, fix *testadmin.Fixture, db database.DB, fn func()) { + t.Helper() + original := fix.Admin.DB + fix.Admin.DB = db + defer func() { fix.Admin.DB = original }() + fn() +} + +func oauth004ID(label string) string { + return fmt.Sprintf("oauth004-%s-%d", label, oauth004Sequence.Add(1)) +} diff --git a/admin/server/auth/handlers.go b/admin/server/auth/handlers.go index ba8c736f34e1..197e5eea566b 100644 --- a/admin/server/auth/handlers.go +++ b/admin/server/auth/handlers.go @@ -659,6 +659,9 @@ func (a *Authenticator) handleAuthorizeRequest(w http.ResponseWriter, r *http.Re // after login redirect back to same path so encode the current URL as a redirect parameter encodedURL := url.QueryEscape(r.URL.String()) http.Redirect(w, r, "/auth/login?redirect="+encodedURL, http.StatusTemporaryRedirect) + // A redirect is the complete response; falling through would append a 400 + // body to the already-committed redirect. + return } if claims.OwnerType() != OwnerTypeUser { http.Error(w, "only users can be authorized", http.StatusBadRequest) diff --git a/admin/server/auth/mcp_oauth.go b/admin/server/auth/mcp_oauth.go index 9e3f17bc8980..05736e76a9e2 100644 --- a/admin/server/auth/mcp_oauth.go +++ b/admin/server/auth/mcp_oauth.go @@ -2,8 +2,10 @@ package auth import ( "encoding/json" + "errors" "fmt" "io" + "net" "net/http" "net/url" "strings" @@ -12,6 +14,8 @@ import ( "go.uber.org/zap" ) +const maxOAuthRegistrationBodyBytes = 1 << 20 + // handleOAuthProtectedResourceMetadata serves the OAuth 2.0 Protected Resource Metadata as per RFC 9728 and MCP OAuth specification. // This endpoint helps MCP clients discover the authorization server for this protected resource. https://www.rfc-editor.org/rfc/rfc9728.html func (a *Authenticator) handleOAuthProtectedResourceMetadata(w http.ResponseWriter, r *http.Request) { @@ -80,9 +84,15 @@ func (a *Authenticator) handleOAuthRegister(w http.ResponseWriter, r *http.Reque return } - // Parse request body - body, err := io.ReadAll(r.Body) + // Bound this public endpoint before decoding so a registration request cannot + // consume an arbitrary amount of server memory. + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxOAuthRegistrationBodyBytes)) if err != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + http.Error(w, "request body too large", http.StatusRequestEntityTooLarge) + return + } internalServerError(w, fmt.Errorf("failed to read request body: %w", err)) return } @@ -115,8 +125,19 @@ func (a *Authenticator) handleOAuthRegister(w http.ResponseWriter, r *http.Reque return } - scope := sanitizeScope(req.Scope) - grantTypes := sanitizeGrantTypes(req.GrantTypes) + // Dynamic registration is untrusted. In particular, it must never mint a + // client eligible for non-expiring access tokens; privileged clients are + // provisioned through trusted internal paths instead. + scope, err := validateDynamicRegistrationScope(req.Scope) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + grantTypes, err := validateDynamicRegistrationGrantTypes(req.GrantTypes) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } if len(grantTypes) == 0 { // Default to authorization_code if none provided grantTypes = []string{authorizationCodeGrantType} @@ -142,10 +163,17 @@ func (a *Authenticator) handleOAuthRegister(w http.ResponseWriter, r *http.Reque ClientURI: req.ClientURI, } + // Marshal before committing the status. Once a response or redirect starts, + // an error must not append a second HTTP error body to it. + respBytes, err := json.Marshal(resp) + if err != nil { + internalServerError(w, fmt.Errorf("failed to encode response: %w", err)) + return + } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusCreated) - if err := json.NewEncoder(w).Encode(resp); err != nil { - internalServerError(w, fmt.Errorf("failed to encode response: %w", err)) + if _, err := w.Write(respBytes); err != nil { + a.logger.Error("Failed to write OAuth registration response", zap.Error(err)) return } @@ -169,6 +197,45 @@ func sanitizeGrantTypes(grants []string) []string { return sanitized } +func validateDynamicRegistrationScope(scope string) (string, error) { + seen := make(map[string]struct{}) + var sanitized []string + for _, token := range strings.Fields(scope) { + switch token { + case "offline_access": + // Safe for public clients when paired with an approved refresh grant. + case longLivedAccessTokenScope: + return "", fmt.Errorf("scope %q requires trusted registration", token) + default: + return "", fmt.Errorf("unsupported scope %q", token) + } + if _, ok := seen[token]; ok { + continue + } + seen[token] = struct{}{} + sanitized = append(sanitized, token) + } + return strings.Join(sanitized, " "), nil +} + +func validateDynamicRegistrationGrantTypes(grants []string) ([]string, error) { + seen := make(map[string]struct{}) + var sanitized []string + for _, grant := range sanitizeGrantTypes(grants) { + switch grant { + case authorizationCodeGrantType, refreshTokenGrantType, deviceCodeGrantType: + default: + return nil, fmt.Errorf("unsupported grant_type %q", grant) + } + if _, ok := seen[grant]; ok { + continue + } + seen[grant] = struct{}{} + sanitized = append(sanitized, grant) + } + return sanitized, nil +} + // sanitizeRedirectURIs validates redirect URIs and normalizes them. // Requires absolute HTTP(S) URLs without fragments. func sanitizeRedirectURIs(uris []string) ([]string, error) { @@ -184,12 +251,18 @@ func sanitizeRedirectURIs(uris []string) ([]string, error) { if err != nil { return nil, fmt.Errorf("invalid redirect_uri %q: %w", raw, err) } + parsed.Scheme = strings.ToLower(parsed.Scheme) if parsed.Scheme != "https" && parsed.Scheme != "http" { return nil, fmt.Errorf("redirect_uri %q must use http or https", raw) } if parsed.Host == "" { return nil, fmt.Errorf("redirect_uri %q must include a host", raw) } + // Authorization codes sent over plain HTTP are exposed in transit. Native + // clients may use HTTP only on the local machine; remote callbacks use TLS. + if parsed.Scheme == "http" && !isLoopbackHostname(parsed.Hostname()) { + return nil, fmt.Errorf("redirect_uri %q must use https unless it is loopback", raw) + } if parsed.Fragment != "" { return nil, fmt.Errorf("redirect_uri %q must not include a fragment", raw) } @@ -203,3 +276,11 @@ func sanitizeRedirectURIs(uris []string) ([]string, error) { } return sanitized, nil } + +func isLoopbackHostname(hostname string) bool { + if strings.EqualFold(hostname, "localhost") { + return true + } + ip := net.ParseIP(hostname) + return ip != nil && ip.IsLoopback() +} diff --git a/admin/server/auth/oauth_registration_test.go b/admin/server/auth/oauth_registration_test.go new file mode 100644 index 000000000000..cdca19bd0f52 --- /dev/null +++ b/admin/server/auth/oauth_registration_test.go @@ -0,0 +1,278 @@ +package auth + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/rilldata/rill/admin" + "github.com/rilldata/rill/admin/database" + "github.com/rilldata/rill/admin/pkg/oauth" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func TestOAuthDynamicRegistrationPolicy(t *testing.T) { + // Dynamic registration normalizes safe metadata and rejects scopes, grants, redirects, or bodies outside policy. + t.Run("persists only approved metadata and unique redirects", func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + req := oauth.ClientRegistrationRequest{ + ClientName: "Local MCP client", + Scope: " offline_access offline_access ", + GrantTypes: []string{ + " " + authorizationCodeGrantType + " ", + refreshTokenGrantType, + authorizationCodeGrantType, + }, + RedirectURIs: []string{ + " http://127.0.0.1:43123/callback ", + "http://127.0.0.1:43123/callback", + "http://[::1]:43123/callback", + "https://app.example.com/callback", + }, + } + + w := registerOAuthClient(t, a, req) + + require.Equal(t, http.StatusCreated, w.Code) + var resp oauth.ClientRegistrationResponse + require.NoError(t, json.NewDecoder(w.Body).Decode(&resp)) + stored, err := db.FindAuthClient(t.Context(), resp.ClientID) + require.NoError(t, err) + require.Equal(t, "offline_access", stored.Scope) + require.Equal(t, []string{authorizationCodeGrantType, refreshTokenGrantType}, stored.GrantTypes) + require.Equal(t, []string{ + "http://127.0.0.1:43123/callback", + "http://[::1]:43123/callback", + "https://app.example.com/callback", + }, stored.RedirectURIs) + require.Equal(t, stored.Scope, resp.Scope) + require.Equal(t, stored.GrantTypes, resp.GrantTypes) + require.Equal(t, stored.RedirectURIs, resp.RedirectURIs) + }) + + tests := []struct { + name string + modify func(*oauth.ClientRegistrationRequest) + body string + }{ + { + name: "rejects privileged scope from an untrusted registrant", + modify: func(req *oauth.ClientRegistrationRequest) { + req.Scope = "offline_access " + longLivedAccessTokenScope + }, + body: "requires trusted registration", + }, + { + name: "rejects unknown scope", + modify: func(req *oauth.ClientRegistrationRequest) { + req.Scope = "openid" + }, + body: "unsupported scope", + }, + { + name: "rejects unknown grant", + modify: func(req *oauth.ClientRegistrationRequest) { + req.GrantTypes = append(req.GrantTypes, "client_credentials") + }, + body: "unsupported grant_type", + }, + { + name: "rejects plain HTTP for a remote callback", + modify: func(req *oauth.ClientRegistrationRequest) { + req.RedirectURIs = []string{"http://app.example.com/callback"} + }, + body: "must use https unless it is loopback", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + req := oauth.ClientRegistrationRequest{ + Scope: "offline_access", + GrantTypes: []string{authorizationCodeGrantType}, + RedirectURIs: []string{"https://app.example.com/callback"}, + } + tt.modify(&req) + + w := registerOAuthClient(t, a, req) + + require.Equal(t, http.StatusBadRequest, w.Code) + require.Contains(t, w.Body.String(), tt.body) + require.Empty(t, db.clients, "invalid metadata must not be persisted") + }) + } + + t.Run("rejects an oversized body before decoding or persistence", func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + req := httptest.NewRequest(http.MethodPost, "/auth/oauth/register", strings.NewReader(strings.Repeat("x", maxOAuthRegistrationBodyBytes+1))) + w := httptest.NewRecorder() + + a.handleOAuthRegister(w, req) + + require.Equal(t, http.StatusRequestEntityTooLarge, w.Code) + require.Empty(t, db.clients) + }) +} + +func TestOAuthAuthorizePolicy(t *testing.T) { + // Authorization must distinguish anonymous login redirects from a valid, grant-bound browser consent action. + const ( + callback = "https://app.example.com/callback" + userID = "user-123" + ) + + t.Run("anonymous response is one login redirect", func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + r := newOAuthAuthorizeRequest(t, anonClaims{}, "client-1", callback) + w := httptest.NewRecorder() + + a.handleAuthorizeRequest(w, r) + + require.Equal(t, http.StatusTemporaryRedirect, w.Code) + require.Contains(t, w.Header().Get("Location"), "/auth/login?redirect=") + require.NotContains(t, w.Body.String(), "only users can be authorized") + require.Equal(t, 1, strings.Count(w.Body.String(), "Temporary Redirect")) + require.Empty(t, db.authorizationCodes) + }) + + t.Run("authenticated browser action consents to a registered ordinary client", func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + client, err := db.InsertAuthClient(t.Context(), "MCP client", "offline_access", []string{authorizationCodeGrantType}, []string{callback}) + require.NoError(t, err) + claims := &oauthPolicyTestClaims{ownerType: OwnerTypeUser, ownerID: userID} + r := newOAuthAuthorizeRequest(t, claims, client.ID, callback) + w := httptest.NewRecorder() + + a.handleAuthorizeRequest(w, r) + + require.Equal(t, http.StatusFound, w.Code) + redirect, err := url.Parse(w.Header().Get("Location")) + require.NoError(t, err) + require.Equal(t, "app.example.com", redirect.Host) + require.Equal(t, "/callback", redirect.Path) + require.NotEmpty(t, redirect.Query().Get("code")) + require.Equal(t, "opaque-state", redirect.Query().Get("state")) + require.Len(t, db.authorizationCodes, 1) + code := db.authorizationCodes[0] + require.Equal(t, userID, code.UserID) + require.Equal(t, client.ID, code.ClientID) + require.Equal(t, callback, code.RedirectURI) + require.Equal(t, "test-challenge", code.CodeChallenge) + require.Equal(t, "S256", code.CodeChallengeMethod) + }) + + t.Run("client without authorization code grant is denied", func(t *testing.T) { + db := newOAuthPolicyTestDB() + a := newOAuthPolicyTestAuthenticator(db) + client, err := db.InsertAuthClient(t.Context(), "Refresh-only client", "offline_access", []string{refreshTokenGrantType}, []string{callback}) + require.NoError(t, err) + claims := &oauthPolicyTestClaims{ownerType: OwnerTypeUser, ownerID: userID} + r := newOAuthAuthorizeRequest(t, claims, client.ID, callback) + w := httptest.NewRecorder() + + a.handleAuthorizeRequest(w, r) + + require.Equal(t, http.StatusBadRequest, w.Code) + require.Contains(t, w.Body.String(), "not permitted to use authorization codes") + require.Empty(t, db.authorizationCodes) + }) +} + +func registerOAuthClient(t *testing.T, a *Authenticator, registration oauth.ClientRegistrationRequest) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(registration) + require.NoError(t, err) + req := httptest.NewRequest(http.MethodPost, "/auth/oauth/register", strings.NewReader(string(body))) + w := httptest.NewRecorder() + a.handleOAuthRegister(w, req) + return w +} + +func newOAuthAuthorizeRequest(t *testing.T, claims Claims, clientID, redirectURI string) *http.Request { + t.Helper() + query := url.Values{ + "client_id": {clientID}, + "redirect_uri": {redirectURI}, + "response_type": {"code"}, + "code_challenge": {"test-challenge"}, + "code_challenge_method": {"S256"}, + "state": {"opaque-state"}, + } + req := httptest.NewRequest(http.MethodGet, "/auth/oauth/authorize?"+query.Encode(), nil) + return req.WithContext(WithClaims(req.Context(), claims)) +} + +func newOAuthPolicyTestAuthenticator(db *oauthPolicyTestDB) *Authenticator { + return &Authenticator{ + logger: zap.NewNop(), + admin: &admin.Service{DB: db}, + } +} + +type oauthPolicyTestClaims struct { + Claims + ownerType OwnerType + ownerID string +} + +func (c *oauthPolicyTestClaims) OwnerType() OwnerType { return c.ownerType } +func (c *oauthPolicyTestClaims) OwnerID() string { return c.ownerID } + +type oauthPolicyTestDB struct { + database.DB + clients map[string]*database.AuthClient + authorizationCodes []*database.AuthorizationCode +} + +func newOAuthPolicyTestDB() *oauthPolicyTestDB { + return &oauthPolicyTestDB{clients: make(map[string]*database.AuthClient)} +} + +func (db *oauthPolicyTestDB) InsertAuthClient(_ context.Context, displayName, scope string, grantTypes, redirectURIs []string) (*database.AuthClient, error) { + client := &database.AuthClient{ + ID: fmt.Sprintf("dynamic-client-%d", len(db.clients)+1), + DisplayName: displayName, + Scope: scope, + GrantTypes: append([]string(nil), grantTypes...), + RedirectURIs: append([]string(nil), redirectURIs...), + CreatedOn: time.Unix(1_700_000_000, 0).UTC(), + } + db.clients[client.ID] = client + return client, nil +} + +func (db *oauthPolicyTestDB) FindAuthClient(_ context.Context, id string) (*database.AuthClient, error) { + client, ok := db.clients[id] + if !ok { + return nil, database.ErrNotFound + } + return client, nil +} + +func (db *oauthPolicyTestDB) InsertAuthorizationCode(_ context.Context, code, userID, clientID, redirectURI, codeChallenge, codeChallengeMethod string, expiration time.Time) (*database.AuthorizationCode, error) { + authCode := &database.AuthorizationCode{ + ID: fmt.Sprintf("authorization-code-%d", len(db.authorizationCodes)+1), + Code: code, + UserID: userID, + ClientID: clientID, + RedirectURI: redirectURI, + CodeChallenge: codeChallenge, + CodeChallengeMethod: codeChallengeMethod, + Expiration: expiration, + } + db.authorizationCodes = append(db.authorizationCodes, authCode) + return authCode, nil +} diff --git a/admin/server/auth/pkce.go b/admin/server/auth/pkce.go index d4b8c5c24f2f..583a47596894 100644 --- a/admin/server/auth/pkce.go +++ b/admin/server/auth/pkce.go @@ -36,6 +36,10 @@ func (a *Authenticator) handlePKCE(w http.ResponseWriter, r *http.Request, clien internalServerError(w, fmt.Errorf("failed to lookup auth client, %w", err)) return } + if !hasGrantType(authClient.GrantTypes, authorizationCodeGrantType) { + http.Error(w, "client is not permitted to use authorization codes", http.StatusBadRequest) + return + } if clientID != database.AuthClientIDRillWebLocal && len(authClient.RedirectURIs) == 0 { http.Error(w, "client has no registered redirect URIs", http.StatusBadRequest) @@ -47,6 +51,10 @@ func (a *Authenticator) handlePKCE(w http.ResponseWriter, r *http.Request, clien return } + // A signed-in user's browser request to an allowed callback is the consent + // boundary for ordinary clients. Public registration cannot create privileged + // clients, so those must have arrived through trusted provisioning. + // Generate a unique authorization code code, err := generateRandomString(16) // 16 bytes, resulting in a 32-character hex string if err != nil { @@ -111,8 +119,16 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter responseVersion := values.Get("token_response_version") - // get the authorization code from the database - authCode, err := a.admin.DB.FindAuthorizationCode(r.Context(), code) + // A successful exchange consumes its code and persists every token in one + // transaction. A failed validation or persistence step rolls everything back. + txCtx, tx, err := a.admin.DB.NewTx(r.Context(), false) + if err != nil { + internalServerError(w, fmt.Errorf("failed to start authorization code transaction, %w", err)) + return + } + defer func() { _ = tx.Rollback() }() + + authCode, err := a.admin.DB.FindAuthorizationCode(txCtx, code) if err != nil { if errors.Is(err, database.ErrNotFound) { http.Error(w, "no such authorization code found", http.StatusBadRequest) @@ -128,26 +144,20 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter return } - // remove the authorization code from the database to prevent reuse - err = a.admin.DB.DeleteAuthorizationCode(r.Context(), code) - if err != nil { - internalServerError(w, fmt.Errorf("failed to delete authorization code, %w", err)) - return - } - - // Check if the client ID matches the stored client ID + // These checks bind the code to the client, callback, lifetime, and PKCE + // secret chosen at authorization time. Invalid attempts mint nothing and do + // not consume the code, preventing an attacker from burning a valid code. if authCode.ClientID != clientID { http.Error(w, "invalid client ID", http.StatusBadRequest) return } - // Check if the redirect URI matches the stored redirect URI if authCode.RedirectURI != redirectURI { http.Error(w, "invalid redirect URI", http.StatusBadRequest) return } - authClient, err := a.admin.DB.FindAuthClient(r.Context(), clientID) + authClient, err := a.admin.DB.FindAuthClient(txCtx, clientID) if err != nil { if errors.Is(err, database.ErrNotFound) { http.Error(w, fmt.Sprintf("invalid client_id %q", clientID), http.StatusBadRequest) @@ -156,27 +166,37 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter internalServerError(w, fmt.Errorf("failed to lookup auth client, %w", err)) return } - a.admin.Used.Client(clientID) - // Check if the authorization code has expired - if time.Now().After(authCode.Expiration) { + if !authCode.Expiration.After(time.Now()) { http.Error(w, "authorization code has expired", http.StatusBadRequest) return } - // Verify the code verifier against the stored code challenge if !verifyCodeChallenge(codeVerifier, authCode.CodeChallenge, authCode.CodeChallengeMethod) { http.Error(w, "invalid code verifier", http.StatusBadRequest) return } + // Deleting inside the transaction is the single-use claim. Concurrent valid + // exchanges can both read the code, but PostgreSQL lets only one delete it; + // the loser observes not-found and cannot persist a token. + err = a.admin.DB.DeleteAuthorizationCode(txCtx, code) + if err != nil { + if errors.Is(err, database.ErrNotFound) { + http.Error(w, "no such authorization code found", http.StatusBadRequest) + } else { + internalServerError(w, fmt.Errorf("failed to consume authorization code, %w", err)) + } + return + } + scope := authClient.Scope longLivedAccess := hasScope(scope, longLivedAccessTokenScope) refreshAllowed := hasGrantType(authClient.GrantTypes, refreshTokenGrantType) var respBytes []byte // Issue long-lived access token for clients explicitly whitelisted if longLivedAccess { - authToken, err := a.admin.IssueUserAuthToken(r.Context(), userID, authCode.ClientID, "", nil, nil, false) + authToken, err := a.admin.IssueUserAuthToken(txCtx, userID, authCode.ClientID, "", nil, nil, false) if err != nil { if errors.Is(err, r.Context().Err()) { http.Error(w, "request cancelled or timeout", http.StatusRequestTimeout) @@ -215,7 +235,7 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter } else { // Issue access token (60 minutes TTL) accessTTL := 60 * time.Minute - accessToken, err := a.admin.IssueUserAuthToken(r.Context(), userID, authCode.ClientID, "Access Token", nil, &accessTTL, false) + accessToken, err := a.admin.IssueUserAuthToken(txCtx, userID, authCode.ClientID, "Access Token", nil, &accessTTL, false) if err != nil { if errors.Is(err, r.Context().Err()) { http.Error(w, "request cancelled or timeout", http.StatusRequestTimeout) @@ -235,7 +255,7 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter if refreshAllowed { // Issue refresh token (365 days TTL) refreshTTL := 365 * 24 * time.Hour - refreshToken, err := a.admin.IssueUserAuthToken(r.Context(), userID, authCode.ClientID, "Refresh Token", nil, &refreshTTL, true) + refreshToken, err := a.admin.IssueUserAuthToken(txCtx, userID, authCode.ClientID, "Refresh Token", nil, &refreshTTL, true) if err != nil { if errors.Is(err, r.Context().Err()) { http.Error(w, "request cancelled or timeout", http.StatusRequestTimeout) @@ -253,6 +273,16 @@ func (a *Authenticator) getAccessTokenForAuthorizationCode(w http.ResponseWriter return } } + + // Commit is the all-or-nothing boundary: the consumed code and every access + // or refresh token become durable together, so no failed exchange can leave + // an orphan token behind. + if err := tx.Commit(); err != nil { + internalServerError(w, fmt.Errorf("failed to commit authorization code exchange, %w", err)) + return + } + a.admin.Used.Client(clientID) + w.Header().Set("Content-Type", "application/json") _, err = w.Write(respBytes) if err != nil { @@ -278,8 +308,17 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * return } - // Validate the refresh token - refreshToken, err := a.admin.ValidateAuthToken(r.Context(), refreshTokenStr) + // Rotation is one transaction: consuming the old refresh token is the + // single-use claim, and both replacement tokens commit with that claim. + txCtx, tx, err := a.admin.DB.NewTx(r.Context(), false) + if err != nil { + internalServerError(w, fmt.Errorf("failed to start refresh token transaction, %w", err)) + return + } + defer func() { _ = tx.Rollback() }() + + // Validate the refresh token. + refreshToken, err := a.admin.ValidateAuthToken(txCtx, refreshTokenStr) if err != nil { http.Error(w, "invalid refresh token", http.StatusUnauthorized) return @@ -296,6 +335,12 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * http.Error(w, "token is not a refresh token", http.StatusBadRequest) return } + // ValidateAuthToken has a short cache, so enforce expiry again at the + // rotation boundary in case the token expired after it was cached. + if userToken.ExpiresOn != nil && !userToken.ExpiresOn.After(time.Now()) { + http.Error(w, "invalid refresh token", http.StatusUnauthorized) + return + } userID := refreshToken.OwnerID() tknClientID := "" @@ -307,7 +352,7 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * return } - authClient, err := a.admin.DB.FindAuthClient(r.Context(), clientID) + authClient, err := a.admin.DB.FindAuthClient(txCtx, clientID) if err != nil { if errors.Is(err, database.ErrNotFound) { http.Error(w, fmt.Sprintf("invalid client_id %q", clientID), http.StatusBadRequest) @@ -316,16 +361,26 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * internalServerError(w, fmt.Errorf("failed to lookup auth client, %w", err)) return } - a.admin.Used.Client(clientID) - if !hasGrantType(authClient.GrantTypes, refreshTokenGrantType) { http.Error(w, "client is not permitted to use refresh tokens", http.StatusBadRequest) return } + // Deleting inside the transaction atomically claims this token. Concurrent + // rotations can both validate it, but only one can delete and commit it. + err = a.admin.DB.DeleteUserAuthToken(txCtx, userToken.ID) + if err != nil { + if errors.Is(err, database.ErrNotFound) { + http.Error(w, "invalid refresh token", http.StatusUnauthorized) + } else { + internalServerError(w, fmt.Errorf("failed to consume refresh token, %w", err)) + } + return + } + // Issue new access token (60 minutes TTL) accessTTL := 60 * time.Minute - accessToken, err := a.admin.IssueUserAuthToken(r.Context(), userID, clientID, "Access Token", nil, &accessTTL, false) + accessToken, err := a.admin.IssueUserAuthToken(txCtx, userID, clientID, "Access Token", nil, &accessTTL, false) if err != nil { if errors.Is(err, r.Context().Err()) { http.Error(w, "request cancelled or timeout", http.StatusRequestTimeout) @@ -337,7 +392,7 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * // Issue new refresh token (365 days TTL) newRefreshTTL := 365 * 24 * time.Hour - newRefreshToken, err := a.admin.IssueUserAuthToken(r.Context(), userID, clientID, "Refresh Token", nil, &newRefreshTTL, true) + newRefreshToken, err := a.admin.IssueUserAuthToken(txCtx, userID, clientID, "Refresh Token", nil, &newRefreshTTL, true) if err != nil { if errors.Is(err, r.Context().Err()) { http.Error(w, "request cancelled or timeout", http.StatusRequestTimeout) @@ -360,20 +415,23 @@ func (a *Authenticator) getAccessTokenForRefreshToken(w http.ResponseWriter, r * internalServerError(w, fmt.Errorf("failed to marshal response, %w", err)) return } + + // Commit before writing the response. If the client disconnects while the + // response is written, the old token stays consumed and cannot be replayed + // to create a second durable branch. + if err := tx.Commit(); err != nil { + internalServerError(w, fmt.Errorf("failed to commit refresh token rotation, %w", err)) + return + } + a.admin.PurgeAuthTokenCache() + a.admin.Used.Client(clientID) + w.Header().Set("Content-Type", "application/json") _, err = w.Write(respBytes) if err != nil { internalServerError(w, fmt.Errorf("failed to write response, %w", err)) return } - - // Delete the old refresh token - err = a.admin.RevokeAuthToken(r.Context(), refreshTokenStr) - if err != nil { - a.logger.Warn("failed to revoke old refresh token", zap.Error(err)) - // Continue anyway - } - a.logger.Debug("Refreshed access and refresh tokens", zap.String("userID", userID), zap.String("clientID", clientID)) } diff --git a/admin/server/auth/pkce_flow_test.go b/admin/server/auth/pkce_flow_test.go new file mode 100644 index 000000000000..922a2ddb777d --- /dev/null +++ b/admin/server/auth/pkce_flow_test.go @@ -0,0 +1,333 @@ +package auth_test + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rilldata/rill/admin/database" + "github.com/rilldata/rill/admin/pkg/oauth" + "github.com/rilldata/rill/admin/testadmin" + "github.com/stretchr/testify/require" +) + +const ( + oauth002AuthorizationCodeGrant = "authorization_code" + oauth002RefreshTokenGrant = "refresh_token" + oauth002RedirectURI = "https://client.example/callback" + oauth002Verifier = "oauth002-verifier-with-enough-entropy-for-pkce" +) + +var ( + oauth002Sequence atomic.Uint64 + oauth002ErrAccessPersistence = errors.New("injected access token persistence failure") + oauth002ErrRefreshPersistence = errors.New("injected refresh token persistence failure") +) + +type oauth002FormResult struct { + status int + body string + err error +} + +type oauth002TokenFailureDB struct { + database.DB + failRefresh bool +} + +func (d *oauth002TokenFailureDB) InsertUserAuthToken(ctx context.Context, opts *database.InsertUserAuthTokenOptions) (*database.UserAuthToken, error) { + if opts.Refresh == d.failRefresh { + if d.failRefresh { + return nil, oauth002ErrRefreshPersistence + } + return nil, oauth002ErrAccessPersistence + } + return d.DB.InsertUserAuthToken(ctx, opts) +} + +func TestOAUTH002AuthorizationCodeExchange(t *testing.T) { + // A real-Postgres exchange must bind every PKCE input, consume once, and persist both tokens atomically. + fix := testadmin.New(t) + tokenURL := fix.ExternalURL() + "/auth/oauth/token" + + t.Run("binds client redirect verifier and expiry before consumption", func(t *testing.T) { + tests := []struct { + name string + changeForm func(url.Values) + expiresOn time.Time + wantBody string + retryable bool + }{ + { + name: "wrong client", + changeForm: func(values url.Values) { + values.Set("client_id", database.AuthClientIDRillWeb) + }, + expiresOn: time.Now().Add(time.Minute), + wantBody: "invalid client ID", + retryable: true, + }, + { + name: "wrong redirect", + changeForm: func(values url.Values) { + values.Set("redirect_uri", "https://client.example/wrong") + }, + expiresOn: time.Now().Add(time.Minute), + wantBody: "invalid redirect URI", + retryable: true, + }, + { + name: "wrong verifier", + changeForm: func(values url.Values) { + values.Set("code_verifier", "wrong-verifier") + }, + expiresOn: time.Now().Add(time.Minute), + wantBody: "invalid code verifier", + retryable: true, + }, + { + name: "expired code", + changeForm: func(url.Values) {}, + expiresOn: time.Now().Add(-time.Minute), + wantBody: "authorization code has expired", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + user, _ := fix.NewUser(t) + client := oauth002NewClient(t, fix.Admin.DB, false) + code := oauth002StoreCode(t, fix.Admin.DB, user.ID, client.ID, tt.expiresOn) + before := oauth002UserTokenCount(t, fix.Admin.DB, user.ID) + + values := oauth002ValidExchangeForm(code, client.ID) + tt.changeForm(values) + res := oauth002PostForm(t.Context(), tokenURL, values) + + require.NoError(t, res.err) + require.Equal(t, http.StatusBadRequest, res.status, res.body) + require.Contains(t, res.body, tt.wantBody) + require.Equal(t, before, oauth002UserTokenCount(t, fix.Admin.DB, user.ID), "an invalid exchange must mint nothing") + _, err := fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.NoError(t, err, "validation failures do not consume the authorization code") + + if tt.retryable { + res = oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + require.Equal(t, before+1, oauth002UserTokenCount(t, fix.Admin.DB, user.ID)) + _, err = fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.ErrorIs(t, err, database.ErrNotFound, "a valid exchange consumes the code") + } + }) + } + }) + + t.Run("two concurrent valid exchanges have one winner", func(t *testing.T) { + user, _ := fix.NewUser(t) + client := oauth002NewClient(t, fix.Admin.DB, false) + code := oauth002StoreCode(t, fix.Admin.DB, user.ID, client.ID, time.Now().Add(time.Minute)) + before := oauth002UserTokenCount(t, fix.Admin.DB, user.ID) + + // Both requests may validate the same row, but the transactional delete is + // the single-use claim, so only its winner can persist a token. + start := make(chan struct{}) + results := make(chan oauth002FormResult, 2) + var group sync.WaitGroup + for range 2 { + group.Add(1) + go func() { + defer group.Done() + <-start + results <- oauth002PostForm(context.Background(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + }() + } + close(start) + group.Wait() + close(results) + + successes := 0 + failures := 0 + for res := range results { + require.NoError(t, res.err) + switch res.status { + case http.StatusOK: + successes++ + case http.StatusBadRequest: + failures++ + require.Contains(t, res.body, "no such authorization code found") + default: + require.Failf(t, "unexpected exchange response", "status=%d body=%s", res.status, res.body) + } + } + require.Equal(t, 1, successes) + require.Equal(t, 1, failures) + require.Equal(t, before+1, oauth002UserTokenCount(t, fix.Admin.DB, user.ID)) + require.Len(t, oauth002ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1) + _, err := fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.ErrorIs(t, err, database.ErrNotFound) + + replay := oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, replay.err) + require.Equal(t, http.StatusBadRequest, replay.status, replay.body) + require.Equal(t, before+1, oauth002UserTokenCount(t, fix.Admin.DB, user.ID), "replay must not mint another token") + }) + + t.Run("access token persistence failure rolls back code consumption", func(t *testing.T) { + user, _ := fix.NewUser(t) + client := oauth002NewClient(t, fix.Admin.DB, false) + code := oauth002StoreCode(t, fix.Admin.DB, user.ID, client.ID, time.Now().Add(time.Minute)) + before := oauth002UserTokenCount(t, fix.Admin.DB, user.ID) + + oauth002WithDB(t, fix, &oauth002TokenFailureDB{DB: fix.Admin.DB}, func() { + res := oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status, res.body) + require.Contains(t, res.body, oauth002ErrAccessPersistence.Error()) + }) + + require.Equal(t, before, oauth002UserTokenCount(t, fix.Admin.DB, user.ID)) + require.Empty(t, oauth002ClientTokens(t, fix.Admin.DB, user.ID, client.ID), "failed persistence must leave no access token") + _, err := fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.NoError(t, err, "rolled-back consumption keeps the code retryable") + + res := oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + require.Len(t, oauth002ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1) + }) + + t.Run("refresh token persistence failure leaves no orphan access token", func(t *testing.T) { + user, _ := fix.NewUser(t) + client := oauth002NewClient(t, fix.Admin.DB, true) + code := oauth002StoreCode(t, fix.Admin.DB, user.ID, client.ID, time.Now().Add(time.Minute)) + before := oauth002UserTokenCount(t, fix.Admin.DB, user.ID) + + // The access insert happens first; failing the refresh insert proves that + // both token rows and code consumption share one rollback boundary. + oauth002WithDB(t, fix, &oauth002TokenFailureDB{DB: fix.Admin.DB, failRefresh: true}, func() { + res := oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status, res.body) + require.Contains(t, res.body, oauth002ErrRefreshPersistence.Error()) + }) + + require.Equal(t, before, oauth002UserTokenCount(t, fix.Admin.DB, user.ID)) + require.Empty(t, oauth002ClientTokens(t, fix.Admin.DB, user.ID, client.ID), "rollback must remove the provisional access token") + _, err := fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.NoError(t, err, "failed all-or-nothing issuance keeps the code retryable") + + res := oauth002PostForm(t.Context(), tokenURL, oauth002ValidExchangeForm(code, client.ID)) + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + var payload oauth.TokenResponse + require.NoError(t, json.Unmarshal([]byte(res.body), &payload)) + require.NotEmpty(t, payload.AccessToken) + require.NotEmpty(t, payload.RefreshToken) + + tokens := oauth002ClientTokens(t, fix.Admin.DB, user.ID, client.ID) + require.Len(t, tokens, 2) + refreshTokens := 0 + for _, token := range tokens { + if token.Refresh { + refreshTokens++ + } + } + require.Equal(t, 1, refreshTokens) + _, err = fix.Admin.DB.FindAuthorizationCode(t.Context(), code) + require.ErrorIs(t, err, database.ErrNotFound) + }) +} + +func oauth002NewClient(t *testing.T, db database.DB, refresh bool) *database.AuthClient { + t.Helper() + grants := []string{oauth002AuthorizationCodeGrant} + if refresh { + grants = append(grants, oauth002RefreshTokenGrant) + } + client, err := db.InsertAuthClient(t.Context(), oauth002ID("client"), "offline_access", grants, []string{oauth002RedirectURI}) + require.NoError(t, err) + return client +} + +func oauth002StoreCode(t *testing.T, db database.DB, userID, clientID string, expiresOn time.Time) string { + t.Helper() + code := oauth002ID("code") + _, err := db.InsertAuthorizationCode(t.Context(), code, userID, clientID, oauth002RedirectURI, oauth002Challenge(oauth002Verifier), "S256", expiresOn) + require.NoError(t, err) + return code +} + +func oauth002ValidExchangeForm(code, clientID string) url.Values { + return url.Values{ + "grant_type": {oauth002AuthorizationCodeGrant}, + "code": {code}, + "client_id": {clientID}, + "redirect_uri": {oauth002RedirectURI}, + "code_verifier": {oauth002Verifier}, + "token_response_version": {"standard"}, + } +} + +func oauth002PostForm(ctx context.Context, endpoint string, values url.Values) oauth002FormResult { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(values.Encode())) + if err != nil { + return oauth002FormResult{err: err} + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return oauth002FormResult{err: err} + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + return oauth002FormResult{status: resp.StatusCode, body: string(body), err: err} +} + +func oauth002Challenge(verifier string) string { + digest := sha256.Sum256([]byte(verifier)) + return base64.RawURLEncoding.EncodeToString(digest[:]) +} + +func oauth002UserTokenCount(t *testing.T, db database.DB, userID string) int { + t.Helper() + tokens, err := db.FindUserAuthTokens(t.Context(), userID, "", 100, nil) + require.NoError(t, err) + return len(tokens) +} + +func oauth002ClientTokens(t *testing.T, db database.DB, userID, clientID string) []*database.UserAuthToken { + t.Helper() + tokens, err := db.FindUserAuthTokens(t.Context(), userID, "", 100, nil) + require.NoError(t, err) + res := make([]*database.UserAuthToken, 0, len(tokens)) + for _, token := range tokens { + if token.AuthClientID != nil && *token.AuthClientID == clientID { + res = append(res, token) + } + } + return res +} + +func oauth002WithDB(t *testing.T, fix *testadmin.Fixture, db database.DB, fn func()) { + t.Helper() + original := fix.Admin.DB + fix.Admin.DB = db + defer func() { fix.Admin.DB = original }() + fn() +} + +func oauth002ID(label string) string { + return fmt.Sprintf("oauth002-%s-%d", label, oauth002Sequence.Add(1)) +} diff --git a/admin/server/auth/refresh_token_test.go b/admin/server/auth/refresh_token_test.go new file mode 100644 index 000000000000..45f465d0ed12 --- /dev/null +++ b/admin/server/auth/refresh_token_test.go @@ -0,0 +1,418 @@ +package auth_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rilldata/rill/admin/database" + "github.com/rilldata/rill/admin/pkg/authtoken" + "github.com/rilldata/rill/admin/pkg/oauth" + "github.com/rilldata/rill/admin/testadmin" + "github.com/stretchr/testify/require" +) + +const oauth003RefreshTokenGrant = "refresh_token" + +var ( + oauth003Sequence atomic.Uint64 + oauth003ErrAccessPersistence = errors.New("injected access token persistence failure") + oauth003ErrRefreshPersistence = errors.New("injected refresh token persistence failure") + oauth003ErrRevokePersistence = errors.New("injected old token revoke failure") + oauth003ErrResponseWrite = errors.New("injected response write failure") +) + +type oauth003FormResult struct { + status int + body string + err error +} + +type oauth003FailureDB struct { + database.DB + failInsertRefresh *bool + failDeleteID string +} + +func (d *oauth003FailureDB) InsertUserAuthToken(ctx context.Context, opts *database.InsertUserAuthTokenOptions) (*database.UserAuthToken, error) { + if d.failInsertRefresh != nil && opts.Refresh == *d.failInsertRefresh { + if opts.Refresh { + return nil, oauth003ErrRefreshPersistence + } + return nil, oauth003ErrAccessPersistence + } + return d.DB.InsertUserAuthToken(ctx, opts) +} + +func (d *oauth003FailureDB) DeleteUserAuthToken(ctx context.Context, id string) error { + if id == d.failDeleteID { + return oauth003ErrRevokePersistence + } + return d.DB.DeleteUserAuthToken(ctx, id) +} + +type oauth003FailingResponseWriter struct { + header http.Header + writes int +} + +func (w *oauth003FailingResponseWriter) Header() http.Header { + return w.header +} + +func (w *oauth003FailingResponseWriter) WriteHeader(int) {} + +func (w *oauth003FailingResponseWriter) Write([]byte) (int, error) { + w.writes++ + return 0, oauth003ErrResponseWrite +} + +func TestOAUTH003RefreshTokenRotation(t *testing.T) { + // Setup: share one real Postgres-backed admin fixture across isolated users and clients. + fix := testadmin.New(t) + tokenURL := fix.ExternalURL() + "/auth/oauth/token" + // Action: each subtest exercises one rotation, replay, concurrency, or failure path. + // Assertion: every path checks both the HTTP result and the durable token rows. + + t.Run("valid rotation is durable and single use", func(t *testing.T) { + // Setup: issue one refresh token for a client that permits refresh grants. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + oldToken := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + + // Action: exchange the refresh token once through the public token endpoint. + res := oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + + // Assertion: the old token is gone and exactly one access/refresh pair is usable. + require.NoError(t, res.err) + require.Equal(t, http.StatusOK, res.status, res.body) + payload := oauth003DecodeResponse(t, res.body) + require.NotEmpty(t, payload.AccessToken) + require.NotEmpty(t, payload.RefreshToken) + require.NotEqual(t, oldToken, payload.RefreshToken) + _, err := fix.Admin.ValidateAuthToken(t.Context(), oldToken) + require.Error(t, err, "the committed rotation must invalidate the old token immediately") + _, err = fix.Admin.ValidateAuthToken(t.Context(), payload.AccessToken) + require.NoError(t, err) + _, err = fix.Admin.ValidateAuthToken(t.Context(), payload.RefreshToken) + require.NoError(t, err) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + }) + + t.Run("invalid tokens mint nothing", func(t *testing.T) { + // Setup: enumerate malformed, expired, and cached-then-revoked token states. + tests := []struct { + name string + setup func(*testing.T, *database.User, *database.AuthClient) string + }{ + { + name: "malformed", + setup: func(*testing.T, *database.User, *database.AuthClient) string { + return "not-an-auth-token" + }, + }, + { + name: "expired", + setup: func(t *testing.T, user *database.User, client *database.AuthClient) string { + return oauth003IssueRefreshToken(t, fix, user.ID, client.ID, -time.Hour) + }, + }, + { + name: "revoked", + setup: func(t *testing.T, user *database.User, client *database.AuthClient) string { + token := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + _, err := fix.Admin.ValidateAuthToken(t.Context(), token) + require.NoError(t, err, "prime the validation cache before revocation") + require.NoError(t, fix.Admin.RevokeAuthToken(t.Context(), token)) + return token + }, + }, + } + + // Action: submit each unusable state through the same public refresh endpoint. + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Setup: prepare a malformed, expired, or already-revoked refresh token. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + token := tt.setup(t, user, client) + + // Action: submit the unusable token to the refresh grant endpoint. + res := oauth003PostRefresh(t.Context(), tokenURL, client.ID, token) + + // Assertion: authentication fails and no replacement pair is persisted. + require.NoError(t, res.err) + require.Equal(t, http.StatusUnauthorized, res.status, res.body) + require.Contains(t, res.body, "invalid refresh token") + require.Empty(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID)) + }) + } + // Assertion: the per-case checks prove no invalid state persisted a replacement. + }) + + t.Run("simultaneous reuse has exactly one winner", func(t *testing.T) { + // Setup: synchronize two requests that carry the same valid refresh token. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + oldToken := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + start := make(chan struct{}) + results := make(chan oauth003FormResult, 2) + var group sync.WaitGroup + for range 2 { + group.Add(1) + go func() { + defer group.Done() + <-start + results <- oauth003PostRefresh(context.Background(), tokenURL, client.ID, oldToken) + }() + } + + // Action: release both requests together without timing sleeps. + close(start) + group.Wait() + close(results) + + // Assertion: one claim commits, one loses, and only one replacement pair remains. + successes := 0 + failures := 0 + for res := range results { + require.NoError(t, res.err) + switch res.status { + case http.StatusOK: + successes++ + case http.StatusUnauthorized: + failures++ + require.Contains(t, res.body, "invalid refresh token") + default: + require.Failf(t, "unexpected rotation response", "status=%d body=%s", res.status, res.body) + } + } + require.Equal(t, 1, successes) + require.Equal(t, 1, failures) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + _, err := fix.Admin.ValidateAuthToken(t.Context(), oldToken) + require.Error(t, err) + }) + + t.Run("ancestor replay cannot fork the active family", func(t *testing.T) { + // Setup: rotate an initial token to establish one active descendant. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + ancestor := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + first := oauth003PostRefresh(t.Context(), tokenURL, client.ID, ancestor) + require.NoError(t, first.err) + require.Equal(t, http.StatusOK, first.status, first.body) + descendant := oauth003DecodeResponse(t, first.body).RefreshToken + + // Action: replay the consumed ancestor, then rotate the legitimate descendant. + replay := oauth003PostRefresh(t.Context(), tokenURL, client.ID, ancestor) + second := oauth003PostRefresh(t.Context(), tokenURL, client.ID, descendant) + + // Assertion: replay mints no branch while the sole descendant advances normally. + require.NoError(t, replay.err) + require.Equal(t, http.StatusUnauthorized, replay.status, replay.body) + require.NoError(t, second.err) + require.Equal(t, http.StatusOK, second.status, second.body) + latest := oauth003DecodeResponse(t, second.body).RefreshToken + require.NotEqual(t, descendant, latest) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 2, 1) + descendantReplay := oauth003PostRefresh(t.Context(), tokenURL, client.ID, descendant) + require.NoError(t, descendantReplay.err) + require.Equal(t, http.StatusUnauthorized, descendantReplay.status, descendantReplay.body) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 2, 1) + }) + + t.Run("replacement persistence failures roll back consumption", func(t *testing.T) { + // Setup: inject failures separately for the new access and refresh token rows. + failAccess := false + failRefresh := true + tests := []struct { + name string + failRefresh *bool + wantError error + }{ + {name: "access insert", failRefresh: &failAccess, wantError: oauth003ErrAccessPersistence}, + {name: "refresh insert", failRefresh: &failRefresh, wantError: oauth003ErrRefreshPersistence}, + } + + // Action: run each insert failure after the old token has been claimed in a transaction. + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Setup: keep one valid old token and inject a replacement-row insert failure. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + oldToken := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + failureDB := &oauth003FailureDB{DB: fix.Admin.DB, failInsertRefresh: tt.failRefresh} + + // Action: attempt rotation while the selected replacement insert fails. + var res oauth003FormResult + oauth003WithDB(t, fix, failureDB, func() { + res = oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + }) + + // Assertion: the transaction leaves only the old token, which remains retryable. + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status, res.body) + require.Contains(t, res.body, tt.wantError.Error()) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 0, 1) + retry := oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + require.NoError(t, retry.err) + require.Equal(t, http.StatusOK, retry.status, retry.body) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + }) + } + // Assertion: every case retries the old token, proving rollback restored one family. + }) + + t.Run("old token revoke failure creates no replacement branch", func(t *testing.T) { + // Setup: inject a database failure for the old token's single-use delete. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + oldToken := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + failureDB := &oauth003FailureDB{DB: fix.Admin.DB, failDeleteID: oauth003TokenID(t, oldToken)} + + // Action: rotate while persistence refuses to revoke the old token. + var res oauth003FormResult + oauth003WithDB(t, fix, failureDB, func() { + res = oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + }) + + // Assertion: no replacements survive, and the unchanged old token can be retried. + require.NoError(t, res.err) + require.Equal(t, http.StatusInternalServerError, res.status, res.body) + require.Contains(t, res.body, oauth003ErrRevokePersistence.Error()) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 0, 1) + retry := oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + require.NoError(t, retry.err) + require.Equal(t, http.StatusOK, retry.status, retry.body) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + }) + + t.Run("response write failure cannot make a second family", func(t *testing.T) { + // Setup: invoke the real route with a writer that rejects every response body. + user, _ := fix.NewUser(t) + client := oauth003NewClient(t, fix.Admin.DB) + oldToken := oauth003IssueRefreshToken(t, fix, user.ID, client.ID, time.Hour) + handler, err := fix.Server.HTTPHandler(t.Context()) + require.NoError(t, err) + form := oauth003RefreshForm(client.ID, oldToken) + req := httptest.NewRequest(http.MethodPost, tokenURL, strings.NewReader(form.Encode())) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + writer := &oauth003FailingResponseWriter{header: make(http.Header)} + + // Action: complete rotation while delivery of its committed response fails. + handler.ServeHTTP(writer, req) + + // Assertion: one committed pair exists and replay cannot recover or fork it. + require.Positive(t, writer.writes) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + _, err = fix.Admin.ValidateAuthToken(t.Context(), oldToken) + require.Error(t, err) + replay := oauth003PostRefresh(t.Context(), tokenURL, client.ID, oldToken) + require.NoError(t, replay.err) + require.Equal(t, http.StatusUnauthorized, replay.status, replay.body) + oauth003RequireTokenShape(t, oauth003ClientTokens(t, fix.Admin.DB, user.ID, client.ID), 1, 1) + }) +} + +func oauth003NewClient(t *testing.T, db database.DB) *database.AuthClient { + t.Helper() + client, err := db.InsertAuthClient(t.Context(), oauth003ID("client"), "offline_access", []string{oauth003RefreshTokenGrant}, []string{"https://client.example/callback"}) + require.NoError(t, err) + return client +} + +func oauth003IssueRefreshToken(t *testing.T, fix *testadmin.Fixture, userID, clientID string, ttl time.Duration) string { + t.Helper() + token, err := fix.Admin.IssueUserAuthToken(t.Context(), userID, clientID, "Refresh Token", nil, &ttl, true) + require.NoError(t, err) + return token.Token().String() +} + +func oauth003PostRefresh(ctx context.Context, endpoint, clientID, refreshToken string) oauth003FormResult { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(oauth003RefreshForm(clientID, refreshToken).Encode())) + if err != nil { + return oauth003FormResult{err: err} + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return oauth003FormResult{err: err} + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + return oauth003FormResult{status: resp.StatusCode, body: string(body), err: err} +} + +func oauth003RefreshForm(clientID, refreshToken string) url.Values { + return url.Values{ + "grant_type": {oauth003RefreshTokenGrant}, + "client_id": {clientID}, + "refresh_token": {refreshToken}, + } +} + +func oauth003DecodeResponse(t *testing.T, body string) oauth.TokenResponse { + t.Helper() + var payload oauth.TokenResponse + require.NoError(t, json.Unmarshal([]byte(body), &payload)) + return payload +} + +func oauth003ClientTokens(t *testing.T, db database.DB, userID, clientID string) []*database.UserAuthToken { + t.Helper() + tokens, err := db.FindUserAuthTokens(t.Context(), userID, "", 100, nil) + require.NoError(t, err) + res := make([]*database.UserAuthToken, 0, len(tokens)) + for _, token := range tokens { + if token.AuthClientID != nil && *token.AuthClientID == clientID { + res = append(res, token) + } + } + return res +} + +func oauth003RequireTokenShape(t *testing.T, tokens []*database.UserAuthToken, wantAccess, wantRefresh int) { + t.Helper() + access := 0 + refresh := 0 + for _, token := range tokens { + if token.Refresh { + refresh++ + } else { + access++ + } + } + require.Equal(t, wantAccess, access, "unexpected durable access-token count") + require.Equal(t, wantRefresh, refresh, "unexpected durable refresh-token count") + require.Len(t, tokens, wantAccess+wantRefresh) +} + +func oauth003TokenID(t *testing.T, token string) string { + t.Helper() + parsed, err := authtoken.FromString(token) + require.NoError(t, err) + return parsed.ID.String() +} + +func oauth003WithDB(t *testing.T, fix *testadmin.Fixture, db database.DB, fn func()) { + t.Helper() + original := fix.Admin.DB + fix.Admin.DB = db + defer func() { fix.Admin.DB = original }() + fn() +} + +func oauth003ID(label string) string { + return fmt.Sprintf("oauth003-%s-%d", label, oauth003Sequence.Add(1)) +} diff --git a/cli/pkg/deviceauth/authenticator.go b/cli/pkg/deviceauth/authenticator.go index 7b6b16420179..eaeb8c904e2f 100644 --- a/cli/pkg/deviceauth/authenticator.go +++ b/cli/pkg/deviceauth/authenticator.go @@ -25,10 +25,20 @@ var ( ErrCodeRejected = fmt.Errorf("confirmation code rejected") ) +const slowDownBackoff = 5 * time.Second + +type pollAction int + +const ( + pollComplete pollAction = iota + pollPending + pollSlowDown +) + // Authenticator is the interface for authentication via device oauth type Authenticator interface { - VerifyDevice(ctx context.Context) (*DeviceVerification, error) - GetAccessTokenForDevice(ctx context.Context, v DeviceVerification) (string, error) + VerifyDevice(ctx context.Context, redirectURL string) (*DeviceVerification, error) + GetAccessTokenForDevice(ctx context.Context, v *DeviceVerification) (*oauth.TokenResponse, error) } // DeviceCodeResponse encapsulates the response for obtaining a device code. @@ -53,10 +63,11 @@ type DeviceVerification struct { // DeviceAuthenticator performs the authentication flow for logging in. type DeviceAuthenticator struct { - client *http.Client - BaseURL *url.URL - Clock clock.Clock - ClientID string + client *http.Client + waitForNextPoll func(context.Context, time.Duration) error + BaseURL *url.URL + Clock clock.Clock + ClientID string } // New returns an instance of the DeviceAuthenticator @@ -81,7 +92,7 @@ func (d *DeviceAuthenticator) VerifyDevice(ctx context.Context, redirectURL stri req, err := d.newFormRequest(ctx, "auth/oauth/device_authorization", url.Values{ "client_id": []string{d.ClientID}, "scope": []string{"full_account"}, - "redirect": []string{url.QueryEscape(redirectURL)}, + "redirect": []string{redirectURL}, }) if err != nil { return nil, err @@ -93,9 +104,13 @@ func (d *DeviceAuthenticator) VerifyDevice(ctx context.Context, redirectURL stri } defer res.Body.Close() - if _, err = checkErrorResponse(res); err != nil { + action, err := checkErrorResponse(res) + if err != nil { return nil, err } + if action != pollComplete { + return nil, errors.New("unexpected retry response while requesting a device code") + } deviceCodeRes := &DeviceCodeResponse{} err = json.NewDecoder(res.Body).Decode(deviceCodeRes) @@ -122,35 +137,72 @@ func (d *DeviceAuthenticator) VerifyDevice(ctx context.Context, redirectURL stri // GetAccessTokenForDevice uses the device verification response to fetch an access token. func (d *DeviceAuthenticator) GetAccessTokenForDevice(ctx context.Context, v *DeviceVerification) (*oauth.TokenResponse, error) { + pollInterval := v.CheckInterval + if pollInterval <= 0 { + pollInterval = 5 * time.Second + } + for { - // This loop begins right after we open the user's browser to send an - // authentication code. We don't request a token immediately because the - // has to complete that authentication flow before we can provide a - // token anyway. - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(v.CheckInterval): - // Ready to check again. + if err := ctx.Err(); err != nil { + return nil, err } - token, err := d.requestToken(ctx, v.DeviceCode, d.ClientID) + // Poll only after the server-provided interval. The wait is capped at the + // device-code expiry so a long interval cannot keep login alive past it. + now := d.Clock.Now() + if !now.Before(v.ExpiresAt) { + return nil, ErrAuthenticationTimedout + } + waitFor := pollInterval + if remaining := v.ExpiresAt.Sub(now); remaining < waitFor { + waitFor = remaining + } + wait := d.waitForNextPoll + if wait == nil { + wait = d.wait + } + if err := wait(ctx, waitFor); err != nil { + return nil, err + } + if !d.Clock.Now().Before(v.ExpiresAt) { + return nil, ErrAuthenticationTimedout + } + + token, action, err := d.requestToken(ctx, v.DeviceCode, d.ClientID) if err != nil { - // Fatal error. return nil, err } - if token != nil { - // Successful authentication. + switch action { + case pollComplete: + // A token is exposed only after the complete success response has been + // decoded and validated. Callers can persist it after this return without + // risking a credential change for a pending or partial response. return token, nil + case pollSlowDown: + // RFC 8628 requires every poll after slow_down to use an interval that + // is at least five seconds longer than the previous one. + pollInterval += slowDownBackoff + case pollPending: + // The user has not completed authorization yet; keep the current interval. } + } +} - if d.Clock.Now().After(v.ExpiresAt) { - return nil, ErrAuthenticationTimedout - } +// wait keeps cancellation responsive while the poller is between HTTP requests. +// It is a field-backed method so tests can substitute a deterministic sleeper. +func (d *DeviceAuthenticator) wait(ctx context.Context, interval time.Duration) error { + timer := d.Clock.Timer(interval) + defer timer.Stop() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil } } -func (d *DeviceAuthenticator) requestToken(ctx context.Context, deviceCode, clientID string) (*oauth.TokenResponse, error) { +func (d *DeviceAuthenticator) requestToken(ctx context.Context, deviceCode, clientID string) (*oauth.TokenResponse, pollAction, error) { req, err := d.newFormRequest(ctx, "auth/oauth/token", url.Values{ "grant_type": []string{"urn:ietf:params:oauth:grant-type:device_code"}, "device_code": []string{deviceCode}, @@ -158,32 +210,34 @@ func (d *DeviceAuthenticator) requestToken(ctx context.Context, deviceCode, clie "token_response_version": []string{"standard"}, // For backward compatibility with older Rill CLI, see utils.go in oauth pkg }) if err != nil { - return nil, fmt.Errorf("error creating request: %w", err) + return nil, pollComplete, fmt.Errorf("error creating request: %w", err) } res, err := d.client.Do(req) if err != nil { - return nil, fmt.Errorf("error performing http request: %w", err) + return nil, pollComplete, fmt.Errorf("error performing http request: %w", err) } defer res.Body.Close() - isRetryable, err := checkErrorResponse(res) + action, err := checkErrorResponse(res) if err != nil { - return nil, err + return nil, pollComplete, err } - // Bail early so the token fetching is retried. - if isRetryable { - return nil, nil + if action != pollComplete { + return nil, action, nil } tokenRes := &oauth.TokenResponse{} err = json.NewDecoder(res.Body).Decode(tokenRes) if err != nil { - return nil, fmt.Errorf("error decoding token response: %w", err) + return nil, pollComplete, fmt.Errorf("error decoding token response: %w", err) + } + if tokenRes.AccessToken == "" { + return nil, pollComplete, errors.New("token response is missing access_token") } - return tokenRes, nil + return tokenRes, pollComplete, nil } // newFormRequest creates a new form URL encoded request @@ -210,29 +264,29 @@ func (d *DeviceAuthenticator) newFormRequest(ctx context.Context, path string, p return req, nil } -// checkErrorResponse returns whether the error is retryable or not and the error itself. -func checkErrorResponse(res *http.Response) (bool, error) { +// checkErrorResponse classifies the two retry instructions used by device polling. +func checkErrorResponse(res *http.Response) (pollAction, error) { if res.StatusCode < http.StatusBadRequest { - // 200 OK, etc. - return false, nil + return pollComplete, nil } - // Client or server error. body, err := io.ReadAll(res.Body) if err != nil { - return false, err + return pollComplete, err } bodyStr := string(bytes.TrimSpace(body)) - // If we're polling and haven't authorized yet or we need to slow down, we don't want to terminate the polling - if bodyStr == "authorization_pending" || bodyStr == "slow_down" { - return true, nil - } - if bodyStr == "expired_token" { - return false, errors.New(bodyStr) - } - if bodyStr == "rejected" { - return false, ErrCodeRejected + if res.StatusCode < http.StatusInternalServerError { + switch bodyStr { + case "authorization_pending": + return pollPending, nil + case "slow_down": + return pollSlowDown, nil + case "expired_token": + return pollComplete, ErrAuthenticationTimedout + case "rejected": + return pollComplete, ErrCodeRejected + } } - return false, errors.New(strconv.Itoa(res.StatusCode) + ": " + bodyStr) + return pollComplete, errors.New(strconv.Itoa(res.StatusCode) + ": " + bodyStr) } diff --git a/cli/pkg/deviceauth/authenticator_test.go b/cli/pkg/deviceauth/authenticator_test.go new file mode 100644 index 000000000000..4602c42c0fc6 --- /dev/null +++ b/cli/pkg/deviceauth/authenticator_test.go @@ -0,0 +1,424 @@ +package deviceauth + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync/atomic" + "testing" + "time" + + "github.com/benbjohnson/clock" + "github.com/rilldata/rill/admin/pkg/oauth" + "github.com/stretchr/testify/require" +) + +type requestSnapshot struct { + method string + path string + contentType string + accept string + form url.Values + parseErr error +} + +var _ Authenticator = (*DeviceAuthenticator)(nil) + +func TestVerifyDeviceContract(t *testing.T) { + t.Parallel() + // Verify the initial device-authorization form and its conversion into deterministic polling state. + + requests := make(chan requestSnapshot, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests <- snapshotRequest(r) + w.Header().Set("Content-Type", oauth.JSONMediaType) + _, _ = io.WriteString(w, `{ + "device_code":"device-code", + "user_code":"ABCD-EFGH", + "verification_uri":"https://auth.example/verify", + "verification_uri_complete":"https://auth.example/verify?user_code=ABCD-EFGH", + "expires_in":60, + "interval":2 + }`) + })) + t.Cleanup(server.Close) + + mockClock := clock.NewMock() + authenticator := newTestAuthenticator(t, server, mockClock) + redirectURL := "http://localhost:9009/auth/callback?project=my project&open=true" + verification, err := authenticator.VerifyDevice(context.Background(), redirectURL) + require.NoError(t, err) + + request := <-requests + require.NoError(t, request.parseErr) + require.Equal(t, http.MethodPost, request.method) + require.Equal(t, "/auth/oauth/device_authorization", request.path) + require.Equal(t, oauth.FormMediaType, request.contentType) + require.Equal(t, oauth.JSONMediaType, request.accept) + require.Equal(t, url.Values{ + "client_id": {authenticator.ClientID}, + "scope": {"full_account"}, + "redirect": {redirectURL}, + }, request.form) + require.Equal(t, &DeviceVerification{ + DeviceCode: "device-code", + UserCode: "ABCD-EFGH", + VerificationURL: "https://auth.example/verify", + VerificationCompleteURL: "https://auth.example/verify?user_code=ABCD-EFGH", + CheckInterval: 2 * time.Second, + ExpiresAt: mockClock.Now().Add(time.Minute), + }, verification) +} + +func TestVerifyDeviceFailures(t *testing.T) { + t.Parallel() + // Malformed and unsuccessful authorization responses must remain errors rather than partial device sessions. + + tests := []struct { + name string + status int + body string + wantErrText string + }{ + { + name: "malformed JSON", + status: http.StatusOK, + body: `{"device_code":`, + wantErrText: "error decoding device code response", + }, + { + name: "server error", + status: http.StatusInternalServerError, + body: "backend unavailable", + wantErrText: "500: backend unavailable", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(tt.status) + _, _ = io.WriteString(w, tt.body) + })) + t.Cleanup(server.Close) + + authenticator := newTestAuthenticator(t, server, clock.NewMock()) + verification, err := authenticator.VerifyDevice(context.Background(), "http://localhost/callback") + require.ErrorContains(t, err, tt.wantErrText) + require.Nil(t, verification) + }) + } +} + +func TestGetAccessTokenForDeviceContract(t *testing.T) { + t.Parallel() + // Model the complete polling protocol, including backoff, expiry, terminal errors, and credential persistence. + + type reply struct { + status int + body string + } + tests := []struct { + name string + replies []reply + interval time.Duration + expiresIn time.Duration + wantWaits []time.Duration + wantRequests int + wantToken string + wantErrIs error + wantErrText string + wantPersisted string + }{ + { + name: "authorization pending then success", + replies: []reply{ + {status: http.StatusUnauthorized, body: "authorization_pending"}, + {status: http.StatusOK, body: `{"access_token":"new-token","token_type":"Bearer"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second, 2 * time.Second}, + wantRequests: 2, + wantToken: "new-token", + wantPersisted: "new-token", + }, + { + name: "slow down increases all later intervals", + replies: []reply{ + {status: http.StatusBadRequest, body: "slow_down"}, + {status: http.StatusUnauthorized, body: "authorization_pending"}, + {status: http.StatusOK, body: `{"access_token":"new-token","token_type":"Bearer"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second, 7 * time.Second, 7 * time.Second}, + wantRequests: 3, + wantToken: "new-token", + wantPersisted: "new-token", + }, + { + name: "rejection is fatal", + replies: []reply{ + {status: http.StatusUnauthorized, body: "rejected"}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second}, + wantRequests: 1, + wantErrIs: ErrCodeRejected, + wantPersisted: "old-token", + }, + { + name: "server reports expiry", + replies: []reply{ + {status: http.StatusUnauthorized, body: "expired_token"}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second}, + wantRequests: 1, + wantErrIs: ErrAuthenticationTimedout, + wantPersisted: "old-token", + }, + { + name: "local expiry caps the final wait", + replies: []reply{ + {status: http.StatusUnauthorized, body: "authorization_pending"}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: 3 * time.Second, + wantWaits: []time.Duration{2 * time.Second, time.Second}, + wantRequests: 1, + wantErrIs: ErrAuthenticationTimedout, + wantPersisted: "old-token", + }, + { + name: "malformed token JSON is fatal", + replies: []reply{ + {status: http.StatusOK, body: `{"access_token":`}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second}, + wantRequests: 1, + wantErrText: "error decoding token response", + wantPersisted: "old-token", + }, + { + name: "empty token response is not success", + replies: []reply{ + {status: http.StatusOK, body: `{}`}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second}, + wantRequests: 1, + wantErrText: "missing access_token", + wantPersisted: "old-token", + }, + { + name: "server error is fatal", + replies: []reply{ + {status: http.StatusInternalServerError, body: "backend unavailable"}, + {status: http.StatusOK, body: `{"access_token":"must-not-be-read"}`}, + }, + interval: 2 * time.Second, + expiresIn: time.Minute, + wantWaits: []time.Duration{2 * time.Second}, + wantRequests: 1, + wantErrText: "500: backend unavailable", + wantPersisted: "old-token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var persisted atomic.Value + persisted.Store("old-token") + var changedBeforeSuccess atomic.Bool + requests := make(chan requestSnapshot, len(tt.replies)+1) + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if persisted.Load().(string) != "old-token" { + changedBeforeSuccess.Store(true) + } + requests <- snapshotRequest(r) + index := int(requestCount.Add(1)) - 1 + if index >= len(tt.replies) { + http.Error(w, "unexpected extra poll", http.StatusInternalServerError) + return + } + reply := tt.replies[index] + w.WriteHeader(reply.status) + _, _ = io.WriteString(w, reply.body) + })) + t.Cleanup(server.Close) + + mockClock := clock.NewMock() + authenticator := newTestAuthenticator(t, server, mockClock) + var waits []time.Duration + authenticator.waitForNextPoll = func(ctx context.Context, interval time.Duration) error { + if err := ctx.Err(); err != nil { + return err + } + waits = append(waits, interval) + mockClock.Add(interval) + return nil + } + verification := &DeviceVerification{ + DeviceCode: "device-code", + CheckInterval: tt.interval, + ExpiresAt: mockClock.Now().Add(tt.expiresIn), + } + + token, err := authenticator.GetAccessTokenForDevice(context.Background(), verification) + // This mirrors the CLI persistence boundary: no local credential is + // replaced until polling returns a complete, usable token. + if err == nil { + persisted.Store(token.AccessToken) + } + + if tt.wantErrIs != nil { + require.ErrorIs(t, err, tt.wantErrIs) + } else if tt.wantErrText != "" { + require.ErrorContains(t, err, tt.wantErrText) + } else { + require.NoError(t, err) + require.Equal(t, tt.wantToken, token.AccessToken) + } + if err != nil { + require.Nil(t, token) + } + require.Equal(t, tt.wantWaits, waits) + require.Equal(t, int32(tt.wantRequests), requestCount.Load()) + require.Equal(t, tt.wantPersisted, persisted.Load()) + require.False(t, changedBeforeSuccess.Load()) + + close(requests) + for request := range requests { + require.NoError(t, request.parseErr) + require.Equal(t, http.MethodPost, request.method) + require.Equal(t, "/auth/oauth/token", request.path) + require.Equal(t, oauth.FormMediaType, request.contentType) + require.Equal(t, oauth.JSONMediaType, request.accept) + require.Equal(t, url.Values{ + "grant_type": {"urn:ietf:params:oauth:grant-type:device_code"}, + "device_code": {"device-code"}, + "client_id": {authenticator.ClientID}, + "token_response_version": {"standard"}, + }, request.form) + } + }) + } +} + +func TestGetAccessTokenForDeviceCancellation(t *testing.T) { + t.Parallel() + // Cancellation must stop both a pending wait and an already in-flight token request without leaking a token. + + t.Run("between polls", func(t *testing.T) { + t.Parallel() + + var requestCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + requestCount.Add(1) + })) + t.Cleanup(server.Close) + + mockClock := clock.NewMock() + authenticator := newTestAuthenticator(t, server, mockClock) + ctx, cancel := context.WithCancel(context.Background()) + authenticator.waitForNextPoll = func(ctx context.Context, _ time.Duration) error { + cancel() + <-ctx.Done() + return ctx.Err() + } + + token, err := authenticator.GetAccessTokenForDevice(ctx, &DeviceVerification{ + DeviceCode: "device-code", + CheckInterval: time.Minute, + ExpiresAt: mockClock.Now().Add(time.Hour), + }) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, token) + require.Zero(t, requestCount.Load()) + }) + + t.Run("in flight", func(t *testing.T) { + t.Parallel() + + requestStarted := make(chan struct{}) + releaseRequest := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + close(requestStarted) + <-releaseRequest + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(releaseRequest) }) + + mockClock := clock.NewMock() + authenticator := newTestAuthenticator(t, server, mockClock) + authenticator.waitForNextPoll = func(_ context.Context, interval time.Duration) error { + mockClock.Add(interval) + return nil + } + ctx, cancel := context.WithCancel(context.Background()) + type result struct { + token *oauth.TokenResponse + err error + } + resultCh := make(chan result, 1) + go func() { + token, err := authenticator.GetAccessTokenForDevice(ctx, &DeviceVerification{ + DeviceCode: "device-code", + CheckInterval: time.Second, + ExpiresAt: mockClock.Now().Add(time.Hour), + }) + resultCh <- result{token: token, err: err} + }() + + <-requestStarted + cancel() + select { + case result := <-resultCh: + require.ErrorIs(t, result.err, context.Canceled) + require.Nil(t, result.token) + case <-time.After(2 * time.Second): + t.Fatal("token request did not stop after cancellation") + } + }) +} + +func newTestAuthenticator(t *testing.T, server *httptest.Server, mockClock *clock.Mock) *DeviceAuthenticator { + t.Helper() + + authenticator, err := New(server.URL + "/") + require.NoError(t, err) + authenticator.client = server.Client() + authenticator.Clock = mockClock + return authenticator +} + +func snapshotRequest(r *http.Request) requestSnapshot { + err := r.ParseForm() + return requestSnapshot{ + method: r.Method, + path: r.URL.Path, + contentType: r.Header.Get("Content-Type"), + accept: r.Header.Get("Accept"), + form: r.PostForm, + parseErr: err, + } +} diff --git a/cli/pkg/pkce/authenticator.go b/cli/pkg/pkce/authenticator.go index 122f109cd0f8..121aba2c031f 100644 --- a/cli/pkg/pkce/authenticator.go +++ b/cli/pkg/pkce/authenticator.go @@ -1,10 +1,12 @@ package pkce import ( + "context" "crypto/rand" "crypto/sha256" "encoding/base64" "encoding/json" + "errors" "fmt" "io" "math/big" @@ -78,8 +80,14 @@ func (a *Authenticator) GetAuthURL(state string) string { } func (a *Authenticator) ExchangeCodeForToken(code string) (string, error) { + return a.ExchangeCodeForTokenContext(context.Background(), code) +} + +// ExchangeCodeForTokenContext lets a caller tie the token request to its callback +// lifecycle, so cancelling the callback also cancels an in-flight exchange. +func (a *Authenticator) ExchangeCodeForTokenContext(ctx context.Context, code string) (string, error) { // Create the token request - req, err := tokenRequest(a.baseAuthURL, code, a.clientID, a.redirectURL, a.codeVerifier) + req, err := tokenRequestWithContext(ctx, a.baseAuthURL, code, a.clientID, a.redirectURL, a.codeVerifier) if err != nil { return "", err } @@ -103,12 +111,20 @@ func (a *Authenticator) ExchangeCodeForToken(code string) (string, error) { if err := json.NewDecoder(resp.Body).Decode(&tokenResponse); err != nil { return "", err } + if tokenResponse.AccessToken == "" { + return "", errors.New("token response is missing access_token") + } - // Return the access token + // Return a token only after the whole exchange succeeds. Callers that persist + // this value can leave existing credentials untouched on every error path. return tokenResponse.AccessToken, nil } func tokenRequest(baseAuthURL, code, clientID, redirectURI, codeVerifier string) (*http.Request, error) { + return tokenRequestWithContext(context.Background(), baseAuthURL, code, clientID, redirectURI, codeVerifier) +} + +func tokenRequestWithContext(ctx context.Context, baseAuthURL, code, clientID, redirectURI, codeVerifier string) (*http.Request, error) { tokenURL := fmt.Sprintf("%s/auth/oauth/token", baseAuthURL) payload := url.Values{ "grant_type": []string{"authorization_code"}, @@ -118,7 +134,8 @@ func tokenRequest(baseAuthURL, code, clientID, redirectURI, codeVerifier string) "code_verifier": []string{codeVerifier}, "token_response_version": []string{"standard"}, // For backward compatibility with older Rill CLI, see utils.go in oauth pkg } - req, err := http.NewRequest( + req, err := http.NewRequestWithContext( + ctx, http.MethodPost, tokenURL, strings.NewReader(payload.Encode()), diff --git a/cli/pkg/pkce/authenticator_test.go b/cli/pkg/pkce/authenticator_test.go index ff2326616674..59c3fca9e550 100644 --- a/cli/pkg/pkce/authenticator_test.go +++ b/cli/pkg/pkce/authenticator_test.go @@ -1,11 +1,23 @@ package pkce import ( - "github.com/stretchr/testify/require" + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync/atomic" "testing" + "time" + + "github.com/rilldata/rill/admin/pkg/oauth" + "github.com/stretchr/testify/require" ) -func Test_generateCodeVerifier(t *testing.T) { +func TestGenerateCodeVerifier(t *testing.T) { + t.Parallel() + // PKCE verifiers must satisfy RFC 7636 length and alphabet constraints on every generation. + for i := 0; i < 1000; i++ { code, err := generateCodeVerifier() require.NoError(t, err) @@ -18,3 +30,193 @@ func Test_generateCodeVerifier(t *testing.T) { } } } + +func TestGetAuthURLContract(t *testing.T) { + t.Parallel() + // The browser URL is the CLI/server contract, including encoded state and the S256 challenge. + + authenticator := &Authenticator{ + baseAuthURL: "https://auth.example", + redirectURL: "http://localhost:9009/auth/callback?open=true", + codeVerifier: "0123456789abcdefghijklmnopqrstuvwxyz-._~ABCDEFG", + clientID: "rill-local", + } + authURL, err := url.Parse(authenticator.GetAuthURL("state with spaces")) + require.NoError(t, err) + require.Equal(t, "https", authURL.Scheme) + require.Equal(t, "auth.example", authURL.Host) + require.Equal(t, "/auth/oauth/authorize", authURL.Path) + require.Equal(t, url.Values{ + "client_id": {"rill-local"}, + "redirect_uri": {"http://localhost:9009/auth/callback?open=true"}, + "response_type": {"code"}, + "code_challenge": {createCodeChallenge(authenticator.codeVerifier)}, + "code_challenge_method": {codeChallengeMethod}, + "state": {"state with spaces"}, + }, authURL.Query()) +} + +func TestExchangeCodeForTokenContract(t *testing.T) { + t.Parallel() + // Exercise successful and malformed token responses while proving failed exchanges do not replace credentials. + + tests := []struct { + name string + status int + body string + wantToken string + wantErrText string + wantPersisted string + }{ + { + name: "success", + status: http.StatusOK, + body: `{"access_token":"new-token","token_type":"Bearer"}`, + wantToken: "new-token", + wantPersisted: "new-token", + }, + { + name: "malformed JSON", + status: http.StatusOK, + body: `{"access_token":`, + wantErrText: "unexpected EOF", + wantPersisted: "old-token", + }, + { + name: "empty token response", + status: http.StatusOK, + body: `{}`, + wantErrText: "missing access_token", + wantPersisted: "old-token", + }, + { + name: "server error", + status: http.StatusInternalServerError, + body: "backend unavailable", + wantErrText: "unexpected status code: 500", + wantPersisted: "old-token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var persisted atomic.Value + persisted.Store("old-token") + var changedBeforeSuccess atomic.Bool + requests := make(chan pkceRequestSnapshot, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if persisted.Load().(string) != "old-token" { + changedBeforeSuccess.Store(true) + } + requests <- snapshotPKCERequest(r) + w.WriteHeader(tt.status) + _, _ = io.WriteString(w, tt.body) + })) + t.Cleanup(server.Close) + + authenticator, err := NewAuthenticator( + server.URL, + "http://localhost:9009/auth/callback?open=true", + "rill-local", + "/project", + ) + require.NoError(t, err) + authenticator.client = server.Client() + authenticator.codeVerifier = "0123456789abcdefghijklmnopqrstuvwxyz-._~ABCDEFG" + + token, err := authenticator.ExchangeCodeForToken("authorization-code") + // This mirrors the callback persistence boundary: the old credential is + // retained unless a complete exchange returns a non-empty token. + if err == nil { + persisted.Store(token) + } + + if tt.wantErrText == "" { + require.NoError(t, err) + require.Equal(t, tt.wantToken, token) + } else { + require.ErrorContains(t, err, tt.wantErrText) + require.Empty(t, token) + } + require.Equal(t, tt.wantPersisted, persisted.Load()) + require.False(t, changedBeforeSuccess.Load()) + + request := <-requests + require.NoError(t, request.parseErr) + require.Equal(t, http.MethodPost, request.method) + require.Equal(t, "/auth/oauth/token", request.path) + require.Equal(t, oauth.FormMediaType, request.contentType) + require.Equal(t, oauth.JSONMediaType, request.accept) + require.Equal(t, url.Values{ + "grant_type": {"authorization_code"}, + "code": {"authorization-code"}, + "client_id": {"rill-local"}, + "redirect_uri": {"http://localhost:9009/auth/callback?open=true"}, + "code_verifier": {authenticator.codeVerifier}, + "token_response_version": {"standard"}, + }, request.form) + }) + } +} + +func TestExchangeCodeForTokenCancellation(t *testing.T) { + t.Parallel() + // Canceling an in-flight HTTP exchange must promptly return without producing a token. + + requestStarted := make(chan struct{}) + releaseRequest := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + close(requestStarted) + <-releaseRequest + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(releaseRequest) }) + + authenticator, err := NewAuthenticator(server.URL, "http://localhost/callback", "rill-local", "/") + require.NoError(t, err) + authenticator.client = server.Client() + + ctx, cancel := context.WithCancel(context.Background()) + type result struct { + token string + err error + } + resultCh := make(chan result, 1) + go func() { + token, err := authenticator.ExchangeCodeForTokenContext(ctx, "authorization-code") + resultCh <- result{token: token, err: err} + }() + + <-requestStarted + cancel() + select { + case result := <-resultCh: + require.ErrorIs(t, result.err, context.Canceled) + require.Empty(t, result.token) + case <-time.After(2 * time.Second): + t.Fatal("PKCE token exchange did not stop after cancellation") + } +} + +type pkceRequestSnapshot struct { + method string + path string + contentType string + accept string + form url.Values + parseErr error +} + +func snapshotPKCERequest(r *http.Request) pkceRequestSnapshot { + err := r.ParseForm() + return pkceRequestSnapshot{ + method: r.Method, + path: r.URL.Path, + contentType: r.Header.Get("Content-Type"), + accept: r.Header.Get("Accept"), + form: r.PostForm, + parseErr: err, + } +} From dfc12a3b38de286fb65977e671287e48b0cbca02 Mon Sep 17 00:00:00 2001 From: Nishant Bangarwa Date: Thu, 23 Jul 2026 10:26:56 +0530 Subject: [PATCH 2/2] chore: remove superseded OAuth helpers --- admin/server/auth/mcp_oauth.go | 5 ----- cli/pkg/pkce/authenticator.go | 4 ---- 2 files changed, 9 deletions(-) diff --git a/admin/server/auth/mcp_oauth.go b/admin/server/auth/mcp_oauth.go index 05736e76a9e2..65a00710ce1f 100644 --- a/admin/server/auth/mcp_oauth.go +++ b/admin/server/auth/mcp_oauth.go @@ -180,11 +180,6 @@ func (a *Authenticator) handleOAuthRegister(w http.ResponseWriter, r *http.Reque a.logger.Info("Registered new OAuth client", zap.String("client_id", client.ID), zap.String("client_name", displayName)) } -// remove extra spaces from space separated scope string -func sanitizeScope(scope string) string { - return strings.Join(strings.Fields(scope), " ") -} - // trims white spaces func sanitizeGrantTypes(grants []string) []string { var sanitized []string diff --git a/cli/pkg/pkce/authenticator.go b/cli/pkg/pkce/authenticator.go index 121aba2c031f..101ebb81a444 100644 --- a/cli/pkg/pkce/authenticator.go +++ b/cli/pkg/pkce/authenticator.go @@ -120,10 +120,6 @@ func (a *Authenticator) ExchangeCodeForTokenContext(ctx context.Context, code st return tokenResponse.AccessToken, nil } -func tokenRequest(baseAuthURL, code, clientID, redirectURI, codeVerifier string) (*http.Request, error) { - return tokenRequestWithContext(context.Background(), baseAuthURL, code, clientID, redirectURI, codeVerifier) -} - func tokenRequestWithContext(ctx context.Context, baseAuthURL, code, clientID, redirectURI, codeVerifier string) (*http.Request, error) { tokenURL := fmt.Sprintf("%s/auth/oauth/token", baseAuthURL) payload := url.Values{