🔒 fix: resolve session persistence post password reset (#5077)

* ✨ feat: Implement session management with CRUD operations and integrate into user workflows * ✨ refactor: Update session model import paths and enhance session creation logic in AuthService * ✨ refactor: Validate session and user ID formats in session management functions * ✨ style: Enhance UI components with improved styling and accessibility features * chore: Update login form tests to use getByTestId instead of getByRole, remove console.log() * chore: Update login form tests to use getByTestId instead of getByRole --------- Co-authored-by: Danny Avila <[email protected]>
danny-avila · Dec 23, 2024 · bdb222d · bdb222d
1 parent 9bca2ae
commit bdb222d
Show file tree

Hide file tree

Showing 17 changed files with 401 additions and 115 deletions.
diff --git a/api/cache/banViolation.js b/api/cache/banViolation.js
@@ -1,7 +1,7 @@
 const { ViolationTypes } = require('librechat-data-provider');
 const { isEnabled, math, removePorts } = require('~/server/utils');
+const { deleteAllUserSessions } = require('~/models');
 const getLogStores = require('./getLogStores');
-const Session = require('~/models/Session');
 const { logger } = require('~/config');
 
 const { BAN_VIOLATIONS, BAN_INTERVAL } = process.env ?? {};
@@ -46,7 +46,7 @@ const banViolation = async (req, res, errorMessage) => {
     return;
   }
 
-  await Session.deleteAllUserSessions(user_id);
+  await deleteAllUserSessions({ userId: user_id });
   res.clearCookie('refreshToken');
 
   const banLogs = getLogStores(ViolationTypes.BAN);

diff --git a/api/models/Session.js b/api/models/Session.js
@@ -1,75 +1,275 @@
 const mongoose = require('mongoose');
 const signPayload = require('~/server/services/signPayload');
 const { hashToken } = require('~/server/utils/crypto');
+const sessionSchema = require('./schema/session');
 const { logger } = require('~/config');
 
+const Session = mongoose.model('Session', sessionSchema);
+
 const { REFRESH_TOKEN_EXPIRY } = process.env ?? {};
-const expires = eval(REFRESH_TOKEN_EXPIRY) ?? 1000 * 60 * 60 * 24 * 7;
-
-const sessionSchema = mongoose.Schema({
-  refreshTokenHash: {
-    type: String,
-    required: true,
-  },
-  expiration: {
-    type: Date,
-    required: true,
-    expires: 0,
-  },
-  user: {
-    type: mongoose.Schema.Types.ObjectId,
-    ref: 'User',
-    required: true,
-  },
-});
-
-sessionSchema.methods.generateRefreshToken = async function () {
+const expires = eval(REFRESH_TOKEN_EXPIRY) ?? 1000 * 60 * 60 * 24 * 7; // 7 days default
+
+/**
+ * Error class for Session-related errors
+ */
+class SessionError extends Error {
+  constructor(message, code = 'SESSION_ERROR') {
+    super(message);
+    this.name = 'SessionError';
+    this.code = code;
+  }
+}
+
+/**
+ * Creates a new session for a user
+ * @param {string} userId - The ID of the user
+ * @param {Object} options - Additional options for session creation
+ * @param {Date} options.expiration - Custom expiration date
+ * @returns {Promise<{session: Session, refreshToken: string}>}
+ * @throws {SessionError}
+ */
+const createSession = async (userId, options = {}) => {
+  if (!userId) {
+    throw new SessionError('User ID is required', 'INVALID_USER_ID');
+  }
+
+  try {
+    const session = new Session({
+      user: userId,
+      expiration: options.expiration || new Date(Date.now() + expires),
+    });
+    const refreshToken = await generateRefreshToken(session);
+    return { session, refreshToken };
+  } catch (error) {
+    logger.error('[createSession] Error creating session:', error);
+    throw new SessionError('Failed to create session', 'CREATE_SESSION_FAILED');
+  }
+};
+
+/**
+ * Finds a session by various parameters
+ * @param {Object} params - Search parameters
+ * @param {string} [params.refreshToken] - The refresh token to search by
+ * @param {string} [params.userId] - The user ID to search by
+ * @param {string} [params.sessionId] - The session ID to search by
+ * @param {Object} [options] - Additional options
+ * @param {boolean} [options.lean=true] - Whether to return plain objects instead of documents
+ * @returns {Promise<Session|null>}
+ * @throws {SessionError}
+ */
+const findSession = async (params, options = { lean: true }) => {
   try {
-    let expiresIn;
-    if (this.expiration) {
-      expiresIn = this.expiration.getTime();
-    } else {
-      expiresIn = Date.now() + expires;
-      this.expiration = new Date(expiresIn);
+    const query = {};
+
+    if (!params.refreshToken && !params.userId && !params.sessionId) {
+      throw new SessionError('At least one search parameter is required', 'INVALID_SEARCH_PARAMS');
     }
 
-    const refreshToken = await signPayload({
-      payload: { id: this.user },
-      secret: process.env.JWT_REFRESH_SECRET,
-      expirationTime: Math.floor((expiresIn - Date.now()) / 1000),
-    });
+    if (params.refreshToken) {
+      const tokenHash = await hashToken(params.refreshToken);
+      query.refreshTokenHash = tokenHash;
+    }
 
-    this.refreshTokenHash = await hashToken(refreshToken);
+    if (params.userId) {
+      query.user = params.userId;
+    }
 
-    await this.save();
+    if (params.sessionId) {
+      const sessionId = params.sessionId.sessionId || params.sessionId;
+      if (!mongoose.Types.ObjectId.isValid(sessionId)) {
+        throw new SessionError('Invalid session ID format', 'INVALID_SESSION_ID');
+      }
+      query._id = sessionId;
+    }
 
-    return refreshToken;
+    // Add expiration check to only return valid sessions
+    query.expiration = { $gt: new Date() };
+
+    const sessionQuery = Session.findOne(query);
+
+    if (options.lean) {
+      return await sessionQuery.lean();
+    }
+
+    return await sessionQuery.exec();
+  } catch (error) {
+    logger.error('[findSession] Error finding session:', error);
+    throw new SessionError('Failed to find session', 'FIND_SESSION_FAILED');
+  }
+};
+
+/**
+ * Updates session expiration
+ * @param {Session|string} session - The session or session ID to update
+ * @param {Date} [newExpiration] - Optional new expiration date
+ * @returns {Promise<Session>}
+ * @throws {SessionError}
+ */
+const updateExpiration = async (session, newExpiration) => {
+  try {
+    const sessionDoc = typeof session === 'string' ? await Session.findById(session) : session;
+
+    if (!sessionDoc) {
+      throw new SessionError('Session not found', 'SESSION_NOT_FOUND');
+    }
+
+    sessionDoc.expiration = newExpiration || new Date(Date.now() + expires);
+    return await sessionDoc.save();
+  } catch (error) {
+    logger.error('[updateExpiration] Error updating session:', error);
+    throw new SessionError('Failed to update session expiration', 'UPDATE_EXPIRATION_FAILED');
+  }
+};
+
+/**
+ * Deletes a session by refresh token or session ID
+ * @param {Object} params - Delete parameters
+ * @param {string} [params.refreshToken] - The refresh token of the session to delete
+ * @param {string} [params.sessionId] - The ID of the session to delete
+ * @returns {Promise<Object>}
+ * @throws {SessionError}
+ */
+const deleteSession = async (params) => {
+  try {
+    if (!params.refreshToken && !params.sessionId) {
+      throw new SessionError(
+        'Either refreshToken or sessionId is required',
+        'INVALID_DELETE_PARAMS',
+      );
+    }
+
+    const query = {};
+
+    if (params.refreshToken) {
+      query.refreshTokenHash = await hashToken(params.refreshToken);
+    }
+
+    if (params.sessionId) {
+      query._id = params.sessionId;
+    }
+
+    const result = await Session.deleteOne(query);
+
+    if (result.deletedCount === 0) {
+      logger.warn('[deleteSession] No session found to delete');
+    }
+
+    return result;
   } catch (error) {
-    logger.error(
-      'Error generating refresh token. Is a `JWT_REFRESH_SECRET` set in the .env file?\n\n',
-      error,
-    );
-    throw error;
+    logger.error('[deleteSession] Error deleting session:', error);
+    throw new SessionError('Failed to delete session', 'DELETE_SESSION_FAILED');
   }
 };
 
-sessionSchema.statics.deleteAllUserSessions = async function (userId) {
+/**
+ * Deletes all sessions for a user
+ * @param {string} userId - The ID of the user
+ * @param {Object} [options] - Additional options
+ * @param {boolean} [options.excludeCurrentSession] - Whether to exclude the current session
+ * @param {string} [options.currentSessionId] - The ID of the current session to exclude
+ * @returns {Promise<Object>}
+ * @throws {SessionError}
+ */
+const deleteAllUserSessions = async (userId, options = {}) => {
   try {
     if (!userId) {
-      return;
+      throw new SessionError('User ID is required', 'INVALID_USER_ID');
+    }
+
+    // Extract userId if it's passed as an object
+    const userIdString = userId.userId || userId;
+
+    if (!mongoose.Types.ObjectId.isValid(userIdString)) {
+      throw new SessionError('Invalid user ID format', 'INVALID_USER_ID_FORMAT');
+    }
+
+    const query = { user: userIdString };
+
+    if (options.excludeCurrentSession && options.currentSessionId) {
+      query._id = { $ne: options.currentSessionId };
     }
-    const result = await this.deleteMany({ user: userId });
-    if (result && result?.deletedCount > 0) {
+
+    const result = await Session.deleteMany(query);
+
+    if (result.deletedCount > 0) {
       logger.debug(
-        `[deleteAllUserSessions] Deleted ${result.deletedCount} sessions for user ${userId}.`,
+        `[deleteAllUserSessions] Deleted ${result.deletedCount} sessions for user ${userIdString}.`,
       );
     }
+
+    return result;
   } catch (error) {
-    logger.error('[deleteAllUserSessions] Error in deleting user sessions:', error);
-    throw error;
+    logger.error('[deleteAllUserSessions] Error deleting user sessions:', error);
+    throw new SessionError('Failed to delete user sessions', 'DELETE_ALL_SESSIONS_FAILED');
   }
 };
 
-const Session = mongoose.model('Session', sessionSchema);
+/**
+ * Generates a refresh token for a session
+ * @param {Session} session - The session to generate a token for
+ * @returns {Promise<string>}
+ * @throws {SessionError}
+ */
+const generateRefreshToken = async (session) => {
+  if (!session || !session.user) {
+    throw new SessionError('Invalid session object', 'INVALID_SESSION');
+  }
+
+  try {
+    const expiresIn = session.expiration ? session.expiration.getTime() : Date.now() + expires;
+
+    if (!session.expiration) {
+      session.expiration = new Date(expiresIn);
+    }
 
-module.exports = Session;
+    const refreshToken = await signPayload({
+      payload: {
+        id: session.user,
+        sessionId: session._id,
+      },
+      secret: process.env.JWT_REFRESH_SECRET,
+      expirationTime: Math.floor((expiresIn - Date.now()) / 1000),
+    });
+
+    session.refreshTokenHash = await hashToken(refreshToken);
+    await session.save();
+
+    return refreshToken;
+  } catch (error) {
+    logger.error('[generateRefreshToken] Error generating refresh token:', error);
+    throw new SessionError('Failed to generate refresh token', 'GENERATE_TOKEN_FAILED');
+  }
+};
+
+/**
+ * Counts active sessions for a user
+ * @param {string} userId - The ID of the user
+ * @returns {Promise<number>}
+ * @throws {SessionError}
+ */
+const countActiveSessions = async (userId) => {
+  try {
+    if (!userId) {
+      throw new SessionError('User ID is required', 'INVALID_USER_ID');
+    }
+
+    return await Session.countDocuments({
+      user: userId,
+      expiration: { $gt: new Date() },
+    });
+  } catch (error) {
+    logger.error('[countActiveSessions] Error counting active sessions:', error);
+    throw new SessionError('Failed to count active sessions', 'COUNT_SESSIONS_FAILED');
+  }
+};
+
+module.exports = {
+  createSession,
+  findSession,
+  updateExpiration,
+  deleteSession,
+  deleteAllUserSessions,
+  generateRefreshToken,
+  countActiveSessions,
+  SessionError,
+};
diff --git a/api/models/index.js b/api/models/index.js
@@ -26,10 +26,18 @@ const {
   deleteMessagesSince,
   deleteMessages,
 } = require('./Message');
+const {
+  createSession,
+  findSession,
+  updateExpiration,
+  deleteSession,
+  deleteAllUserSessions,
+  generateRefreshToken,
+  countActiveSessions,
+} = require('./Session');
 const { getConvoTitle, getConvo, saveConvo, deleteConvos } = require('./Conversation');
 const { getPreset, getPresets, savePreset, deletePresets } = require('./Preset');
 const { createToken, findToken, updateToken, deleteTokens } = require('./Token');
-const Session = require('./Session');
 const Balance = require('./Balance');
 const User = require('./User');
 const Key = require('./Key');
@@ -75,8 +83,15 @@ module.exports = {
   updateToken,
   deleteTokens,
 
+  createSession,
+  findSession,
+  updateExpiration,
+  deleteSession,
+  deleteAllUserSessions,
+  generateRefreshToken,
+  countActiveSessions,
+
   User,
   Key,
-  Session,
   Balance,
 };
diff --git a/api/models/schema/session.js b/api/models/schema/session.js
@@ -0,0 +1,20 @@
+const mongoose = require('mongoose');
+
+const sessionSchema = mongoose.Schema({
+  refreshTokenHash: {
+    type: String,
+    required: true,
+  },
+  expiration: {
+    type: Date,
+    required: true,
+    expires: 0,
+  },
+  user: {
+    type: mongoose.Schema.Types.ObjectId,
+    ref: 'User',
+    required: true,
+  },
+});
+
+module.exports = sessionSchema;