fix: prevent reuse of flow state (#2483)

When a user has been assigned to a flow state during a PKCE flow,
prevent the reuse of the state.

Author: fadymak

internal/api/apierrors/errorcode.go

@@ -23,6 +23,7 @@ const (
 	ErrorCodeRefreshTokenAlreadyUsed           ErrorCode = "refresh_token_already_used"
 	ErrorCodeFlowStateNotFound                 ErrorCode = "flow_state_not_found"
 	ErrorCodeFlowStateExpired                  ErrorCode = "flow_state_expired"
+	ErrorCodeFlowStateAlreadyUsed              ErrorCode = "flow_state_already_used"
 	ErrorCodeOAuthClientStateNotFound          ErrorCode = "oauth_client_state_not_found"
 	ErrorCodeOAuthClientStateExpired           ErrorCode = "oauth_client_state_expired"
 	ErrorCodeOAuthInvalidState                 ErrorCode = "oauth_invalid_state"

internal/api/external.go

@@ -222,6 +222,13 @@ func (a *API) internalExternalProviderCallback(w http.ResponseWriter, r *http.Re
 		}
 		if flowState != nil && flowState.IsPKCE() {
 			// PKCE flow: update flow state with user ID and tokens
+			// Re-fetch with FOR UPDATE lock inside the transaction to prevent concurrent claims
+			if flowState, terr = models.FindFlowStateByIDForUpdate(tx, flowState.ID.String()); terr != nil {
+				return terr
+			}
+			if flowState.UserID != nil {
+				return apierrors.NewBadRequestError(apierrors.ErrorCodeFlowStateAlreadyUsed, "State has already been used")
+			}
 			flowState.ProviderAccessToken = providerAccessToken
 			flowState.ProviderRefreshToken = providerRefreshToken
 			flowState.UserID = &(user.ID)
@@ -543,6 +550,11 @@ func (a *API) loadExternalStateFromUUID(ctx context.Context, db *storage.Connect
 		return ctx, apierrors.NewBadRequestError(apierrors.ErrorCodeBadOAuthState, "OAuth state has expired")
 	}
 
+	// UserID is nil at creation and set during callback, so non-nil means already consumed.
+	if flowState.IsPKCE() && flowState.UserID != nil {
+		return ctx, apierrors.NewBadRequestError(apierrors.ErrorCodeFlowStateAlreadyUsed, "State has already been used")
+	}
+
 	ctx = withExternalProviderType(ctx, flowState.ProviderType, flowState.EmailOptional)
 
 	if flowState.InviteToken != nil && *flowState.InviteToken != "" {

internal/api/external_test.go

@@ -1,6 +1,8 @@
 package api
 
 import (
+	"crypto/sha256"
+	"encoding/base64"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -408,3 +410,55 @@ func (ts *ExternalTestSuite) TestOAuthState_InvalidFormat() {
 	// Should redirect to site URL with error since state is invalid
 	ts.Require().Equal(http.StatusSeeOther, w.Code)
 }
+
+// TestPKCEFlowStateReuseRejected verifies that a PKCE flow state cannot be reused
+// after the OAuth callback has been completed
+func (ts *ExternalTestSuite) TestPKCEFlowStateReuseRejected() {
+	code := "authcode"
+	server := setupGenericOAuthServer(ts, code)
+	defer server.Close()
+
+	codeVerifier := "testtesttesttesttesttesttesttesttesttesttesttesttesttest"
+	hashedCodeVerifier := sha256.Sum256([]byte(codeVerifier))
+	codeChallenge := base64.RawURLEncoding.EncodeToString(hashedCodeVerifier[:])
+
+	// Step 1: Initiate PKCE authorization flow and extract the state parameter
+	w := performPKCEAuthorizationRequest(ts, "github", codeChallenge, "s256")
+	ts.Require().Equal(http.StatusFound, w.Code)
+	u, err := url.Parse(w.Header().Get("Location"))
+	ts.Require().NoError(err)
+	state := u.Query().Get("state")
+	ts.Require().NotEmpty(state)
+
+	// Step 2: First callback completes successfully (sets UserID on the flow state)
+	callbackURL, err := url.Parse("http://localhost/callback")
+	ts.Require().NoError(err)
+	v := callbackURL.Query()
+	v.Set("code", code)
+	v.Set("state", state)
+	callbackURL.RawQuery = v.Encode()
+
+	req := httptest.NewRequest(http.MethodGet, callbackURL.String(), nil)
+	w = httptest.NewRecorder()
+	ts.API.handler.ServeHTTP(w, req)
+	ts.Require().Equal(http.StatusFound, w.Code)
+
+	firstRedirect, err := url.Parse(w.Header().Get("Location"))
+	ts.Require().NoError(err)
+	firstQuery, err := url.ParseQuery(firstRedirect.RawQuery)
+	ts.Require().NoError(err)
+	ts.Require().NotEmpty(firstQuery.Get("code"), "first callback should return an auth code")
+
+	// Step 3: Second callback with the same state must be rejected
+	req = httptest.NewRequest(http.MethodGet, callbackURL.String(), nil)
+	w = httptest.NewRecorder()
+	ts.API.handler.ServeHTTP(w, req)
+
+	// The callback redirects errors to the redirect URL with error parameters
+	redirectURL, err := url.Parse(w.Header().Get("Location"))
+	ts.Require().NoError(err)
+	errorQuery, err := url.ParseQuery(redirectURL.RawQuery)
+	ts.Require().NoError(err)
+	ts.Contains(errorQuery.Get("error_description"), "already been used",
+		"second callback with same state should be rejected as already used")
+}

internal/api/samlacs.go

@@ -314,8 +314,15 @@ func (a *API) handleSamlAcs(w http.ResponseWriter, r *http.Request) error {
 
 		if flowState != nil && flowState.IsPKCE() {
 			// PKCE flow: update flow state with user ID
+			// Re-fetch with FOR UPDATE lock inside the transaction to prevent concurrent claims
+			if flowState, terr = models.FindFlowStateByIDForUpdate(tx, flowState.ID.String()); terr != nil {
+				return terr
+			}
+			if flowState.UserID != nil {
+				return apierrors.NewBadRequestError(apierrors.ErrorCodeFlowStateAlreadyUsed, "State has already been used")
+			}
 			flowState.UserID = &(user.ID)
-			if terr := tx.Update(flowState); terr != nil {
+			if terr = tx.Update(flowState); terr != nil {
 				return terr
 			}
 		}

internal/models/flow_state.go

@@ -164,7 +164,27 @@ func FindFlowStateByID(tx *storage.Connection, id string) (*FlowState, error) {
 		}
 		return nil, errors.Wrap(err, "error finding flow state")
 	}
+	return obj, nil
+}
 
+// FindFlowStateByIDForUpdate finds a flow state by ID and locks the row with
+// FOR UPDATE SKIP LOCKED to prevent concurrent modifications. If the row is
+// already locked by another transaction, SKIP LOCKED causes the query to
+// return no rows instead of blocking, which surfaces as FlowStateNotFoundError.
+// The lock is held until the transaction commits or rolls back.
+func FindFlowStateByIDForUpdate(tx *storage.Connection, id string) (*FlowState, error) {
+	obj := &FlowState{}
+	// Pop does not provide a way to execute FOR UPDATE queries,
+	// so we use a raw query to lock the row first.
+	if err := tx.RawQuery(
+		fmt.Sprintf("SELECT * FROM %q WHERE id = ? LIMIT 1 FOR UPDATE SKIP LOCKED", obj.TableName()),
+		id,
+	).First(obj); err != nil {
+		if errors.Cause(err) == sql.ErrNoRows {
+			return nil, FlowStateNotFoundError{}
+		}
+		return nil, errors.Wrap(err, "error finding flow state")
+	}
 	return obj, nil
 }