From 046e850a657e0790e9c8d3045d1400de84f62dc6 Mon Sep 17 00:00:00 2001
From: Kraiem Taha Yassine <tahayk2@gmail.com>
Date: Thu, 2 May 2024 14:02:04 +0200
Subject: [PATCH] fix(chalice): changed SSO and added logs (#2155)

---
 ee/api/chalicelib/core/users.py |  7 +++++--
 ee/api/routers/saml.py          | 25 +++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/ee/api/chalicelib/core/users.py b/ee/api/chalicelib/core/users.py
index af6f51bda..65239537a 100644
--- a/ee/api/chalicelib/core/users.py
+++ b/ee/api/chalicelib/core/users.py
@@ -1,4 +1,5 @@
 import json
+import logging
 import secrets
 
 from decouple import config
@@ -7,12 +8,14 @@ from starlette import status
 
 import schemas
 from chalicelib.core import authorizers, metadata, projects
+from chalicelib.core import roles
 from chalicelib.core import tenants, assist
 from chalicelib.utils import email_helper, smtp
 from chalicelib.utils import helper
 from chalicelib.utils import pg_client
 from chalicelib.utils.TimeUTC import TimeUTC
-from chalicelib.core import roles
+
+logger = logging.getLogger(__name__)
 
 
 def __generate_invitation_token():
@@ -861,7 +864,7 @@ def authenticate_sso(email, internal_id, exp=None):
                                                              jwt_jti=jwt_r_jti),
             "refreshTokenMaxAge": config("JWT_REFRESH_EXPIRATION", cast=int),
         }
-
+    logger.warning(f"SSO user not found with email: {email} and internal_id: {internal_id}")
     return None
 
 
diff --git a/ee/api/routers/saml.py b/ee/api/routers/saml.py
index aeaa107f2..7a0f9caf4 100644
--- a/ee/api/routers/saml.py
+++ b/ee/api/routers/saml.py
@@ -149,8 +149,29 @@ async def process_sso_assertion_tk(tenantKey: str, request: Request):
     session = req["cookie"]["session"]
     auth = init_saml_auth(req)
 
-    redirect_to_link2 = json.loads(req.get("post_data", {}) \
-                                   .get('RelayState', '{}')).get("iFrame")
+    post_data = req.get("post_data")
+    if post_data is None:
+        post_data = {}
+    elif isinstance(post_data, str):
+        post_data = json.loads(post_data)
+    elif not isinstance(post_data, dict):
+        logger.error("Received invalid post_data")
+        logger.error("type: {}".format(type(post_data)))
+        logger.error(post_data)
+        post_data = {}
+
+    redirect_to_link2 = None
+    relay_state = post_data.get('RelayState')
+    if relay_state:
+        if isinstance(relay_state, str):
+            relay_state = json.loads(relay_state)
+        elif not isinstance(relay_state, dict):
+            logger.error("Received invalid relay_state")
+            logger.error("type: {}".format(type(relay_state)))
+            logger.error(relay_state)
+            relay_state = {}
+        redirect_to_link2 = relay_state.get("iFrame")
+
     request_id = None
     if 'AuthNRequestID' in session:
         request_id = session['AuthNRequestID']