From fba59381a0a6bb39290a4720ba582d9a3a28caec Mon Sep 17 00:00:00 2001
From: Julien CLEMENT <33104892+Jujure@users.noreply.github.com>
Date: Fri, 21 Mar 2025 20:10:34 +0100
Subject: [PATCH] Generate/load pickle key on SSO (#29568)

* Generate/load pickle key when logged in with SSO

* add comments

* Refactor pickle key loading/creation

* Coding style fixes and fix racy loadOrCreatePickleKey

* fix outdated documentation comment

* fix prettier

Signed-off-by: Julien CLEMENT <julien.clement@epita.fr>

---------

Signed-off-by: Julien CLEMENT <julien.clement@epita.fr>
---
 src/Lifecycle.ts | 49 ++++++++++++++++++++++++++++++++++++------------
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/Lifecycle.ts b/src/Lifecycle.ts
index 3ef7c56e82..a7f4d6d31b 100644
--- a/src/Lifecycle.ts
+++ b/src/Lifecycle.ts
@@ -406,6 +406,39 @@ export function attemptTokenLogin(
         });
 }
 
+/**
+ * Load the pickle key inside the credentials or create it if it does not exist for this device.
+ *
+ * @param credentials Holds the device to load/store the pickle key
+ *
+ * @returns {Promise} promise which resolves to the loaded or generated pickle key or undefined if
+ *    none was loaded nor generated
+ */
+async function loadOrCreatePickleKey(credentials: IMatrixClientCreds): Promise<string | undefined> {
+    // Try to load the pickle key
+    const userId = credentials.userId;
+    const deviceId = credentials.deviceId;
+    let pickleKey = (await PlatformPeg.get()?.getPickleKey(userId, deviceId ?? "")) ?? undefined;
+    if (!pickleKey) {
+        // Create it if it did not exist
+        pickleKey =
+            userId && deviceId
+                ? ((await PlatformPeg.get()?.createPickleKey(userId, deviceId)) ?? undefined)
+                : undefined;
+        if (pickleKey) {
+            logger.log(`Created pickle key for ${credentials.userId}|${credentials.deviceId}`);
+        } else {
+            logger.log("Pickle key not created");
+        }
+    } else {
+        logger.log(
+            `Pickle key already exists for ${credentials.userId}|${credentials.deviceId} do not create a new one`,
+        );
+    }
+
+    return pickleKey;
+}
+
 /**
  * Called after a successful token login or OIDC authorization.
  * Clear storage then save new credentials in storage
@@ -413,6 +446,8 @@ export function attemptTokenLogin(
  */
 async function onSuccessfulDelegatedAuthLogin(credentials: IMatrixClientCreds): Promise<void> {
     await clearStorage();
+    // SSO does not go through setLoggedIn so we need to load/create the pickle key here too
+    credentials.pickleKey = await loadOrCreatePickleKey(credentials);
     await persistCredentials(credentials);
 
     // remember that we just logged in
@@ -655,18 +690,8 @@ async function handleLoadSessionFailure(e: unknown): Promise<boolean> {
 export async function setLoggedIn(credentials: IMatrixClientCreds): Promise<MatrixClient> {
     credentials.freshLogin = true;
     stopMatrixClient();
-    const pickleKey =
-        credentials.userId && credentials.deviceId
-            ? await PlatformPeg.get()?.createPickleKey(credentials.userId, credentials.deviceId)
-            : null;
-
-    if (pickleKey) {
-        logger.log(`Created pickle key for ${credentials.userId}|${credentials.deviceId}`);
-    } else {
-        logger.log("Pickle key not created");
-    }
-
-    return doSetLoggedIn({ ...credentials, pickleKey: pickleKey ?? undefined }, true, true);
+    credentials.pickleKey = await loadOrCreatePickleKey(credentials);
+    return doSetLoggedIn(credentials, true, true);
 }
 
 /**