From 254ea4b0c85209c1f1bf12667b10015979566d22 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Mon, 1 Nov 2021 17:03:55 +0100
Subject: [PATCH] social_auth: Save authentication method information in the
 session.

The immediate application of this will be for SAML SP-initiated logout,
where information about which IdP was used for authenticating the
session needs to be accessed. Aside of that, this seems like generally
valuable session information to keep that other features may benefit
from in the future.
---
 zerver/views/auth.py |  4 ++++
 zproject/backends.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/zerver/views/auth.py b/zerver/views/auth.py
index 3c3f29cfb1..23e3606ad9 100644
--- a/zerver/views/auth.py
+++ b/zerver/views/auth.py
@@ -339,6 +339,10 @@ def login_or_register_remote_user(request: HttpRequest, result: ExternalAuthResu
     * A zulip:// URL to send control back to the mobile or desktop apps if they
       are doing authentication using the mobile_flow_otp or desktop_flow_otp flow.
     """
+
+    for key, value in result.data_dict.get("params_to_store_in_authenticated_session", {}).items():
+        request.session[key] = value
+
     user_profile = result.user_profile
     if user_profile is None or user_profile.is_mirror_dummy:
         return register_remote_user(request, result)
diff --git a/zproject/backends.py b/zproject/backends.py
index 4b2c4939eb..2556a53877 100644
--- a/zproject/backends.py
+++ b/zproject/backends.py
@@ -1295,6 +1295,7 @@ class ExternalAuthDataDict(TypedDict, total=False):
     desktop_flow_otp: Optional[str]
     multiuse_object_key: str
     full_name_validated: bool
+    params_to_store_in_authenticated_session: Dict[str, str]
 
 
 class ExternalAuthResult:
@@ -1797,6 +1798,7 @@ def social_auth_finish(
         full_name_validated=full_name_validated,
         mobile_flow_otp=mobile_flow_otp,
         desktop_flow_otp=desktop_flow_otp,
+        params_to_store_in_authenticated_session=backend.get_params_to_store_in_authenticated_session(),
     )
     if user_profile is None:
         data_dict.update(dict(full_name=full_name, email=email_address))
@@ -1875,6 +1877,13 @@ class SocialAuthMixin(ZulipAuthMixin, ExternalAuthMethod, BaseAuth):
     def should_auto_signup(self) -> bool:
         return False
 
+    def get_params_to_store_in_authenticated_session(self) -> Dict[str, str]:
+        """
+        Specifies a dict of keys:values to be saved in the user's session
+        after successfully authenticating.
+        """
+        return {"authentication_method": self.name}
+
     @classmethod
     def dict_representation(cls, realm: Optional[Realm] = None) -> List[ExternalAuthMethodDictT]:
         return [
@@ -2701,6 +2710,11 @@ class SAMLAuthBackend(SocialAuthMixin, SAMLAuth):
         assert isinstance(auto_signup, bool)
         return auto_signup
 
+    def get_params_to_store_in_authenticated_session(self) -> Dict[str, str]:
+        idp_name = self.strategy.session_get("saml_idp_name")
+
+        return {"authentication_method": f"saml:{idp_name}"}
+
 
 def patch_saml_auth_require_messages_signed(auth: OneLogin_Saml2_Auth) -> None:
     """