Repository: supabase/gotrue
Pull request #1446 (2024-02-21 08:08:52)

feat: refactor PKCE FlowState to reduce duplicate code
## What kind of change does this PR introduce?

- Removes `NewFlowStateWithUserID` - it is sufficient to have one method to create a flow state
- Compresses some of the PKCE checks into a single function



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b0581293..a4376b14 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,26 @@
 # Changelog
 
+## [2.144.0](https://github.com/supabase/gotrue/compare/v2.143.0...v2.144.0) (2024-03-04)
+
+
+### Features
+
+* add configuration for custom sms sender hook ([#1428](https://github.com/supabase/gotrue/issues/1428)) ([1ea56b6](https://github.com/supabase/gotrue/commit/1ea56b62d47edb0766d9e445406ecb43d387d920))
+* anonymous sign-ins  ([#1460](https://github.com/supabase/gotrue/issues/1460)) ([130df16](https://github.com/supabase/gotrue/commit/130df165270c69c8e28aaa1b9421342f997c1ff3))
+* clean up test setup in MFA tests ([#1452](https://github.com/supabase/gotrue/issues/1452)) ([7185af8](https://github.com/supabase/gotrue/commit/7185af8de4a269cdde2629054d222333d3522ebe))
+* pass transaction to `invokeHook`, fixing pool exhaustion ([#1465](https://github.com/supabase/gotrue/issues/1465)) ([b536d36](https://github.com/supabase/gotrue/commit/b536d368f35adb31f937169e3f093d28352fa7be))
+* refactor resource owner password grant ([#1443](https://github.com/supabase/gotrue/issues/1443)) ([e63ad6f](https://github.com/supabase/gotrue/commit/e63ad6ff0f67d9a83456918a972ecb5109125628))
+* use dummy instance id to improve performance on refresh token queries ([#1454](https://github.com/supabase/gotrue/issues/1454)) ([656474e](https://github.com/supabase/gotrue/commit/656474e1b9ff3d5129190943e8c48e456625afe5))
+
+
+### Bug Fixes
+
+* expose `provider` under `amr` in access token ([#1456](https://github.com/supabase/gotrue/issues/1456)) ([e9f38e7](https://github.com/supabase/gotrue/commit/e9f38e76d8a7b93c5c2bb0de918a9b156155f018))
+* improve MFA QR Code resilience so as to support providers like 1Password ([#1455](https://github.com/supabase/gotrue/issues/1455)) ([6522780](https://github.com/supabase/gotrue/commit/652278046c9dd92f5cecd778735b058ef3fb41c7))
+* refactor request params to use generics ([#1464](https://github.com/supabase/gotrue/issues/1464)) ([e1cdf5c](https://github.com/supabase/gotrue/commit/e1cdf5c4b5c1bf467094f4bdcaa2e42a5cc51c20))
+* revert refactor resource owner password grant ([#1466](https://github.com/supabase/gotrue/issues/1466)) ([fa21244](https://github.com/supabase/gotrue/commit/fa21244fa929709470c2e1fc4092a9ce947399e7))
+* update file name so migration to Drop IP Address is applied ([#1447](https://github.com/supabase/gotrue/issues/1447)) ([f29e89d](https://github.com/supabase/gotrue/commit/f29e89d7d2c48ee8fd5bf8279a7fa3db0ad4d842))
+
 ## [2.143.0](https://github.com/supabase/gotrue/compare/v2.142.0...v2.143.0) (2024-02-19)
 
 
diff --git a/internal/api/external.go b/internal/api/external.go
index 6982f2c6..cf9da1a0 100644
--- a/internal/api/external.go
+++ b/internal/api/external.go
@@ -79,12 +79,8 @@ func (a *API) GetExternalProviderRedirectURL(w http.ResponseWriter, r *http.Requ
 	flowType := getFlowFromChallenge(codeChallenge)
 
 	flowStateID := ""
-	if flowType == models.PKCEFlow {
-		codeChallengeMethodType, err := models.ParseCodeChallengeMethod(codeChallengeMethod)
-		if err != nil {
-			return "", err
-		}
-		flowState, err := models.NewFlowState(providerType, codeChallenge, codeChallengeMethodType, models.OAuth)
+	if isPKCEFlow(flowType) {
+		flowState, err := generateFlowState(providerType, models.OAuth, codeChallengeMethod, codeChallenge, nil)
 		if err != nil {
 			return "", err
 		}
diff --git a/internal/api/magic_link.go b/internal/api/magic_link.go
index e1b12caa..322cf8b8 100644
--- a/internal/api/magic_link.go
+++ b/internal/api/magic_link.go
@@ -125,18 +125,21 @@ func (a *API) MagicLink(w http.ResponseWriter, r *http.Request) error {
 
 		return sendJSON(w, http.StatusOK, make(map[string]string))
 	}
+	var flowState *models.FlowState
+
+	if isPKCEFlow(flowType) {
+		flowState, err = generateFlowState(models.MagicLink.String(), models.MagicLink, params.CodeChallengeMethod, params.CodeChallenge, &user.ID)
+		if err != nil {
+			return err
+		}
+	}
 
 	err = db.Transaction(func(tx *storage.Connection) error {
 		if terr := models.NewAuditLogEntry(r, tx, user, models.UserRecoveryRequestedAction, "", nil); terr != nil {
 			return terr
 		}
-
-		if isPKCEFlow(flowType) {
-			codeChallengeMethod, terr := models.ParseCodeChallengeMethod(params.CodeChallengeMethod)
-			if terr != nil {
-				return terr
-			}
-			if terr := models.NewFlowStateWithUserID(tx, models.MagicLink.String(), params.CodeChallenge, codeChallengeMethod, models.MagicLink, &user.ID); terr != nil {
+		if flowState != nil {
+			if terr := tx.Create(flowState); terr != nil {
 				return terr
 			}
 		}
diff --git a/internal/api/pkce.go b/internal/api/pkce.go
index c5ed5e54..6b3166df 100644
--- a/internal/api/pkce.go
+++ b/internal/api/pkce.go
@@ -3,6 +3,7 @@ package api
 import (
 	"regexp"
 
+	"github.com/gofrs/uuid"
 	"github.com/supabase/auth/internal/models"
 	"github.com/supabase/auth/internal/storage"
 )
@@ -77,3 +78,14 @@ func getFlowFromChallenge(codeChallenge string) models.FlowType {
 		return models.ImplicitFlow
 	}
 }
+
+// Should only be used with Auth Code of PKCE Flows
+func generateFlowState(providerType string, authenticationMethod models.AuthenticationMethod, codeChallengeMethodParam string, codeChallenge string, userID *uuid.UUID) (*models.FlowState, error) {
+	codeChallengeMethod, err := models.ParseCodeChallengeMethod(codeChallengeMethodParam)
+	if err != nil {
+		return nil, err
+	}
+	flowState := models.NewFlowState(providerType, codeChallenge, codeChallengeMethod, authenticationMethod, userID)
+	return flowState, nil
+
+}
diff --git a/internal/api/recover.go b/internal/api/recover.go
index dcf574d1..80cac6b8 100644
--- a/internal/api/recover.go
+++ b/internal/api/recover.go
@@ -56,6 +56,13 @@ func (a *API) Recover(w http.ResponseWriter, r *http.Request) error {
 		}
 		return internalServerError("Unable to process request").WithInternalError(err)
 	}
+	var flowState *models.FlowState
+	if isPKCEFlow(flowType) {
+		flowState, err = generateFlowState(models.Recovery.String(), models.Recovery, params.CodeChallengeMethod, params.CodeChallenge, &(user.ID))
+		if err != nil {
+			return err
+		}
+	}
 
 	err = db.Transaction(func(tx *storage.Connection) error {
 		if terr := models.NewAuditLogEntry(r, tx, user, models.UserRecoveryRequestedAction, "", nil); terr != nil {
@@ -63,12 +70,8 @@ func (a *API) Recover(w http.ResponseWriter, r *http.Request) error {
 		}
 		mailer := a.Mailer(ctx)
 		referrer := utilities.GetReferrer(r, config)
-		if isPKCEFlow(flowType) {
-			codeChallengeMethod, terr := models.ParseCodeChallengeMethod(params.CodeChallengeMethod)
-			if terr != nil {
-				return terr
-			}
-			if terr := models.NewFlowStateWithUserID(tx, models.Recovery.String(), params.CodeChallenge, codeChallengeMethod, models.Recovery, &(user.ID)); terr != nil {
+		if flowState != nil {
+			if terr := tx.Create(flowState); terr != nil {
 				return terr
 			}
 		}
diff --git a/internal/api/signup.go b/internal/api/signup.go
index 7093fe3b..9bcac4b4 100644
--- a/internal/api/signup.go
+++ b/internal/api/signup.go
@@ -128,16 +128,9 @@ func (a *API) Signup(w http.ResponseWriter, r *http.Request) error {
 		return err
 	}
 
-	var codeChallengeMethod models.CodeChallengeMethod
 	var err error
 	flowType := getFlowFromChallenge(params.CodeChallenge)
 
-	if isPKCEFlow(flowType) {
-		if codeChallengeMethod, err = models.ParseCodeChallengeMethod(params.CodeChallengeMethod); err != nil {
-			return err
-		}
-	}
-
 	var user *models.User
 	var grantParams models.GrantParams
 
@@ -237,8 +230,12 @@ func (a *API) Signup(w http.ResponseWriter, r *http.Request) error {
 				}); terr != nil {
 					return terr
 				}
-				if ok := isPKCEFlow(flowType); ok {
-					if terr := models.NewFlowStateWithUserID(tx, params.Provider, params.CodeChallenge, codeChallengeMethod, models.EmailSignup, &user.ID); terr != nil {
+				if isPKCEFlow(flowType) {
+					flowState, terr := generateFlowState(params.Provider, models.EmailSignup, params.CodeChallengeMethod, params.CodeChallenge, &user.ID)
+					if terr != nil {
+						return terr
+					}
+					if terr := tx.Create(flowState); terr != nil {
 						return terr
 					}
 				}
diff --git a/internal/api/sso.go b/internal/api/sso.go
index 0b4fd890..c1ec8aeb 100644
--- a/internal/api/sso.go
+++ b/internal/api/sso.go
@@ -60,12 +60,8 @@ func (a *API) SingleSignOn(w http.ResponseWriter, r *http.Request) error {
 	flowType := getFlowFromChallenge(params.CodeChallenge)
 	var flowStateID *uuid.UUID
 	flowStateID = nil
-	if flowType == models.PKCEFlow {
-		codeChallengeMethodType, err := models.ParseCodeChallengeMethod(codeChallengeMethod)
-		if err != nil {
-			return err
-		}
-		flowState, err := models.NewFlowState(models.SSOSAML.String(), codeChallenge, codeChallengeMethodType, models.SSOSAML)
+	if isPKCEFlow(flowType) {
+		flowState, err := generateFlowState(models.SSOSAML.String(), models.SSOSAML, codeChallengeMethod, codeChallenge, nil)
 		if err != nil {
 			return err
 		}
diff --git a/internal/api/user.go b/internal/api/user.go
index ddf497f7..e2aaedeb 100644
--- a/internal/api/user.go
+++ b/internal/api/user.go
@@ -197,11 +197,12 @@ func (a *API) UserUpdate(w http.ResponseWriter, r *http.Request) error {
 			referrer := utilities.GetReferrer(r, config)
 			flowType := getFlowFromChallenge(params.CodeChallenge)
 			if isPKCEFlow(flowType) {
-				codeChallengeMethod, terr := models.ParseCodeChallengeMethod(params.CodeChallengeMethod)
+				flowState, terr := generateFlowState(models.EmailChange.String(), models.EmailChange, params.CodeChallengeMethod, params.CodeChallenge, &user.ID)
 				if terr != nil {
 					return terr
 				}
-				if terr := models.NewFlowStateWithUserID(tx, models.EmailChange.String(), params.CodeChallenge, codeChallengeMethod, models.EmailChange, &user.ID); terr != nil {
+
+				if terr := tx.Create(flowState); terr != nil {
 					return terr
 				}
 			}
diff --git a/internal/models/flow_state.go b/internal/models/flow_state.go
index 6aced0b5..04e880ec 100644
--- a/internal/models/flow_state.go
+++ b/internal/models/flow_state.go
@@ -81,21 +81,7 @@ func (FlowState) TableName() string {
 	return tableName
 }
 
-func NewFlowState(providerType, codeChallenge string, codeChallengeMethod CodeChallengeMethod, authenticationMethod AuthenticationMethod) (*FlowState, error) {
-	id := uuid.Must(uuid.NewV4())
-	authCode := uuid.Must(uuid.NewV4())
-	flowState := &FlowState{
-		ID:                   id,
-		ProviderType:         providerType,
-		CodeChallenge:        codeChallenge,
-		CodeChallengeMethod:  codeChallengeMethod.String(),
-		AuthCode:             authCode.String(),
-		AuthenticationMethod: authenticationMethod.String(),
-	}
-	return flowState, nil
-}
-
-func NewFlowStateWithUserID(tx *storage.Connection, providerType, codeChallenge string, codeChallengeMethod CodeChallengeMethod, authenticationMethod AuthenticationMethod, userID *uuid.UUID) error {
+func NewFlowState(providerType, codeChallenge string, codeChallengeMethod CodeChallengeMethod, authenticationMethod AuthenticationMethod, userID *uuid.UUID) *FlowState {
 	id := uuid.Must(uuid.NewV4())
 	authCode := uuid.Must(uuid.NewV4())
 	flowState := &FlowState{
@@ -107,7 +93,7 @@ func NewFlowStateWithUserID(tx *storage.Connection, providerType, codeChallenge
 		AuthenticationMethod: authenticationMethod.String(),
 		UserID:               userID,
 	}
-	return tx.Create(flowState)
+	return flowState
 }
 
 func FindFlowStateByAuthCode(tx *storage.Connection, authCode string) (*FlowState, error) {


Test patch:
diff --git a/internal/api/token_test.go b/internal/api/token_test.go
index 0c8cc237..b12a79a8 100644
--- a/internal/api/token_test.go
+++ b/internal/api/token_test.go
@@ -306,8 +306,7 @@ func (ts *TokenTestSuite) TestTokenPKCEGrantFailure() {
 	invalidVerifier := codeVerifier + "123"
 	codeChallenge := sha256.Sum256([]byte(codeVerifier))
 	challenge := base64.RawURLEncoding.EncodeToString(codeChallenge[:])
-	flowState, err := models.NewFlowState("github", challenge, models.SHA256, models.OAuth)
-	require.NoError(ts.T(), err)
+	flowState := models.NewFlowState("github", challenge, models.SHA256, models.OAuth, nil)
 	flowState.AuthCode = authCode
 	require.NoError(ts.T(), ts.API.db.Create(flowState))
 	cases := []struct {
diff --git a/internal/api/verify_test.go b/internal/api/verify_test.go
index 51118f35..c782446f 100644
--- a/internal/api/verify_test.go
+++ b/internal/api/verify_test.go
@@ -661,8 +661,8 @@ func (ts *VerifyTestSuite) TestVerifyPKCEOTP() {
 			var buffer bytes.Buffer
 			require.NoError(ts.T(), json.NewEncoder(&buffer).Encode(c.payload))
 			codeChallenge := "codechallengecodechallengcodechallengcodechallengcodechallenge" + c.payload.Type
-			err := models.NewFlowStateWithUserID(ts.API.db, c.authenticationMethod.String(), codeChallenge, models.SHA256, c.authenticationMethod, &u.ID)
-			require.NoError(ts.T(), err)
+			flowState := models.NewFlowState(c.authenticationMethod.String(), codeChallenge, models.SHA256, c.authenticationMethod, &u.ID)
+			require.NoError(ts.T(), ts.API.db.Create(flowState))
 
 			requestUrl := fmt.Sprintf("http://localhost/verify?type=%v&token=%v", c.payload.Type, c.payload.Token)
 			req := httptest.NewRequest(http.MethodGet, requestUrl, &buffer)