allow extra state to be persisted between login and callback - for example to store a 'return to' url to redirect the user to after they login

Author: jessemcl-flwls

authlib/integrations/starlette_client/apps.py

@@ -19,11 +19,12 @@ async def save_authorize_data(self, request, **kwargs):
         else:
             raise RuntimeError('Missing state value')
 
-    async def authorize_redirect(self, request, redirect_uri=None, **kwargs):
+    async def authorize_redirect(self, request, redirect_uri=None, extra_state=None, **kwargs):
         """Create a HTTP Redirect for Authorization Endpoint.
 
         :param request: HTTP request instance from Starlette view.
         :param redirect_uri: Callback or redirect URI for authorization.
+        :param extra_state: Extra state data to be stored in session.
         :param kwargs: Extra parameters to include.
         :return: A HTTP redirect response.
         """
@@ -32,6 +33,8 @@ async def authorize_redirect(self, request, redirect_uri=None, **kwargs):
         if redirect_uri and isinstance(redirect_uri, URL):
             redirect_uri = str(redirect_uri)
         rv = await self.create_authorization_url(redirect_uri, **kwargs)
+        if extra_state is not None:
+            rv['extra_state'] = extra_state
         await self.save_authorize_data(request, redirect_uri=redirect_uri, **rv)
         return RedirectResponse(rv['url'], status_code=302)
 
@@ -83,4 +86,6 @@ async def authorize_access_token(self, request, **kwargs):
         if 'id_token' in token and 'nonce' in state_data:
             userinfo = await self.parse_id_token(token, nonce=state_data['nonce'], claims_options=claims_options)
             token['userinfo'] = userinfo
+            if 'extra_state' in state_data:
+                token['extra_state'] = state_data['extra_state']
         return token