Repository: supabase/gotrue
Pull request #1190 (2023-07-20 09:06:45)

feat: serialized access to session in `refresh_token` grant
Improves the `refresh_token` grant flow by serializing access over the session ID, so that a session cannot concurrently be refreshed.

It achieves this by adding a boolean `forUpdate` parameter to `models.FindSessionByID()` and `models.FindUserWithRefreshToken()`. This in turn uses a [`SELECT ... FOR UPDATE`](https://www.postgresql.org/docs/current/sql-select.html#SQL-FOR-UPDATE-SHARE) query that locks the row from use with other flows that select it with a `FOR UPDATE` clause.


Patch:
diff --git a/internal/api/auth.go b/internal/api/auth.go
index 4ac0a700..8e904dce 100644
--- a/internal/api/auth.go
+++ b/internal/api/auth.go
@@ -113,7 +113,7 @@ func (a *API) maybeLoadUserOrSession(ctx context.Context) (context.Context, erro
 		if err != nil {
 			return ctx, err
 		}
-		session, err = models.FindSessionByID(db, sessionId)
+		session, err = models.FindSessionByID(db, sessionId, false)
 		if err != nil && !models.IsNotFoundError(err) {
 			return ctx, err
 		}
diff --git a/internal/api/external.go b/internal/api/external.go
index c9fce207..1f13983b 100644
--- a/internal/api/external.go
+++ b/internal/api/external.go
@@ -30,12 +30,6 @@ type ExternalProviderClaims struct {
 	FlowStateID string `json:"flow_state_id"`
 }
 
-// ExternalSignupParams are the parameters the Signup endpoint accepts
-type ExternalSignupParams struct {
-	Provider string `json:"provider"`
-	Code     string `json:"code"`
-}
-
 // ExternalProviderRedirect redirects the request to the corresponding oauth provider
 func (a *API) ExternalProviderRedirect(w http.ResponseWriter, r *http.Request) error {
 	ctx := r.Context()
diff --git a/internal/api/pkce.go b/internal/api/pkce.go
index 167eb4f0..ca926cca 100644
--- a/internal/api/pkce.go
+++ b/internal/api/pkce.go
@@ -9,11 +9,9 @@ import (
 )
 
 const (
-	PKCE                          = "pkce"
 	PKCEPrefix                    = "pkce_"
 	MinCodeChallengeLength        = 43
 	MaxCodeChallengeLength        = 128
-	InvalidFlowTypeErrorMessage   = "Invalid flow type. Flow Type must be either implicit or pkce"
 	InvalidPKCEParamsErrorMessage = "PKCE flow requires code_challenge_method and code_challenge"
 )
 
diff --git a/internal/api/token.go b/internal/api/token.go
index 3521f22f..77abf34d 100644
--- a/internal/api/token.go
+++ b/internal/api/token.go
@@ -251,7 +251,7 @@ func generateAccessToken(tx *storage.Connection, user *models.User, sessionId *u
 	sid := ""
 	if sessionId != nil {
 		sid = sessionId.String()
-		session, terr := models.FindSessionByID(tx, *sessionId)
+		session, terr := models.FindSessionByID(tx, *sessionId, false)
 		if terr != nil {
 			return "", terr
 		}
@@ -311,7 +311,7 @@ func (a *API) issueRefreshToken(ctx context.Context, conn *storage.Connection, u
 			return internalServerError("Database error granting user").WithInternalError(terr)
 		}
 
-		session, terr := models.FindSessionByID(tx, *refreshToken.SessionId)
+		session, terr := models.FindSessionByID(tx, *refreshToken.SessionId, false)
 		if terr != nil {
 			return terr
 		}
@@ -350,7 +350,7 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 		return nil, internalServerError("Cannot read SessionId claim as UUID").WithInternalError(err)
 	}
 	err = tx.Transaction(func(tx *storage.Connection) error {
-		session, terr := models.FindSessionByID(tx, sessionId)
+		session, terr := models.FindSessionByID(tx, sessionId, false)
 		if terr != nil {
 			return terr
 		}
@@ -358,7 +358,7 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 		if terr != nil {
 			return terr
 		}
-		session, terr = models.FindSessionByID(tx, sessionId)
+		session, terr = models.FindSessionByID(tx, sessionId, false)
 		if terr != nil {
 			return terr
 		}
diff --git a/internal/api/token_refresh.go b/internal/api/token_refresh.go
index 2416037d..eca28210 100644
--- a/internal/api/token_refresh.go
+++ b/internal/api/token_refresh.go
@@ -36,7 +36,7 @@ func (a *API) RefreshTokenGrant(ctx context.Context, w http.ResponseWriter, r *h
 		return oauthError("invalid_request", "refresh_token required")
 	}
 
-	user, token, session, err := models.FindUserWithRefreshToken(db, params.RefreshToken)
+	user, _, session, err := models.FindUserWithRefreshToken(db, params.RefreshToken, false)
 	if err != nil {
 		if models.IsNotFoundError(err) {
 			return oauthError("invalid_grant", "Invalid Refresh Token: Refresh Token Not Found")
@@ -60,39 +60,51 @@ func (a *API) RefreshTokenGrant(ctx context.Context, w http.ResponseWriter, r *h
 		}
 	}
 
-	if token.Revoked {
-		a.clearCookieTokens(config, w)
-		// For a revoked refresh token to be reused, it has to fall within the reuse interval.
+	// basic checks above passed, now we need to serialize access to the
+	// session in a transaction so that there's no concurrent modification
 
-		reuseUntil := token.UpdatedAt.Add(
-			time.Second * time.Duration(config.Security.RefreshTokenReuseInterval))
+	var tokenString string
+	var newTokenResponse *AccessTokenResponse
 
-		if time.Now().After(reuseUntil) {
-			// not OK to reuse this token
+	err = db.Transaction(func(tx *storage.Connection) error {
+		user, token, _, terr := models.FindUserWithRefreshToken(db, params.RefreshToken, true /* forUpdate */)
+		if terr != nil {
+			if models.IsNotFoundError(terr) {
+				return oauthError("invalid_grant", "Invalid Refresh Token: Refresh Token Not Found")
+			}
+			return internalServerError(terr.Error())
+		}
 
-			if config.Security.RefreshTokenRotationEnabled {
-				// Revoke all tokens in token family
-				err = db.Transaction(func(tx *storage.Connection) error {
-					var terr error
-					if terr = models.RevokeTokenFamily(tx, token); terr != nil {
-						return terr
+		// refresh token row and session are locked, cannot be concurrently refreshed
+
+		if token.Revoked {
+			a.clearCookieTokens(config, w)
+			// For a revoked refresh token to be reused, it has to fall within the reuse interval.
+
+			reuseUntil := token.UpdatedAt.Add(
+				time.Second * time.Duration(config.Security.RefreshTokenReuseInterval))
+
+			if time.Now().After(reuseUntil) {
+				// not OK to reuse this token
+
+				if config.Security.RefreshTokenRotationEnabled {
+					// Revoke all tokens in token family
+					err = db.Transaction(func(tx *storage.Connection) error {
+						var terr error
+						if terr = models.RevokeTokenFamily(tx, token); terr != nil {
+							return terr
+						}
+						return nil
+					})
+					if err != nil {
+						return internalServerError(err.Error())
 					}
-					return nil
-				})
-				if err != nil {
-					return internalServerError(err.Error())
 				}
-			}
 
-			return oauthError("invalid_grant", "Invalid Refresh Token: Already Used").WithInternalMessage("Possible abuse attempt: %v", token.ID)
+				return oauthError("invalid_grant", "Invalid Refresh Token: Already Used").WithInternalMessage("Possible abuse attempt: %v", token.ID)
+			}
 		}
-	}
-
-	var tokenString string
-	var newTokenResponse *AccessTokenResponse
 
-	err = db.Transaction(func(tx *storage.Connection) error {
-		var terr error
 		if terr = models.NewAuditLogEntry(r, tx, user, models.TokenRefreshedAction, "", nil); terr != nil {
 			return terr
 		}
diff --git a/internal/crypto/crypto.go b/internal/crypto/crypto.go
index 2b4fb96c..70980c16 100644
--- a/internal/crypto/crypto.go
+++ b/internal/crypto/crypto.go
@@ -37,22 +37,3 @@ func GenerateOtp(digits int) (string, error) {
 	otp := fmt.Sprintf(expr, val.String())
 	return otp, nil
 }
-
-// generateOtpFromCharset generates a random n-length otp from a charset
-func generateOtpFromCharset(length int, charset string) (string, error) {
-	b := make([]byte, length)
-	for i := range b {
-		val, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
-		if err != nil {
-			return "", errors.WithMessage(err, "Error generating otp from charset")
-		}
-		b[i] = charset[val.Int64()]
-	}
-	return string(b), nil
-}
-
-// GenerateNanoId generates a random n-length alphanumeric otp
-func GenerateNanoId(length int) (string, error) {
-	const charset = "0123456789abcdefghijklmnopqrstuvwxyz"
-	return generateOtpFromCharset(length, charset)
-}
diff --git a/internal/models/refresh_token.go b/internal/models/refresh_token.go
index 50168100..3a15b0b0 100644
--- a/internal/models/refresh_token.go
+++ b/internal/models/refresh_token.go
@@ -89,19 +89,6 @@ func RevokeTokenFamily(tx *storage.Connection, token *RefreshToken) error {
 	return nil
 }
 
-// GetValidChildToken returns the child token of the token provided if the child is not revoked.
-func GetValidChildToken(tx *storage.Connection, token *RefreshToken) (*RefreshToken, error) {
-	refreshToken := &RefreshToken{}
-	err := tx.Q().Where("parent = ? and revoked = false", token.Token).First(refreshToken)
-	if err != nil {
-		if errors.Cause(err) == sql.ErrNoRows {
-			return nil, RefreshTokenNotFoundError{}
-		}
-		return nil, err
-	}
-	return refreshToken, nil
-}
-
 func FindTokenBySessionID(tx *storage.Connection, sessionId *uuid.UUID) (*RefreshToken, error) {
 	refreshToken := &RefreshToken{}
 	err := tx.Q().Where("session_id = ?", sessionId).Order("created_at asc").First(refreshToken)
diff --git a/internal/models/sessions.go b/internal/models/sessions.go
index ad6fa1e2..264fae76 100644
--- a/internal/models/sessions.go
+++ b/internal/models/sessions.go
@@ -2,6 +2,7 @@ package models
 
 import (
 	"database/sql"
+	"fmt"
 	"sort"
 	"strings"
 	"time"
@@ -85,8 +86,22 @@ func NewSession() (*Session, error) {
 	return session, nil
 }
 
-func FindSessionByID(tx *storage.Connection, id uuid.UUID) (*Session, error) {
+func FindSessionByID(tx *storage.Connection, id uuid.UUID, forUpdate bool) (*Session, error) {
 	session := &Session{}
+
+	if forUpdate {
+		// pop does not provide us with a way to execute FOR UPDATE
+		// queries which lock the rows affected by the query from
+		// being accessed by any other transaction that also uses FOR
+		// UPDATE
+		if err := tx.RawQuery(fmt.Sprintf("SELECT id FROM %q WHERE id = ? LIMIT 1 FOR UPDATE;", session.TableName()), id).Exec(); err != nil {
+			if errors.Cause(err) == sql.ErrNoRows {
+				return nil, SessionNotFoundError{}
+			}
+		}
+	}
+
+	// once the rows are locked (if forUpdate was true), we can query again using pop
 	if err := tx.Eager().Q().Where("id = ?", id).First(session); err != nil {
 		if errors.Cause(err) == sql.ErrNoRows {
 			return nil, SessionNotFoundError{}
diff --git a/internal/models/user.go b/internal/models/user.go
index aafb570f..2bab520a 100644
--- a/internal/models/user.go
+++ b/internal/models/user.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"database/sql"
 	"encoding/base64"
+	"fmt"
 	"strings"
 	"time"
 
@@ -435,8 +436,24 @@ func FindUserByEmailChangeToken(tx *storage.Connection, token string) (*User, er
 }
 
 // FindUserWithRefreshToken finds a user from the provided refresh token.
-func FindUserWithRefreshToken(tx *storage.Connection, token string) (*User, *RefreshToken, *Session, error) {
+func FindUserWithRefreshToken(tx *storage.Connection, token string, forUpdate bool) (*User, *RefreshToken, *Session, error) {
 	refreshToken := &RefreshToken{}
+
+	if forUpdate {
+		// pop does not provide us with a way to execute FOR UPDATE
+		// queries which lock the rows affected by the query from
+		// being accessed by any other transaction that also uses FOR
+		// UPDATE
+		if err := tx.RawQuery(fmt.Sprintf("SELECT id FROM %q WHERE token = ? LIMIT 1 FOR UPDATE;", refreshToken.TableName()), token).Exec(); err != nil {
+			if errors.Cause(err) == sql.ErrNoRows {
+				return nil, nil, nil, RefreshTokenNotFoundError{}
+			}
+
+			return nil, nil, nil, errors.Wrap(err, "error finding refresh token for update")
+		}
+	}
+
+	// once the rows are locked (if forUpdate was true), we can query again using pop
 	if err := tx.Where("token = ?", token).First(refreshToken); err != nil {
 		if errors.Cause(err) == sql.ErrNoRows {
 			return nil, nil, nil, RefreshTokenNotFoundError{}
@@ -455,7 +472,7 @@ func FindUserWithRefreshToken(tx *storage.Connection, token string) (*User, *Ref
 		sessionId := *refreshToken.SessionId
 
 		if sessionId != uuid.Nil {
-			session, err = FindSessionByID(tx, sessionId)
+			session, err = FindSessionByID(tx, sessionId, forUpdate)
 			if err != nil {
 				if !IsNotFoundError(err) {
 					return nil, nil, nil, errors.Wrap(err, "error finding session from refresh token")


Test patch:
diff --git a/internal/api/mfa_test.go b/internal/api/mfa_test.go
index 52c07c7e..4df987db 100644
--- a/internal/api/mfa_test.go
+++ b/internal/api/mfa_test.go
@@ -258,7 +258,7 @@ func (ts *MFATestSuite) TestMFAVerifyFactor() {
 
 			if v.expectedHTTPCode == http.StatusOK {
 				// Ensure alternate session has been deleted
-				_, err = models.FindSessionByID(ts.API.db, secondarySession.ID)
+				_, err = models.FindSessionByID(ts.API.db, secondarySession.ID, false)
 				require.EqualError(ts.T(), err, models.SessionNotFoundError{}.Error())
 			}
 			if !v.validChallenge {
@@ -330,7 +330,7 @@ func (ts *MFATestSuite) TestUnenrollVerifiedFactor() {
 			if v.expectedHTTPCode == http.StatusOK {
 				_, err = models.FindFactorByFactorID(ts.API.db, f.ID)
 				require.EqualError(ts.T(), err, models.FactorNotFoundError{}.Error())
-				session, _ := models.FindSessionByID(ts.API.db, secondarySession.ID)
+				session, _ := models.FindSessionByID(ts.API.db, secondarySession.ID, false)
 				require.Equal(ts.T(), models.AAL1.String(), session.GetAAL())
 				require.Nil(ts.T(), session.FactorID)
 
@@ -373,7 +373,7 @@ func (ts *MFATestSuite) TestUnenrollUnverifiedFactor() {
 	require.Equal(ts.T(), http.StatusOK, w.Code)
 	_, err = models.FindFactorByFactorID(ts.API.db, f.ID)
 	require.EqualError(ts.T(), err, models.FactorNotFoundError{}.Error())
-	session, _ := models.FindSessionByID(ts.API.db, secondarySession.ID)
+	session, _ := models.FindSessionByID(ts.API.db, secondarySession.ID, false)
 	require.Equal(ts.T(), models.AAL1.String(), session.GetAAL())
 	require.Nil(ts.T(), session.FactorID)
 
diff --git a/internal/api/user_test.go b/internal/api/user_test.go
index c36d9cf6..994b1b09 100644
--- a/internal/api/user_test.go
+++ b/internal/api/user_test.go
@@ -209,7 +209,7 @@ func (ts *UserTestSuite) TestUserUpdatePassword() {
 	require.NoError(ts.T(), err)
 
 	// create a session and modify it's created_at time to simulate a session that is not recently logged in
-	notRecentlyLoggedIn, err := models.FindSessionByID(ts.API.db, *r2.SessionId)
+	notRecentlyLoggedIn, err := models.FindSessionByID(ts.API.db, *r2.SessionId, true)
 	require.NoError(ts.T(), err)
 
 	// cannot use Update here because Update doesn't removes the created_at field
diff --git a/internal/models/refresh_token_test.go b/internal/models/refresh_token_test.go
index d1875da3..b255ce12 100644
--- a/internal/models/refresh_token_test.go
+++ b/internal/models/refresh_token_test.go
@@ -52,7 +52,7 @@ func (ts *RefreshTokenTestSuite) TestGrantRefreshTokenSwap() {
 	s, err := GrantRefreshTokenSwap(&http.Request{}, ts.db, u, r)
 	require.NoError(ts.T(), err)
 
-	_, nr, _, err := FindUserWithRefreshToken(ts.db, r.Token)
+	_, nr, _, err := FindUserWithRefreshToken(ts.db, r.Token, false)
 	require.NoError(ts.T(), err)
 
 	require.Equal(ts.T(), r.ID, nr.ID)
@@ -68,7 +68,7 @@ func (ts *RefreshTokenTestSuite) TestLogout() {
 	require.NoError(ts.T(), err)
 
 	require.NoError(ts.T(), Logout(ts.db, u.ID))
-	u, r, _, err = FindUserWithRefreshToken(ts.db, r.Token)
+	u, r, _, err = FindUserWithRefreshToken(ts.db, r.Token, false)
 	require.Errorf(ts.T(), err, "expected error when there are no refresh tokens to authenticate. user: %v token: %v", u, r)
 
 	require.True(ts.T(), IsNotFoundError(err), "expected NotFoundError")
diff --git a/internal/models/sessions_test.go b/internal/models/sessions_test.go
index efc30255..f7e9d7b8 100644
--- a/internal/models/sessions_test.go
+++ b/internal/models/sessions_test.go
@@ -1,13 +1,14 @@
 package models
 
 import (
+	"testing"
+	"time"
+
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"github.com/supabase/gotrue/internal/conf"
 	"github.com/supabase/gotrue/internal/storage"
 	"github.com/supabase/gotrue/internal/storage/test"
-	"testing"
-	"time"
 )
 
 type SessionsTestSuite struct {
@@ -39,6 +40,20 @@ func TestSession(t *testing.T) {
 	suite.Run(t, ts)
 }
 
+func (ts *SessionsTestSuite) TestFindBySessionIDWithForUpdate() {
+	u, err := FindUserByEmailAndAudience(ts.db, "test@example.com", ts.Config.JWT.Aud)
+	require.NoError(ts.T(), err)
+	session, err := NewSession()
+	require.NoError(ts.T(), err)
+	session.UserID = u.ID
+	require.NoError(ts.T(), ts.db.Create(session))
+
+	found, err := FindSessionByID(ts.db, session.ID, true)
+	require.NoError(ts.T(), err)
+
+	require.Equal(ts.T(), session.ID, found.ID)
+}
+
 func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	totalDistinctClaims := 2
 	u, err := FindUserByEmailAndAudience(ts.db, "test@example.com", ts.Config.JWT.Aud)
@@ -54,7 +69,7 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	firstClaimAddedTime := time.Now()
 	err = AddClaimToSession(ts.db, session, TOTPSignIn)
 	require.NoError(ts.T(), err)
-	session, err = FindSessionByID(ts.db, session.ID)
+	session, err = FindSessionByID(ts.db, session.ID, false)
 	require.NoError(ts.T(), err)
 
 	aal, amr, err := session.CalculateAALAndAMR(ts.db)
@@ -65,7 +80,7 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	err = AddClaimToSession(ts.db, session, TOTPSignIn)
 	require.NoError(ts.T(), err)
 
-	session, err = FindSessionByID(ts.db, session.ID)
+	session, err = FindSessionByID(ts.db, session.ID, false)
 	require.NoError(ts.T(), err)
 
 	aal, amr, err = session.CalculateAALAndAMR(ts.db)
@@ -81,5 +96,4 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 		}
 	}
 	require.True(ts.T(), found)
-
 }
diff --git a/internal/models/user_test.go b/internal/models/user_test.go
index 0949a65e..564e6fd8 100644
--- a/internal/models/user_test.go
+++ b/internal/models/user_test.go
@@ -151,7 +151,7 @@ func (ts *UserTestSuite) TestFindUserWithRefreshToken() {
 	r, err := GrantAuthenticatedUser(ts.db, u, GrantParams{})
 	require.NoError(ts.T(), err)
 
-	n, nr, s, err := FindUserWithRefreshToken(ts.db, r.Token)
+	n, nr, s, err := FindUserWithRefreshToken(ts.db, r.Token, true /* forUpdate */)
 	require.NoError(ts.T(), err)
 	require.Equal(ts.T(), r.ID, nr.ID)
 	require.Equal(ts.T(), u.ID, n.ID)