From b2966b158a5c27b3c8904ba87c955bbe89730e68 Mon Sep 17 00:00:00 2001
From: Bernd Schoolmann <mail@quexten.com>
Date: Fri, 6 Sep 2024 00:15:09 +0200
Subject: [PATCH] [PM-11461] Fix jit proivisioning for snap desktop (#10794)

* Fix snap jit privisioning

* Fix server keeping connection open on favicon request
---
 .../sso-localhost-callback.service.ts         | 31 +++++++++++--------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/apps/desktop/src/platform/services/sso-localhost-callback.service.ts b/apps/desktop/src/platform/services/sso-localhost-callback.service.ts
index 5efe73e2ad..cd4c7df66e 100644
--- a/apps/desktop/src/platform/services/sso-localhost-callback.service.ts
+++ b/apps/desktop/src/platform/services/sso-localhost-callback.service.ts
@@ -18,10 +18,10 @@ export class SSOLocalhostCallbackService {
     private messagingService: MessageSender,
   ) {
     ipcMain.handle("openSsoPrompt", async (event, { codeChallenge, state }) => {
-      const { ssoCode } = await this.openSsoPrompt(codeChallenge, state);
+      const { ssoCode, recvState } = await this.openSsoPrompt(codeChallenge, state);
       this.messagingService.send("ssoCallback", {
         code: ssoCode,
-        state: state,
+        state: recvState,
         redirectUri: this.ssoRedirectUri,
       });
     });
@@ -30,24 +30,20 @@ export class SSOLocalhostCallbackService {
   private async openSsoPrompt(
     codeChallenge: string,
     state: string,
-  ): Promise<{ ssoCode: string; orgIdentifier: string }> {
+  ): Promise<{ ssoCode: string; recvState: string }> {
     const env = await firstValueFrom(this.environmentService.environment$);
 
     return new Promise((resolve, reject) => {
       const callbackServer = http.createServer((req, res) => {
-        // after 5 minutes, close the server
-        setTimeout(
-          () => {
-            callbackServer.close(() => reject());
-          },
-          5 * 60 * 1000,
-        );
-
         const urlString = "http://localhost" + req.url;
         const url = new URL(urlString);
         const code = url.searchParams.get("code");
+        if (code == null) {
+          res.writeHead(404);
+          res.end("not found");
+          return;
+        }
         const receivedState = url.searchParams.get("state");
-        const orgIdentifier = this.getOrgIdentifierFromState(receivedState);
         res.setHeader("Content-Type", "text/html");
         if (code != null && receivedState != null && this.checkState(receivedState, state)) {
           res.writeHead(200);
@@ -60,7 +56,7 @@ export class SSOLocalhostCallbackService {
           callbackServer.close(() =>
             resolve({
               ssoCode: code,
-              orgIdentifier: orgIdentifier,
+              recvState: receivedState,
             }),
           );
         } else {
@@ -74,6 +70,7 @@ export class SSOLocalhostCallbackService {
           callbackServer.close(() => reject());
         }
       });
+
       let foundPort = false;
       const webUrl = env.getWebVaultUrl();
       for (let port = 8065; port <= 8070; port++) {
@@ -102,6 +99,14 @@ export class SSOLocalhostCallbackService {
       if (!foundPort) {
         reject();
       }
+
+      // after 5 minutes, close the server
+      setTimeout(
+        () => {
+          callbackServer.close(() => reject());
+        },
+        5 * 60 * 1000,
+      );
     });
   }