From 7e5364d40001c562da96beac055832aa9247173d Mon Sep 17 00:00:00 2001
From: Jan Prochazka <jan.prochazka@sprinx.cz>
Date: Thu, 8 Aug 2024 11:45:21 +0200
Subject: [PATCH] msentra auth

---
 packages/api/src/controllers/connections.js       | 12 ++++++++++++
 packages/tools/src/driverBase.ts                  |  1 +
 packages/types/engines.d.ts                       |  1 +
 packages/web/src/utility/api.ts                   | 15 ++++++++++++---
 plugins/dbgate-plugin-mssql/src/backend/driver.js |  3 +++
 5 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/packages/api/src/controllers/connections.js b/packages/api/src/controllers/connections.js
index f8b3461e..7476f380 100644
--- a/packages/api/src/controllers/connections.js
+++ b/packages/api/src/controllers/connections.js
@@ -443,4 +443,16 @@ module.exports = {
     const loginResp = await getAuthProviderById(amoid).login(null, null, { conid });
     return loginResp;
   },
+
+  volatileDbloginFromAuth_meta: true,
+  async volatileDbloginFromAuth({ conid }, req) {
+    const connection = await this.getCore({ conid });
+    const driver = requireEngineDriver(connection);
+    const accessToken = await driver.getAccessTokenFromAuth(connection, req);
+    if (accessToken) {
+      const volatile = await this.saveVolatile({ conid, accessToken });
+      return volatile;
+    }
+    return null;
+  },
 };
diff --git a/packages/tools/src/driverBase.ts b/packages/tools/src/driverBase.ts
index 626ed0b2..20044d7a 100644
--- a/packages/tools/src/driverBase.ts
+++ b/packages/tools/src/driverBase.ts
@@ -146,4 +146,5 @@ export const driverBase = {
   },
   showConnectionField: (field, values) => false,
   showConnectionTab: field => true,
+  getAccessTokenFromAuth: async (connection, req) => null,
 };
diff --git a/packages/types/engines.d.ts b/packages/types/engines.d.ts
index e8bfc815..58561864 100644
--- a/packages/types/engines.d.ts
+++ b/packages/types/engines.d.ts
@@ -151,6 +151,7 @@ export interface EngineDriver {
   stopProfiler(pool, profiler): Promise<void>;
   getRedirectAuthUrl(connection, options): Promise<string>;
   getAuthTokenFromCode(connection, options): Promise<string>;
+  getAccessTokenFromAuth(connection, req): Promise<string | null>;
 
   analyserClass?: any;
   dumperClass?: any;
diff --git a/packages/web/src/utility/api.ts b/packages/web/src/utility/api.ts
index b8c4eea4..c2fea27c 100644
--- a/packages/web/src/utility/api.ts
+++ b/packages/web/src/utility/api.ts
@@ -79,13 +79,22 @@ function wantEventSource() {
   }
 }
 
-function processApiResponse(route, args, resp) {
+async function processApiResponse(route, args, resp) {
   // if (apiLogging) {
   //   console.log('<<< API RESPONSE', route, args, resp);
   // }
 
   if (resp?.missingCredentials) {
     if (resp.detail.redirectToDbLogin) {
+      const volatile = await apiCall('connections/volatile-dblogin-from-auth', { conid: resp.detail.conid });
+      if (volatile) {
+        setVolatileConnectionRemapping(resp.detail.conid, volatile._id);
+        await callServerPing();
+        dispatchCacheChange({ key: `server-status-changed` });
+        batchDispatchCacheTriggers(x => x.conid == resp.detail.conid);
+        return null;
+      }
+
       const state = `dbg-dblogin:${strmid}:${resp.detail.conid}`;
       localStorage.setItem('dbloginState', state);
       openWebLink(
@@ -145,7 +154,7 @@ export async function apiCall(route: string, args: {} = undefined) {
   const electron = getElectron();
   if (electron) {
     const resp = await electron.invoke(route.replace('/', '-'), args);
-    return processApiResponse(route, args, resp);
+    return await processApiResponse(route, args, resp);
   } else {
     const resp = await fetch(`${resolveApi()}/${route}`, {
       method: 'POST',
@@ -174,7 +183,7 @@ export async function apiCall(route: string, args: {} = undefined) {
     }
 
     const json = await resp.json();
-    return processApiResponse(route, args, json);
+    return await processApiResponse(route, args, json);
   }
 }
 
diff --git a/plugins/dbgate-plugin-mssql/src/backend/driver.js b/plugins/dbgate-plugin-mssql/src/backend/driver.js
index 055890d3..3aa38f8a 100644
--- a/plugins/dbgate-plugin-mssql/src/backend/driver.js
+++ b/plugins/dbgate-plugin-mssql/src/backend/driver.js
@@ -144,6 +144,9 @@ const driver = {
   getAuthTokenFromCode(connection, options) {
     return azureAuth.azureGetAuthTokenFromCode(options);
   },
+  getAccessTokenFromAuth: (connection, req) => {
+    return req?.user?.msentraToken;
+  },
 };
 
 driver.initialize = dbgateEnv => {