🚧 fix: Add Per-User Throttle to 2FA Continuation Attempts (#13583)

* fix: refine auth continuation throttling

* chore: import order

api/server/middleware/limiters/twoFactorTempLimiter.js

@@ -1,5 +1,6 @@
-const rateLimit = require('express-rate-limit');
 const jwt = require('jsonwebtoken');
+const { createHash } = require('crypto');
+const rateLimit = require('express-rate-limit');
 const { ViolationTypes } = require('librechat-data-provider');
 const { limiterCache, removePorts } = require('@librechat/api');
 const { logViolation } = require('~/cache');
@@ -18,6 +19,23 @@ const score = TWO_FACTOR_TEMP_VIOLATION_SCORE ?? LOGIN_VIOLATION_SCORE;
 const windowInMinutes = windowMs / 60000;
 const message = `Too many verification attempts, please try again after ${windowInMinutes} minutes.`;
 
+const hashLimiterKey = (value) => createHash('sha256').update(value).digest('hex');
+
+const getUserLimiterKey = (req) => {
+  const userId = req.user?.id ?? req.user?._id;
+  if (userId) {
+    return `user:${userId.toString()}`;
+  }
+
+  const tempToken = req.body?.tempToken;
+  if (typeof tempToken === 'string' && tempToken) {
+    return `temp:${hashLimiterKey(tempToken)}`;
+  }
+
+  const ip = removePorts(req);
+  return ip ? `ip:${ip}` : 'ip:unknown';
+};
+
 const getTempTokenUserId = (tempToken) => {
   if (!tempToken) {
     return null;
@@ -31,11 +49,12 @@ const getTempTokenUserId = (tempToken) => {
   }
 };
 
-const handler = async (req, res) => {
+const createHandler = (limiter) => async (req, res) => {
   const type = ViolationTypes.LOGINS;
   const errorMessage = {
     type,
     max,
+    limiter,
     windowInMinutes,
   };
 
@@ -50,14 +69,33 @@ const handler = async (req, res) => {
   return res.status(429).json({ message });
 };
 
-const limiterOptions = {
+const ipLimiterOptions = {
   windowMs,
   max,
-  handler,
+  handler: createHandler('ip'),
   keyGenerator: removePorts,
   store: limiterCache('two_factor_temp_limiter'),
 };
 
-const twoFactorTempLimiter = rateLimit(limiterOptions);
+const userLimiterOptions = {
+  windowMs,
+  max,
+  handler: createHandler('user'),
+  keyGenerator: getUserLimiterKey,
+  store: limiterCache('two_factor_temp_user_limiter'),
+};
+
+const twoFactorTempIpLimiter = rateLimit(ipLimiterOptions);
+const twoFactorTempUserLimiter = rateLimit(userLimiterOptions);
+
+const twoFactorTempLimiter = (req, res, next) => {
+  twoFactorTempIpLimiter(req, res, (err) => {
+    if (err) {
+      return next(err);
+    }
+
+    return twoFactorTempUserLimiter(req, res, next);
+  });
+};
 
 module.exports = twoFactorTempLimiter;

api/server/middleware/limiters/twoFactorTempLimiter.test.js

@@ -0,0 +1,111 @@
+const jwt = require('jsonwebtoken');
+const express = require('express');
+const request = require('supertest');
+
+const originalEnv = process.env;
+const jwtSecret = 'test-two-factor-secret';
+
+const createToken = (userId) =>
+  jwt.sign({ userId, twoFAPending: true }, jwtSecret, { expiresIn: '5m' });
+
+const createApp = () => {
+  jest.resetModules();
+  process.env = {
+    ...originalEnv,
+    JWT_SECRET: jwtSecret,
+    LOGIN_MAX: '2',
+    LOGIN_WINDOW: '5',
+    TWO_FACTOR_TEMP_MAX: '2',
+    TWO_FACTOR_TEMP_WINDOW: '5',
+  };
+
+  jest.doMock('@librechat/api', () => ({
+    limiterCache: jest.fn(() => undefined),
+    removePorts: (req) => req?.['ip'],
+  }));
+  jest.doMock('~/cache', () => ({
+    logViolation: jest.fn().mockResolvedValue(undefined),
+  }));
+
+  const setTwoFactorTempUser = require('../setTwoFactorTempUser');
+  const twoFactorTempLimiter = require('./twoFactorTempLimiter');
+  const { logViolation } = require('~/cache');
+
+  const app = express();
+  app.set('trust proxy', 1);
+  app.use(express.json());
+  app.post('/verify', setTwoFactorTempUser, twoFactorTempLimiter, (req, res) =>
+    res.status(204).end(),
+  );
+
+  return { app, logViolation };
+};
+
+describe('twoFactorTempLimiter', () => {
+  afterEach(() => {
+    jest.dontMock('@librechat/api');
+    jest.dontMock('~/cache');
+    process.env = originalEnv;
+  });
+
+  it('limits a valid temp-token user across rotating source IPs', async () => {
+    const { app, logViolation } = createApp();
+    const tempToken = createToken('user-1');
+
+    await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '203.0.113.1')
+      .send({ tempToken, token: '000000' })
+      .expect(204);
+    await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '203.0.113.2')
+      .send({ tempToken, token: '000001' })
+      .expect(204);
+
+    const response = await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '203.0.113.3')
+      .send({ tempToken, token: '000002' })
+      .expect(429);
+
+    expect(response.body).toEqual({
+      message: 'Too many verification attempts, please try again after 5 minutes.',
+    });
+    expect(logViolation).toHaveBeenCalledTimes(1);
+    expect(logViolation.mock.calls[0][0].user).toEqual({ id: 'user-1' });
+    expect(logViolation.mock.calls[0][3]).toMatchObject({
+      limiter: 'user',
+      max: '2',
+      windowInMinutes: 5,
+    });
+  });
+
+  it('keeps the existing source IP limit before the user limit', async () => {
+    const { app, logViolation } = createApp();
+
+    await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '198.51.100.1')
+      .send({ tempToken: createToken('user-a'), token: '000000' })
+      .expect(204);
+    await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '198.51.100.1')
+      .send({ tempToken: createToken('user-b'), token: '000001' })
+      .expect(204);
+
+    await request(app)
+      .post('/verify')
+      .set('X-Forwarded-For', '198.51.100.1')
+      .send({ tempToken: createToken('user-c'), token: '000002' })
+      .expect(429);
+
+    expect(logViolation).toHaveBeenCalledTimes(1);
+    expect(logViolation.mock.calls[0][3]).toMatchObject({
+      limiter: 'ip',
+      max: '2',
+      windowInMinutes: 5,
+    });
+  });
+});