🛡️ feat: Model Validation Middleware (#1841)

* refactor: add ViolationTypes enum and add new violation for illegal model requests

* feat: validateModel middleware to protect the backend against illicit requests for unlisted models
This commit is contained in:
Danny Avila 2024-02-19 22:47:39 -05:00 committed by GitHub
parent d8038e3b19
commit a8a19c6caa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 539 additions and 377 deletions

View file

@ -3,6 +3,7 @@ const checkBan = require('./checkBan');
const uaParser = require('./uaParser');
const setHeaders = require('./setHeaders');
const loginLimiter = require('./loginLimiter');
const validateModel = require('./validateModel');
const requireJwtAuth = require('./requireJwtAuth');
const uploadLimiters = require('./uploadLimiters');
const registerLimiter = require('./registerLimiter');
@ -32,6 +33,7 @@ module.exports = {
validateMessageReq,
buildEndpointOption,
validateRegistration,
validateModel,
moderateText,
noIndex,
};

View file

@ -1,5 +1,5 @@
const rateLimit = require('express-rate-limit');
const { CacheKeys } = require('librechat-data-provider');
const { ViolationTypes } = require('librechat-data-provider');
const logViolation = require('~/cache/logViolation');
const getEnvironmentVariables = () => {
@ -35,7 +35,7 @@ const createFileUploadHandler = (ip = true) => {
} = getEnvironmentVariables();
return async (req, res) => {
const type = CacheKeys.FILE_UPLOAD_LIMIT;
const type = ViolationTypes.FILE_UPLOAD_LIMIT;
const errorMessage = {
type,
max: ip ? fileUploadIpMax : fileUploadUserMax,

View file

@ -0,0 +1,50 @@
const { EModelEndpoint, CacheKeys, ViolationTypes } = require('librechat-data-provider');
const { logViolation, getLogStores } = require('~/cache');
const { handleError } = require('~/server/utils');
/**
* Validates the model of the request.
*
* @async
* @param {Express.Request} req - The Express request object.
* @param {Express.Response} res - The Express response object.
* @param {Function} next - The Express next function.
*/
const validateModel = async (req, res, next) => {
const { model, endpoint } = req.body;
if (!model) {
return handleError(res, { text: 'Model not provided' });
}
const cache = getLogStores(CacheKeys.CONFIG_STORE);
const modelsConfig = await cache.get(CacheKeys.MODELS_CONFIG);
if (!modelsConfig) {
return handleError(res, { text: 'Models not loaded' });
}
const availableModels = modelsConfig[endpoint];
if (!availableModels) {
return handleError(res, { text: 'Endpoint models not loaded' });
}
let validModel = !!availableModels.find((availableModel) => availableModel === model);
if (endpoint === EModelEndpoint.gptPlugins) {
validModel = validModel && availableModels.includes(req.body.agentOptions?.model);
}
if (validModel) {
return next();
}
const { ILLEGAL_MODEL_REQ_SCORE: score = 5 } = process.env ?? {};
const type = ViolationTypes.ILLEGAL_MODEL_REQUEST;
const errorMessage = {
type,
};
await logViolation(req, res, type, errorMessage, score);
return handleError(res, { text: 'Illegal model request' });
};
module.exports = validateModel;