more type fixes

supertokens · Nov 23, 2023 · 43e8f7c · 43e8f7c
1 parent 1a672b9
commit 43e8f7c
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 37 deletions.
diff --git a/recipe/session/cookieAndHeaders.go b/recipe/session/cookieAndHeaders.go
@@ -55,15 +55,15 @@ type TokenInfo struct {
 	Up  interface{} `json:"up"`
 }
 
-func ClearSessionFromAllTokenTransferMethods(config sessmodels.TypeNormalisedInput, req *http.Request, res http.ResponseWriter, request *http.Request, userContext supertokens.UserContext) error {
+func ClearSessionFromAllTokenTransferMethods(config sessmodels.TypeNormalisedInput, req *http.Request, res http.ResponseWriter, userContext supertokens.UserContext) error {
 	// We are clearing the session in all transfermethods to be sure to override cookies in case they have been already added to the response.
 	// This is done to handle the following use-case:
 	// If the app overrides signInPOST to check the ban status of the user after the original implementation and throwing an UNAUTHORISED error
 	// In this case: the SDK has attached cookies to the response, but none was sent with the request
 	// We can't know which to clear since we can't reliably query or remove the set-cookie header added to the response (causes issues in some frameworks, i.e.: hapi)
 	// The safe solution in this case is to overwrite all the response cookies/headers with an empty value, which is what we are doing here
 	for _, transferMethod := range AvailableTokenTransferMethods {
-		err := ClearSession(config, res, transferMethod, request, userContext)
+		err := ClearSession(config, res, transferMethod, req, userContext)
 		if err != nil {
 			return err
 		}

diff --git a/recipe/session/recipe.go b/recipe/session/recipe.go
@@ -183,29 +183,29 @@ func (r *Recipe) getAllCORSHeaders() []string {
 	return resp
 }
 
-func (r *Recipe) handleError(err error, req *http.Request, res http.ResponseWriter) (bool, error) {
+func (r *Recipe) handleError(err error, req *http.Request, res http.ResponseWriter, userContext supertokens.UserContext) (bool, error) {
 	if defaultErrors.As(err, &errors.UnauthorizedError{}) {
 		supertokens.LogDebugMessage("errorHandler: returning UNAUTHORISED")
 		unauthErr := err.(errors.UnauthorizedError)
 		if unauthErr.ClearTokens == nil || *unauthErr.ClearTokens {
 			supertokens.LogDebugMessage("errorHandler: Clearing tokens because of UNAUTHORISED response")
-			ClearSessionFromAllTokenTransferMethods(r.Config, req, res)
+			ClearSessionFromAllTokenTransferMethods(r.Config, req, res, userContext)
 		}
 		return true, r.Config.ErrorHandlers.OnUnauthorised(err.Error(), req, res)
 	} else if defaultErrors.As(err, &errors.TryRefreshTokenError{}) {
 		supertokens.LogDebugMessage("errorHandler: returning TRY_REFRESH_TOKEN")
 		return true, r.Config.ErrorHandlers.OnTryRefreshToken(err.Error(), req, res)
 	} else if defaultErrors.As(err, &errors.TokenTheftDetectedError{}) {
 		supertokens.LogDebugMessage("errorHandler: clearing tokens because of TOKEN_THEFT_DETECTED response")
-		ClearSessionFromAllTokenTransferMethods(r.Config, req, res)
+		ClearSessionFromAllTokenTransferMethods(r.Config, req, res, userContext)
 		errs := err.(errors.TokenTheftDetectedError)
 		return true, r.Config.ErrorHandlers.OnTokenTheftDetected(errs.Payload.SessionHandle, errs.Payload.UserID, req, res)
 	} else if defaultErrors.As(err, &errors.InvalidClaimError{}) {
 		supertokens.LogDebugMessage("errorHandler: returning INVALID_CLAIMS")
 		errs := err.(errors.InvalidClaimError)
 		return true, r.Config.ErrorHandlers.OnInvalidClaim(errs.InvalidClaims, req, res)
 	} else {
-		return r.OpenIdRecipe.RecipeModule.HandleError(err, req, res)
+		return r.OpenIdRecipe.RecipeModule.HandleError(err, req, res, userContext)
 	}
 }
 

diff --git a/recipe/session/session.go b/recipe/session/session.go
@@ -83,7 +83,7 @@ func newSessionContainer(config sessmodels.TypeNormalisedInput, session *Session
 			// If we instead clear the cookies only when revokeSession
 			// returns true, it can cause this kind of a bug:
 			// https://github.com/supertokens/supertokens-node/issues/343
-			ClearSession(config, session.requestResponseInfo.Res, session.requestResponseInfo.TokenTransferMethod)
+			ClearSession(config, session.requestResponseInfo.Res, session.requestResponseInfo.TokenTransferMethod, session.requestResponseInfo.Req, supertokens.SetRequestInUserContextIfNotDefined(userContext, session.requestResponseInfo.Req))
 		}
 		return nil
 	}
@@ -195,7 +195,7 @@ func newSessionContainer(config sessmodels.TypeNormalisedInput, session *Session
 			session.accessTokenUpdated = true
 
 			if session.requestResponseInfo != nil {
-				setTokenErr := SetAccessTokenInResponse(config, session.requestResponseInfo.Res, session.accessToken, session.frontToken, session.requestResponseInfo.TokenTransferMethod)
+				setTokenErr := SetAccessTokenInResponse(config, session.requestResponseInfo.Res, session.accessToken, session.frontToken, session.requestResponseInfo.TokenTransferMethod, session.requestResponseInfo.Req, supertokens.SetRequestInUserContextIfNotDefined(userContext, session.requestResponseInfo.Req))
 				if setTokenErr != nil {
 					return setTokenErr
 				}
@@ -279,6 +279,32 @@ func newSessionContainer(config sessmodels.TypeNormalisedInput, session *Session
 		return sessionContainer.MergeIntoAccessTokenPayloadWithContext(update, userContext)
 	}
 
+	sessionContainer.AttachToRequestResponseWithContext = func(info sessmodels.RequestResponseInfo, userContext supertokens.UserContext) error {
+		session.requestResponseInfo = &info
+
+		if session.accessTokenUpdated {
+			err := SetAccessTokenInResponse(config, info.Res, session.accessToken, session.frontToken, info.TokenTransferMethod, session.requestResponseInfo.Req, supertokens.SetRequestInUserContextIfNotDefined(userContext, session.requestResponseInfo.Req))
+
+			if err != nil {
+				return err
+			}
+
+			if session.refreshToken != nil {
+				err = setToken(config, info.Res, sessmodels.RefreshToken, session.refreshToken.Token, session.refreshToken.Expiry, info.TokenTransferMethod, session.requestResponseInfo.Req, supertokens.SetRequestInUserContextIfNotDefined(userContext, session.requestResponseInfo.Req))
+
+				if err != nil {
+					return err
+				}
+			}
+
+			if session.antiCSRFToken != nil {
+				setAntiCsrfTokenInHeaders(info.Res, *session.antiCSRFToken)
+			}
+		}
+
+		return nil
+	}
+
 	sessionContainer.RevokeSession = func() error {
 		return sessionContainer.RevokeSessionWithContext(&map[string]interface{}{})
 	}
@@ -346,30 +372,8 @@ func newSessionContainer(config sessmodels.TypeNormalisedInput, session *Session
 		}
 	}
 
-	sessionContainer.AttachToRequestResponse = func(info sessmodels.RequestResponseInfo) error {
-		session.requestResponseInfo = &info
-
-		if session.accessTokenUpdated {
-			err := SetAccessTokenInResponse(config, info.Res, session.accessToken, session.frontToken, info.TokenTransferMethod)
-
-			if err != nil {
-				return err
-			}
-
-			if session.refreshToken != nil {
-				err = setToken(config, info.Res, sessmodels.RefreshToken, session.refreshToken.Token, session.refreshToken.Expiry, info.TokenTransferMethod)
-
-				if err != nil {
-					return err
-				}
-			}
-
-			if session.antiCSRFToken != nil {
-				setAntiCsrfTokenInHeaders(info.Res, *session.antiCSRFToken)
-			}
-		}
-
-		return nil
+	sessionContainer.AttachToRequestResponseWithContext = func(info sessmodels.RequestResponseInfo, userContext supertokens.UserContext) error {
+		return sessionContainer.AttachToRequestResponseWithContext(info, &map[string]interface{}{})
 	}
 
 	return sessionContainer

diff --git a/recipe/session/utils.go b/recipe/session/utils.go
@@ -287,20 +287,20 @@ func GetCurrTimeInMS() uint64 {
 	return uint64(time.Now().UnixNano() / 1000000)
 }
 
-func SetAccessTokenInResponse(config sessmodels.TypeNormalisedInput, res http.ResponseWriter, accessToken string, frontToken string, tokenTransferMethod sessmodels.TokenTransferMethod) error {
+func SetAccessTokenInResponse(config sessmodels.TypeNormalisedInput, res http.ResponseWriter, accessToken string, frontToken string, tokenTransferMethod sessmodels.TokenTransferMethod, request *http.Request, userContext supertokens.UserContext) error {
 	setFrontTokenInHeaders(res, frontToken)
 	// We set the expiration to 100 years, because we can't really access the expiration of the refresh token everywhere we are setting it.
 	// This should be safe to do, since this is only the validity of the cookie (set here or on the frontend) but we check the expiration of the JWT anyway.
 	// Even if the token is expired the presence of the token indicates that the user could have a valid refresh
 	// Setting them to infinity would require special case handling on the frontend and just adding 100 years seems enough.
-	setToken(config, res, sessmodels.AccessToken, accessToken, GetCurrTimeInMS()+uint64(accessTokenCookiesExpiryDurationMillis), tokenTransferMethod)
+	setToken(config, res, sessmodels.AccessToken, accessToken, GetCurrTimeInMS()+uint64(accessTokenCookiesExpiryDurationMillis), tokenTransferMethod, request, userContext)
 
 	if config.ExposeAccessTokenToFrontendInCookieBasedAuth && tokenTransferMethod == sessmodels.CookieTransferMethod {
 		// We set the expiration to 100 years, because we can't really access the expiration of the refresh token everywhere we are setting it.
 		// This should be safe to do, since this is only the validity of the cookie (set here or on the frontend) but we check the expiration of the JWT anyway.
 		// Even if the token is expired the presence of the token indicates that the user could have a valid refresh
 		// Setting them to infinity would require special case handling on the frontend and just adding 100 years seems enough.
-		setToken(config, res, sessmodels.AccessToken, accessToken, GetCurrTimeInMS()+uint64(accessTokenCookiesExpiryDurationMillis), sessmodels.HeaderTransferMethod)
+		setToken(config, res, sessmodels.AccessToken, accessToken, GetCurrTimeInMS()+uint64(accessTokenCookiesExpiryDurationMillis), sessmodels.HeaderTransferMethod, request, userContext)
 	}
 	return nil
 }

diff --git a/supertokens/recipeModule.go b/supertokens/recipeModule.go
@@ -28,7 +28,7 @@ type RecipeModule struct {
 	GetAllCORSHeaders             func() []string
 	GetAPIsHandled                func() ([]APIHandled, error)
 	ReturnAPIIdIfCanHandleRequest func(path NormalisedURLPath, method string, userContext UserContext) (*string, string, error)
-	HandleError                   func(err error, req *http.Request, res http.ResponseWriter) (bool, error)
+	HandleError                   func(err error, req *http.Request, res http.ResponseWriter, userContext UserContext) (bool, error)
 	OnSuperTokensAPIError         func(err error, req *http.Request, res http.ResponseWriter)
 }
 
@@ -39,7 +39,7 @@ func MakeRecipeModule(
 	getAllCORSHeaders func() []string,
 	getAPIsHandled func() ([]APIHandled, error),
 	returnAPIIdIfCanHandleRequest func(path NormalisedURLPath, method string, userContext UserContext) (*string, string, error),
-	handleError func(err error, req *http.Request, res http.ResponseWriter) (bool, error),
+	handleError func(err error, req *http.Request, res http.ResponseWriter, userContext UserContext) (bool, error),
 	onSuperTokensAPIError func(err error, req *http.Request, res http.ResponseWriter)) RecipeModule {
 	if handleError == nil {
 		// Execution will come here only if there is a bug in the code