Fix Electron SSO handling to support multiple profiles

Signed-off-by: Michael Telatynski <7t3chguy@gmail.com>

src/electron-main.js

@@ -35,7 +35,7 @@ const tray = require('./tray');
 const vectorMenu = require('./vectormenu');
 const webContentsHandler = require('./webcontents-handler');
 const updater = require('./updater');
-const protocolInit = require('./protocol');
+const {getProfileFromDeeplink, protocolInit, recordSSOSession} = require('./protocol');
 
 const windowStateKeeper = require('electron-window-state');
 const Store = require('electron-store');
@@ -86,7 +86,11 @@ if (argv["help"]) {
     app.exit();
 }
 
-if (argv['profile-dir']) {
+// check if we are passed a profile in the SSO callback url
+const userDataPathInProtocol = getProfileFromDeeplink(argv["_"]);
+if (userDataPathInProtocol) {
+    app.setPath('userData', userDataPathInProtocol);
+} else if (argv['profile-dir']) {
     app.setPath('userData', argv['profile-dir']);
 } else if (argv['profile']) {
     app.setPath('userData', `${app.getPath('userData')}-${argv['profile']}`);
@@ -339,6 +343,9 @@ ipcMain.on('ipcCall', async function(ev, payload) {
             }
             break;
         }
+        case 'startSSOFlow':
+            recordSSOSession(args[0]);
+            break;
 
         default:
             mainWindow.webContents.send('ipcReply', {

src/protocol.js

@@ -14,40 +14,89 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-const {app} = require('electron');
+const {app} = require("electron");
+const path = require("path");
+const fs = require("fs");
+
+const PROTOCOL = "riot://";
+const SEARCH_PARAM = "riot-desktop-ssoid";
+const STORE_FILE_NAME = "sso-sessions.json";
+
+// we getPath userData before electron-main changes it, so this is the default value
+const storePath = path.join(app.getPath("userData"), STORE_FILE_NAME);
 
 const processUrl = (url) => {
     if (!global.mainWindow) return;
     console.log("Handling link: ", url);
-    global.mainWindow.loadURL(url.replace("riot://", "vector://"));
+    global.mainWindow.loadURL(url.replace(PROTOCOL, "vector://"));
 };
 
-module.exports = () => {
-    // get all args except `hidden` as it'd mean the app would not get focused
-    // XXX: passing args to protocol handlers only works on Windows,
-    // so unpackaged deep-linking and --profile passing won't work on Mac/Linux
-    const args = process.argv.slice(1).filter(arg => arg !== "--hidden" && arg !== "-hidden");
-    if (app.isPackaged) {
-        app.setAsDefaultProtocolClient('riot', process.execPath, args);
-    } else if (process.platform === 'win32') { // on Mac/Linux this would just cause the electron binary to open
-        // special handler for running without being packaged, e.g `electron .` by passing our app path to electron
-        app.setAsDefaultProtocolClient('riot', process.execPath, [app.getAppPath(), ...args]);
-    }
-
-    if (process.platform === 'darwin') {
-        // Protocol handler for macos
-        app.on('open-url', function(ev, url) {
-            ev.preventDefault();
-            processUrl(url);
-        });
-    } else {
-        // Protocol handler for win32/Linux
-        app.on('second-instance', (ev, commandLine) => {
-            const url = commandLine[commandLine.length - 1];
-            if (!url.startsWith("riot://")) return;
-            processUrl(url);
-        });
+const readStore = () => {
+    try {
+        const s = fs.readFileSync(storePath, { encoding: "utf8" });
+        const o = JSON.parse(s);
+        return typeof o === "object" ? o : {};
+    } catch (e) {
+        return {};
     }
 };
 
+const writeStore = (data) => {
+    fs.writeFileSync(storePath, JSON.stringify(data));
+};
 
+module.exports = {
+    recordSSOSession: (sessionID) => {
+        const userDataPath = app.getPath('userData');
+        const store = readStore();
+        for (const key in store) {
+            // ensure each instance only has one (the latest) session ID to prevent the file growing unbounded
+            if (store[key] === userDataPath) {
+                delete store[key];
+                break;
+            }
+        }
+        store[sessionID] = userDataPath;
+        writeStore(store);
+    },
+    getProfileFromDeeplink: (args) => {
+        // check if we are passed a profile in the SSO callback url
+        const deeplinkUrl = args.find(arg => arg.startsWith('riot://'));
+        if (deeplinkUrl && deeplinkUrl.includes(SEARCH_PARAM)) {
+            const parsedUrl = new URL(deeplinkUrl);
+            if (parsedUrl.protocol === 'riot:') {
+                const ssoID = parsedUrl.searchParams.get(SEARCH_PARAM);
+                const store = readStore();
+                console.log("Forwarding to profile: ", store[ssoID]);
+                return store[ssoID];
+            }
+        }
+    },
+    protocolInit: () => {
+        // get all args except `hidden` as it'd mean the app would not get focused
+        // XXX: passing args to protocol handlers only works on Windows, so unpackaged deep-linking
+        // --profile/--profile-dir are passed via the SEARCH_PARAM var in the callback url
+        const args = process.argv.slice(1).filter(arg => arg !== "--hidden" && arg !== "-hidden");
+        if (app.isPackaged) {
+            app.setAsDefaultProtocolClient('riot', process.execPath, args);
+        } else if (process.platform === 'win32') { // on Mac/Linux this would just cause the electron binary to open
+            // special handler for running without being packaged, e.g `electron .` by passing our app path to electron
+            app.setAsDefaultProtocolClient('riot', process.execPath, [app.getAppPath(), ...args]);
+        }
+
+        if (process.platform === 'darwin') {
+            // Protocol handler for macos
+            app.on('open-url', function(ev, url) {
+                ev.preventDefault();
+                processUrl(url);
+            });
+        } else {
+            // Protocol handler for win32/Linux
+            app.on('second-instance', (ev, commandLine) => {
+                const url = commandLine[commandLine.length - 1];
+                if (!url.startsWith(PROTOCOL)) return;
+                processUrl(url);
+            });
+        }
+    },
+};