Merge pull request #67 from vector-im/t3chguy/poc_riot_desktop_sso_multi_profile

Fix Electron SSO handling to support multiple profiles

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']}`);
@@ -349,6 +353,9 @@ ipcMain.on('ipcCall', async function(ev, payload) {
             }
             break;
         }
+        case 'startSSOFlow':
+            recordSSOSession(args[0]);
+            break;
 
         default:
             mainWindow.webContents.send('ipcReply', {

src/protocol.js

@@ -14,18 +14,68 @@ 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 = () => {
+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,
+        // XXX: passing args to protocol handlers only works on Windows, so unpackaged deep-linking
-    // so unpackaged deep-linking and --profile passing won't work on Mac/Linux
+        // --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);
@@ -44,10 +94,9 @@ module.exports = () => {
             // Protocol handler for win32/Linux
             app.on('second-instance', (ev, commandLine) => {
                 const url = commandLine[commandLine.length - 1];
-            if (!url.startsWith("riot://")) return;
+                if (!url.startsWith(PROTOCOL)) return;
                 processUrl(url);
             });
         }
+    },
 };
-
-