Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

πŸ”’ fix: resolve session persistence post password reset #5077

Merged
merged 6 commits into from
Dec 23, 2024
4 changes: 2 additions & 2 deletions api/cache/banViolation.js
Original file line number Diff line number Diff line change
@@ -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 ?? {};
Expand Down Expand Up @@ -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);
Expand Down
296 changes: 248 additions & 48 deletions api/models/Session.js
Original file line number Diff line number Diff line change
@@ -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,
};
19 changes: 17 additions & 2 deletions api/models/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down Expand Up @@ -75,8 +83,15 @@ module.exports = {
updateToken,
deleteTokens,

createSession,
findSession,
updateExpiration,
deleteSession,
deleteAllUserSessions,
generateRefreshToken,
countActiveSessions,

User,
Key,
Session,
Balance,
};
20 changes: 20 additions & 0 deletions api/models/schema/session.js
Original file line number Diff line number Diff line change
@@ -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;
Loading
Loading