feat: implement session restore when cookies are not reliable

When cookies are created with SameSite policy, they won't be available
during the authentication flow which uses POST such as OpenID or SAML.
This adds support in Strategy to get session ID and restore it later in
the login flow.

See python-social-auth/social-app-django#481

Author: nijel

social_core/backends/base.py

@@ -224,7 +224,7 @@ def continue_pipeline(self, partial):
             self, pipeline_index=partial.next_step, *partial.args, **partial.kwargs
         )
 
-    def auth_extra_arguments(self):
+    def auth_extra_arguments(self) -> dict[str, str]:
         """Return extra arguments needed on auth process. The defaults can be
         overridden by GET parameters."""
         extra_arguments = self.setting("AUTH_EXTRA_ARGUMENTS", {}).copy()

social_core/backends/oauth.py

@@ -280,8 +280,8 @@ def oauth_authorization_request(self, token):
             token = parse_qs(token)
         params = self.auth_extra_arguments() or {}
         params.update(self.get_scope_argument())
-        params[self.OAUTH_TOKEN_PARAMETER_NAME] = token.get(
-            self.OAUTH_TOKEN_PARAMETER_NAME
+        params[self.OAUTH_TOKEN_PARAMETER_NAME] = cast(
+            "str", token.get(self.OAUTH_TOKEN_PARAMETER_NAME)
         )
         state = self.get_or_create_state()
         params[self.REDIRECT_URI_PARAMETER_NAME] = self.get_redirect_uri(state)

social_core/backends/open_id.py

@@ -144,20 +144,25 @@ def extra_data(self, user, uid, response, details=None, *args, **kwargs):
         values.update(from_details)
         return values
 
+    def get_return_to(self) -> str:
+        params: dict[str, str] = {}
+        if session_id := self.strategy.get_session_id():
+            params[self.strategy.SESSION_SAVE_KEY] = session_id
+
+        return url_add_parameters(self.strategy.absolute_uri(self.redirect_uri), params)
+
     def auth_url(self):
         """Return auth URL returned by service"""
         openid_request = self.setup_request(self.auth_extra_arguments())
         # Construct completion URL, including page we should redirect to
-        return_to = self.strategy.absolute_uri(self.redirect_uri)
-        return openid_request.redirectURL(self.trust_root(), return_to)
+        return openid_request.redirectURL(self.trust_root(), self.get_return_to())
 
     def auth_html(self):
         """Return auth HTML returned by service"""
         openid_request = self.setup_request(self.auth_extra_arguments())
-        return_to = self.strategy.absolute_uri(self.redirect_uri)
         form_tag = {"id": "openid_message"}
         return openid_request.htmlMarkup(
-            self.trust_root(), return_to, form_tag_attrs=form_tag
+            self.trust_root(), self.get_return_to(), form_tag_attrs=form_tag
         )
 
     def trust_root(self):
@@ -167,7 +172,7 @@ def trust_root(self):
     def continue_pipeline(self, partial):
         """Continue previous halted pipeline"""
         response = self.consumer().complete(
-            dict(self.data.items()), self.strategy.absolute_uri(self.redirect_uri)
+            dict(self.data.items()), self.get_return_to()
         )
         return self.strategy.authenticate(
             self,
@@ -180,9 +185,11 @@ def continue_pipeline(self, partial):
     def auth_complete(self, *args, **kwargs):
         """Complete auth process"""
         response = self.consumer().complete(
-            dict(self.data.items()), self.strategy.absolute_uri(self.redirect_uri)
+            dict(self.data.items()), self.get_return_to()
         )
         self.process_error(response)
+        if session_id := self.data.get(self.strategy.SESSION_SAVE_KEY):
+            self.strategy.restore_session(session_id, kwargs)
         return self.strategy.authenticate(self, response=response, *args, **kwargs)
 
     def process_error(self, data):

social_core/backends/saml.py

@@ -286,6 +286,8 @@ def auth_url(self):
             "idp": idp_name,
             "next": self.data.get("next"),
         }
+        if session_id := self.strategy.get_session_id():
+            relay_state[self.strategy.SESSION_SAVE_KEY] = session_id
         return auth.login(return_to=json.dumps(relay_state))
 
     def get_user_details(self, response):
@@ -326,7 +328,9 @@ def auth_complete(self, *args, **kwargs):
             idp_name = relay_state_str
         else:
             idp_name = relay_state["idp"]
-            if next_url := relay_state.get("next"):
+            if session_id := relay_state.get(self.strategy.SESSION_SAVE_KEY):
+                self.strategy.restore_session(session_id, kwargs)
+            elif next_url := relay_state.get("next"):
                 # The do_complete action expects the "next" URL to be in session state or the request params.
                 self.strategy.session_set(kwargs.get("redirect_name", "next"), next_url)
 

social_core/exceptions.py

@@ -10,9 +10,22 @@ class SocialAuthBaseException(ValueError):
     """Base class for pipeline exceptions."""
 
 
+class StrategyMissingFeatureError(SocialAuthBaseException):
+    """Strategy does not support this."""
+
+    def __init__(self, strategy_name: str, feature_name: str):
+        self.strategy_name = strategy_name
+        self.feature_name = feature_name
+        super().__init__()
+
+    def __str__(self):
+        return f"Strategy {self.strategy_name} does not support {self.feature_name}"
+
+
 class WrongBackend(SocialAuthBaseException):
     def __init__(self, backend_name: str):
         self.backend_name = backend_name
+        super().__init__()
 
     def __str__(self):
         return f'Incorrect authentication service "{self.backend_name}"'

social_core/strategy.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import secrets
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from .backends.utils import get_backend
+from .exceptions import StrategyMissingFeatureError
 from .pipeline import DEFAULT_AUTH_PIPELINE, DEFAULT_DISCONNECT_PIPELINE
 from .pipeline.utils import partial_load, partial_prepare, partial_store
 from .store import OpenIdSessionWrapper, OpenIdStore
@@ -35,6 +36,7 @@ def render_string(self, html, context):
 class BaseStrategy:
     ALLOWED_CHARS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
     DEFAULT_TEMPLATE_STRATEGY = BaseTemplateStrategy
+    SESSION_SAVE_KEY = "psa_session_id"
 
     def __init__(self, storage=None, tpl=None):
         self.storage = storage
@@ -61,6 +63,20 @@ def session_setdefault(self, name, value):
         self.session_set(name, value)
         return self.session_get(name)
 
+    def get_session_id(self) -> str | None:
+        """
+        Return session ID to be used by restore_session.
+        """
+        return None
+
+    def restore_session(self, session_id: str, kwargs: dict[str, Any]) -> None:
+        """
+        Restores session and updates kwargs to match it.
+
+        This is only called if get_session_id returns a value.
+        """
+        raise StrategyMissingFeatureError(self.__class__.__name__, "session restore")
+
     def openid_session_dict(self, name):
         # Many frameworks are switching the session serialization from Pickle
         # to JSON to avoid code execution risks. Flask did this from Flask