From 90a02ebc2f6e769e96813197c795ed7cf7ad7628 Mon Sep 17 00:00:00 2001
From: Jakob Ackermann <jakob.ackermann@overleaf.com>
Date: Thu, 25 Apr 2024 13:54:01 +0100
Subject: [PATCH] Merge pull request #17949 from overleaf/jpa-set-nx-xx

[web] stricter writes to redis when creating and updating sessions

GitOrigin-RevId: 79723e0d38884bf723c7a2ba32993e4daa2612a0
---
 .../src/infrastructure/CustomSessionStore.js  | 45 ++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/services/web/app/src/infrastructure/CustomSessionStore.js b/services/web/app/src/infrastructure/CustomSessionStore.js
index 522ecfde86..ee8623ac0e 100644
--- a/services/web/app/src/infrastructure/CustomSessionStore.js
+++ b/services/web/app/src/infrastructure/CustomSessionStore.js
@@ -2,14 +2,30 @@ const session = require('express-session')
 const RedisStore = require('connect-redis')(session)
 const metrics = require('@overleaf/metrics')
 const logger = require('@overleaf/logger')
+const Settings = require('@overleaf/settings')
 const SessionManager = require('../Features/Authentication/SessionManager')
 
 const MAX_SESSION_SIZE_THRESHOLD = 4096
 
 // Define a custom session store to record session metrics and log large
 // anonymous sessions for debugging purposes
+// Also make the SET calls more robust/consistent by adding flags
+// - XX: ensure update in place, expect that the old session value is still in redis at that key
+// - NX: ensure initial set, expect that there is no other session at that key already
 class CustomSessionStore extends RedisStore {
   static largestSessionSize = 3 * 1024 // ignore sessions smaller than 3KB
+  #initialSetStore
+  #updateInPlaceStore
+
+  constructor({ client }) {
+    super({ client })
+    this.#initialSetStore = new RedisStore({
+      client: new CustomSetRedisClient(client, 'NX'),
+    })
+    this.#updateInPlaceStore = new RedisStore({
+      client: new CustomSetRedisClient(client, 'XX'),
+    })
+  }
 
   static metric(method, sess) {
     let type // type of session: 'logged-in', 'anonymous', or 'na' (not available)
@@ -51,7 +67,14 @@ class CustomSessionStore extends RedisStore {
 
   set(sid, sess, cb) {
     CustomSessionStore.metric('set', sess)
-    super.set(sid, sess, cb)
+    const originalId = sess.req.signedCookies[Settings.cookieName]
+    if (sid === originalId || sid === sess.req.newSessionId) {
+      this.#updateInPlaceStore.set(sid, sess, cb)
+    } else {
+      // Multiple writes can get issued with the new sid. Keep track of it.
+      Object.defineProperty(sess.req, 'newSessionId', { value: sid })
+      this.#initialSetStore.set(sid, sess, cb)
+    }
   }
 
   touch(sid, sess, cb) {
@@ -81,4 +104,24 @@ function redactSession(sess) {
   )
 }
 
+class CustomSetRedisClient {
+  #client
+  #flag
+  constructor(client, flag) {
+    this.#client = client
+    this.#flag = flag
+  }
+
+  set(args, cb) {
+    args.push(this.#flag)
+    this.#client.set(args, (err, ok) => {
+      metrics.inc('session.store.set', 1, {
+        path: this.#flag,
+        status: err ? 'error' : ok ? 'success' : 'failure',
+      })
+      cb(err, ok)
+    })
+  }
+}
+
 module.exports = CustomSessionStore