From de6108f95d57d0e571acb2b64f1fe0f94ab6008e Mon Sep 17 00:00:00 2001
From: zadam <zadam.apps@gmail.com>
Date: Fri, 7 May 2021 22:23:49 +0200
Subject: [PATCH] protected session is now global application state to avoid
 weird issues with multiple tabs/windows/reloads

---
 src/public/app/services/app_context.js        |  2 -
 src/public/app/services/glob.js               |  2 -
 src/public/app/services/protected_session.js  | 42 +++++++++----------
 .../app/services/protected_session_holder.js  | 23 +++-------
 src/public/app/widgets/standard_top_widget.js |  6 ++-
 .../app/widgets/type_widgets/editable_text.js |  4 +-
 src/routes/api/login.js                       | 14 +++----
 src/routes/index.js                           |  4 +-
 src/routes/routes.js                          |  1 -
 src/services/protected_session.js             | 32 +++-----------
 src/views/desktop.ejs                         |  1 +
 11 files changed, 49 insertions(+), 82 deletions(-)

diff --git a/src/public/app/services/app_context.js b/src/public/app/services/app_context.js
index 17d92b58d..0b91cec77 100644
--- a/src/public/app/services/app_context.js
+++ b/src/public/app/services/app_context.js
@@ -119,8 +119,6 @@ const appContext = new AppContext(window.glob.isMainWindow);
 
 // we should save all outstanding changes before the page/app is closed
 $(window).on('beforeunload', () => {
-    protectedSessionHolder.resetSessionCookie();
-
     let allSaved = true;
 
     appContext.beforeUnloadListeners = appContext.beforeUnloadListeners.filter(wr => !!wr.deref());
diff --git a/src/public/app/services/glob.js b/src/public/app/services/glob.js
index 50bfbd75f..5e6670651 100644
--- a/src/public/app/services/glob.js
+++ b/src/public/app/services/glob.js
@@ -67,8 +67,6 @@ function setupGlobs() {
         return false;
     };
 
-    protectedSessionHolder.setProtectedSessionId(null);
-
     for (const appCssNoteId of glob.appCssNoteIds || []) {
         libraryLoader.requireCss(`api/notes/download/${appCssNoteId}`);
     }
diff --git a/src/public/app/services/protected_session.js b/src/public/app/services/protected_session.js
index 42550a14e..e0a63d6e2 100644
--- a/src/public/app/services/protected_session.js
+++ b/src/public/app/services/protected_session.js
@@ -1,16 +1,16 @@
-import utils from './utils.js';
 import server from './server.js';
 import protectedSessionHolder from './protected_session_holder.js';
 import toastService from "./toast.js";
 import ws from "./ws.js";
 import appContext from "./app_context.js";
 import treeCache from "./tree_cache.js";
+import utils from "./utils.js";
 
 let protectedSessionDeferred = null;
 
 async function leaveProtectedSession() {
     if (protectedSessionHolder.isProtectedSessionAvailable()) {
-        protectedSessionHolder.resetProtectedSession();
+        await protectedSessionHolder.resetProtectedSession();
     }
 }
 
@@ -41,37 +41,37 @@ async function reloadData() {
 }
 
 async function setupProtectedSession(password) {
-    const response = await enterProtectedSessionOnServer(password);
+    const response = await server.post('login/protected', { password: password });
 
     if (!response.success) {
         toastService.showError("Wrong password.", 3000);
         return;
     }
 
-    protectedSessionHolder.setProtectedSessionId(response.protectedSessionId);
-    protectedSessionHolder.touchProtectedSession();
+    protectedSessionHolder.enableProtectedSession();
+}
 
-    await reloadData();
+ws.subscribeToMessages(async message => {
+    if (message.type === 'protectedSessionLogin') {
+        await reloadData();
 
-    await appContext.triggerEvent('treeCacheReloaded');
+        await appContext.triggerEvent('treeCacheReloaded');
 
-    appContext.triggerEvent('protectedSessionStarted');
+        appContext.triggerEvent('protectedSessionStarted');
 
-    if (protectedSessionDeferred !== null) {
-        import("../dialogs/protected_session.js").then(dialog => dialog.close());
+        if (protectedSessionDeferred !== null) {
+            import("../dialogs/protected_session.js").then(dialog => dialog.close());
 
-        protectedSessionDeferred.resolve(true);
-        protectedSessionDeferred = null;
+            protectedSessionDeferred.resolve(true);
+            protectedSessionDeferred = null;
+        }
+
+        toastService.showMessage("Protected session has been started.");
     }
-
-    toastService.showMessage("Protected session has been started.");
-}
-
-async function enterProtectedSessionOnServer(password) {
-    return await server.post('login/protected', {
-        password: password
-    });
-}
+    else if (message.type === 'protectedSessionLogout') {
+        utils.reloadApp();
+    }
+});
 
 async function protectNote(noteId, protect, includingSubtree) {
     await enterProtectedSession();
diff --git a/src/public/app/services/protected_session_holder.js b/src/public/app/services/protected_session_holder.js
index be2dfc500..86ff20b2b 100644
--- a/src/public/app/services/protected_session_holder.js
+++ b/src/public/app/services/protected_session_holder.js
@@ -1,9 +1,6 @@
-import utils from "./utils.js";
 import options from './options.js';
 import server from "./server.js";
 
-const PROTECTED_SESSION_ID_KEY = 'protectedSessionId';
-
 let lastProtectedSessionOperationDate = 0;
 
 setInterval(() => {
@@ -15,32 +12,23 @@ setInterval(() => {
     }
 }, 10000);
 
-function setProtectedSessionId(id) {
-    // using session cookie so that it disappears after browser/tab is closed
-    utils.setSessionCookie(PROTECTED_SESSION_ID_KEY, id);
-}
+function enableProtectedSession() {
+    glob.isProtectedSessionAvailable = true;
 
-function resetSessionCookie() {
-    utils.setSessionCookie(PROTECTED_SESSION_ID_KEY, null);
+    touchProtectedSession();
 }
 
 async function resetProtectedSession() {
-    resetSessionCookie();
-
     await server.post("logout/protected");
-
-    utils.reloadApp();
 }
 
 function isProtectedSessionAvailable() {
-    return !!utils.getCookie(PROTECTED_SESSION_ID_KEY);
+    return glob.isProtectedSessionAvailable;
 }
 
 function touchProtectedSession() {
     if (isProtectedSessionAvailable()) {
         lastProtectedSessionOperationDate = Date.now();
-
-        setProtectedSessionId(utils.getCookie(PROTECTED_SESSION_ID_KEY));
     }
 }
 
@@ -51,8 +39,7 @@ function touchProtectedSessionIfNecessary(note) {
 }
 
 export default {
-    setProtectedSessionId,
-    resetSessionCookie,
+    enableProtectedSession,
     resetProtectedSession,
     isProtectedSessionAvailable,
     touchProtectedSession,
diff --git a/src/public/app/widgets/standard_top_widget.js b/src/public/app/widgets/standard_top_widget.js
index 0b14723f6..5e346ef9c 100644
--- a/src/public/app/widgets/standard_top_widget.js
+++ b/src/public/app/widgets/standard_top_widget.js
@@ -1,5 +1,6 @@
 import BasicWidget from "./basic_widget.js";
 import HistoryNavigationWidget from "./history_navigation.js";
+import protectedSessionHolder from "../services/protected_session_holder.js";
 import protectedSessionService from "../services/protected_session.js";
 import QuickSearchWidget from "./quick_search.js";
 
@@ -68,8 +69,7 @@ const TPL = `
         </button>
     
         <button class="btn btn-sm leave-protected-session-button noborder"
-                title="Leave protected session so that protected notes are not accessible any more."
-                style="display: none;">
+                title="Leave protected session so that protected notes are not accessible any more.">
             <span class="bx bx-log-out"></span>
     
             Leave protected session
@@ -96,9 +96,11 @@ export default class StandardTopWidget extends BasicWidget {
 
         this.$enterProtectedSessionButton = this.$widget.find(".enter-protected-session-button");
         this.$enterProtectedSessionButton.on('click', protectedSessionService.enterProtectedSession);
+        this.$enterProtectedSessionButton.toggle(!protectedSessionHolder.isProtectedSessionAvailable());
 
         this.$leaveProtectedSessionButton = this.$widget.find(".leave-protected-session-button");
         this.$leaveProtectedSessionButton.on('click', protectedSessionService.leaveProtectedSession);
+        this.$leaveProtectedSessionButton.toggle(protectedSessionHolder.isProtectedSessionAvailable());
 
         return this.$widget;
     }
diff --git a/src/public/app/widgets/type_widgets/editable_text.js b/src/public/app/widgets/type_widgets/editable_text.js
index cd429cb3f..876272eb7 100644
--- a/src/public/app/widgets/type_widgets/editable_text.js
+++ b/src/public/app/widgets/type_widgets/editable_text.js
@@ -8,7 +8,7 @@ import treeService from "../../services/tree.js";
 import noteCreateService from "../../services/note_create.js";
 import AbstractTextTypeWidget from "./abstract_text_type_widget.js";
 
-const ENABLE_INSPECTOR = true;
+const ENABLE_INSPECTOR = false;
 
 const mentionSetup = {
     feeds: [
@@ -121,7 +121,7 @@ export default class EditableTextTypeWidget extends AbstractTextTypeWidget {
         });
 
         this.textEditor.model.document.on('change:data', () => this.spacedUpdate.scheduleUpdate());
-console.log('glob.isDev', glob.isDev);
+
         if (glob.isDev && ENABLE_INSPECTOR) {
             await import(/* webpackIgnore: true */'../../../libraries/ckeditor/inspector.js');
             CKEditorInspector.attach(this.textEditor);
diff --git a/src/routes/api/login.js b/src/routes/api/login.js
index 159e43be1..268917dc2 100644
--- a/src/routes/api/login.js
+++ b/src/routes/api/login.js
@@ -8,11 +8,11 @@ const passwordEncryptionService = require('../../services/password_encryption');
 const protectedSessionService = require('../../services/protected_session');
 const appInfo = require('../../services/app_info');
 const eventService = require('../../services/events');
-const cls = require('../../services/cls');
 const sqlInit = require('../../services/sql_init');
 const sql = require('../../services/sql');
 const optionService = require('../../services/options');
 const ApiToken = require('../../entities/api_token');
+const ws = require("../../services/ws.js");
 
 function loginSync(req) {
     if (!sqlInit.schemaExists()) {
@@ -65,16 +65,14 @@ function loginToProtectedSession(req) {
 
     const decryptedDataKey = passwordEncryptionService.getDataKey(password);
 
-    const protectedSessionId = protectedSessionService.setDataKey(decryptedDataKey);
-
-    // this is set here so that event handlers have access to the protected session
-    cls.set('protectedSessionId', protectedSessionId);
+    protectedSessionService.setDataKey(decryptedDataKey);
 
     eventService.emit(eventService.ENTER_PROTECTED_SESSION);
 
+    ws.sendMessageToAllClients({ type: 'protectedSessionLogin' });
+
     return {
-        success: true,
-        protectedSessionId: protectedSessionId
+        success: true
     };
 }
 
@@ -82,6 +80,8 @@ function logoutFromProtectedSession() {
     protectedSessionService.resetDataKey();
 
     eventService.emit(eventService.LEAVE_PROTECTED_SESSION);
+
+    ws.sendMessageToAllClients({ type: 'protectedSessionLogout' });
 }
 
 function token(req) {
diff --git a/src/routes/index.js b/src/routes/index.js
index 0de11466a..2a05fc602 100644
--- a/src/routes/index.js
+++ b/src/routes/index.js
@@ -7,6 +7,7 @@ const config = require('../services/config');
 const optionService = require('../services/options');
 const log = require('../services/log');
 const env = require('../services/env');
+const protectedSessionService = require("../services/protected_session.js");
 
 function index(req, res) {
     const options = optionService.getOptionsMap();
@@ -30,7 +31,8 @@ function index(req, res) {
         appCssNoteIds: getAppCssNoteIds(),
         isDev: env.isDev(),
         isMainWindow: !req.query.extra,
-        extraHoistedNoteId: req.query.extraHoistedNoteId
+        extraHoistedNoteId: req.query.extraHoistedNoteId,
+        isProtectedSessionAvailable: protectedSessionService.isProtectedSessionAvailable()
     });
 }
 
diff --git a/src/routes/routes.js b/src/routes/routes.js
index d12f1b003..0f109aa3d 100644
--- a/src/routes/routes.js
+++ b/src/routes/routes.js
@@ -93,7 +93,6 @@ function route(method, path, middleware, routeHandler, resultHandler, transactio
                 cls.set('sourceId', req.headers['trilium-source-id']);
                 cls.set('localNowDateTime', req.headers['trilium-local-now-datetime']);
                 cls.set('hoistedNoteId', req.headers['trilium-hoisted-note-id'] || 'root');
-                protectedSessionService.setProtectedSessionId(req);
 
                 const cb = () => routeHandler(req, res, next);
 
diff --git a/src/services/protected_session.js b/src/services/protected_session.js
index b86af84e6..7dfddc2df 100644
--- a/src/services/protected_session.js
+++ b/src/services/protected_session.js
@@ -1,42 +1,24 @@
 "use strict";
 
-const utils = require('./utils');
 const log = require('./log');
 const dataEncryptionService = require('./data_encryption');
-const cls = require('./cls');
 
-let dataKeyMap = {};
+let dataKey = null;
 
 function setDataKey(decryptedDataKey) {
-    const protectedSessionId = utils.randomSecureToken(32);
-
-    dataKeyMap[protectedSessionId] = Array.from(decryptedDataKey); // can't store buffer in session
-
-    return protectedSessionId;
-}
-
-function setProtectedSessionId(req) {
-    cls.set('protectedSessionId', req.cookies.protectedSessionId);
-}
-
-function getProtectedSessionId() {
-    return cls.get('protectedSessionId');
+    dataKey = Array.from(decryptedDataKey);
 }
 
 function getDataKey() {
-    const protectedSessionId = getProtectedSessionId();
-
-    return dataKeyMap[protectedSessionId];
+    return dataKey;
 }
 
 function resetDataKey() {
-    dataKeyMap = {};
+    dataKey = null;
 }
 
 function isProtectedSessionAvailable() {
-    const protectedSessionId = getProtectedSessionId();
-
-    return !!dataKeyMap[protectedSessionId];
+    return !!dataKey;
 }
 
 function decryptNotes(notes) {
@@ -74,12 +56,10 @@ function decryptString(cipherText) {
 
 module.exports = {
     setDataKey,
-    getDataKey,
     resetDataKey,
     isProtectedSessionAvailable,
     encrypt,
     decrypt,
     decryptString,
-    decryptNotes,
-    setProtectedSessionId
+    decryptNotes
 };
diff --git a/src/views/desktop.ejs b/src/views/desktop.ejs
index de9adfd94..0253cadeb 100644
--- a/src/views/desktop.ejs
+++ b/src/views/desktop.ejs
@@ -56,6 +56,7 @@
         appCssNoteIds: <%- JSON.stringify(appCssNoteIds) %>,
         isMainWindow: <%= isMainWindow %>,
         extraHoistedNoteId: '<%= extraHoistedNoteId %>',
+        isProtectedSessionAvailable: <%= isProtectedSessionAvailable %>,
     };
 </script>