tests: Harden OAuth device, PKCE, and refresh-token flows

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)
 	}
 }

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;

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" {