refactor: fully working E2EE

small issue to fix. when full response is received it replaces the text with the text from the DB. and then the decryption is not yet implement.

api/models/Message.js

@@ -2,6 +2,7 @@ const { z } = require('zod');
 const Message = require('./schema/messageSchema');
 const { logger } = require('~/config');
 
+// Validate conversation ID as a UUID (if your conversation IDs follow UUID format)
 const idSchema = z.string().uuid();
 
 /**
@@ -28,8 +29,11 @@ const idSchema = z.string().uuid();
  * @param {string} [params.plugin] - Plugin associated with the message.
  * @param {string[]} [params.plugins] - An array of plugins associated with the message.
  * @param {string} [params.model] - The model used to generate the message.
- * @param {Object} [metadata] - Additional metadata for this operation
- * @param {string} [metadata.context] - The context of the operation
+ * @param {string} [params.iv] - (Optional) Base64-encoded initialization vector for encryption.
+ * @param {string} [params.authTag] - (Optional) Base64-encoded authentication tag from AES-GCM.
+ * @param {string} [params.encryptedKey] - (Optional) Base64-encoded AES key encrypted with RSA.
+ * @param {Object} [metadata] - Additional metadata for this operation.
+ * @param {string} [metadata.context] - The context of the operation.
  * @returns {Promise<TMessage>} The updated or newly inserted message document.
  * @throws {Error} If there is an error in saving the message.
  */
@@ -51,6 +55,9 @@ async function saveMessage(req, params, metadata) {
       ...params,
       user: req.user.id,
       messageId: params.newMessageId || params.messageId,
+      iv: params.iv ?? null,
+      authTag: params.authTag ?? null,
+      encryptedKey: params.encryptedKey ?? null,
     };
 
     if (req?.body?.isTemporary) {
@@ -90,7 +97,12 @@ async function bulkSaveMessages(messages, overrideTimestamp = false) {
     const bulkOps = messages.map((message) => ({
       updateOne: {
         filter: { messageId: message.messageId },
-        update: message,
+        update: {
+          ...message,
+          iv: message.iv ?? null,
+          authTag: message.authTag ?? null,
+          encryptedKey: message.encryptedKey ?? null,
+        },
         timestamps: !overrideTimestamp,
         upsert: true,
       },
@@ -119,14 +131,7 @@ async function bulkSaveMessages(messages, overrideTimestamp = false) {
  * @returns {Promise<Object>} The updated or newly inserted message document.
  * @throws {Error} If there is an error in saving the message.
  */
-async function recordMessage({
-  user,
-  endpoint,
-  messageId,
-  conversationId,
-  parentMessageId,
-  ...rest
-}) {
+async function recordMessage({ user, endpoint, messageId, conversationId, parentMessageId, ...rest }) {
   try {
     // No parsing of convoId as may use threadId
     const message = {
@@ -136,6 +141,9 @@ async function recordMessage({
       conversationId,
       parentMessageId,
       ...rest,
+      iv: rest.iv ?? null,
+      authTag: rest.authTag ?? null,
+      encryptedKey: rest.encryptedKey ?? null,
     };
 
     return await Message.findOneAndUpdate({ user, messageId }, message, {
@@ -190,12 +198,15 @@ async function updateMessageText(req, { messageId, text }) {
 async function updateMessage(req, message, metadata) {
   try {
     const { messageId, ...update } = message;
+    // Ensure encryption fields are explicitly updated (if provided)
+    update.iv = update.iv ?? null;
+    update.authTag = update.authTag ?? null;
+    update.encryptedKey = update.encryptedKey ?? null;
+
     const updatedMessage = await Message.findOneAndUpdate(
       { messageId, user: req.user.id },
       update,
-      {
-        new: true,
-      },
+      { new: true },
     );
 
     if (!updatedMessage) {
@@ -225,11 +236,11 @@ async function updateMessage(req, message, metadata) {
  *
  * @async
  * @function deleteMessagesSince
- * @param {Object} params - The parameters object.
  * @param {Object} req - The request object.
+ * @param {Object} params - The parameters object.
  * @param {string} params.messageId - The unique identifier for the message.
  * @param {string} params.conversationId - The identifier of the conversation.
- * @returns {Promise<Number>} The number of deleted messages.
+ * @returns {Promise<number>} The number of deleted messages.
  * @throws {Error} If there is an error in deleting messages.
  */
 async function deleteMessagesSince(req, { messageId, conversationId }) {
@@ -263,7 +274,6 @@ async function getMessages(filter, select) {
     if (select) {
       return await Message.find(filter).select(select).sort({ createdAt: 1 }).lean();
     }
-
     return await Message.find(filter).sort({ createdAt: 1 }).lean();
   } catch (err) {
     logger.error('Error getting messages:', err);
@@ -281,10 +291,7 @@ async function getMessages(filter, select) {
  */
 async function getMessage({ user, messageId }) {
   try {
-    return await Message.findOne({
-      user,
-      messageId,
-    }).lean();
+    return await Message.findOne({ user, messageId }).lean();
   } catch (err) {
     logger.error('Error getting message:', err);
     throw err;

api/models/schema/messageSchema.js

@@ -54,11 +54,6 @@ const messageSchema = mongoose.Schema(
       type: String,
       meiliIndex: true,
     },
-    // If the message is encrypted (and stored in 'text'),
-    // then this field should hold the IV used during encryption.
-    messageEncryptionIV: {
-      type: String,
-    },
     summary: {
       type: String,
     },
@@ -142,6 +137,18 @@ const messageSchema = mongoose.Schema(
     expiredAt: {
       type: Date,
     },
+    iv: {
+      type: String,
+      default: null,
+    },
+    authTag: {
+      type: String,
+      default: null,
+    },
+    encryptedKey: {
+      type: String,
+      default: null,
+    },
   },
   { timestamps: true },
 );

api/models/schema/shareSchema.js

@@ -23,13 +23,6 @@ const shareSchema = mongoose.Schema(
       type: Boolean,
       default: true,
     },
-    // --- Field for re-encrypting the conversation key for the forked user ---
-    encryptionKeys: [
-      {
-        user: { type: String, index: true },
-        encryptedKey: { type: String },
-      },
-    ],
   },
   { timestamps: true },
 );

api/models/schema/userSchema.js

@@ -27,10 +27,10 @@ const { SystemRoles } = require('librechat-data-provider');
  * @property {Array} [plugins=[]] - List of plugins used by the user
  * @property {Array.<MongoSession>} [refreshToken] - List of sessions with refresh tokens
  * @property {Date} [expiresAt] - Optional expiration date of the file
- * @property {string} [encryptionPublicKey] - The user's public key for E2EE (client-generated)
- * @property {string} [encryptedPrivateKey] - The user's private key encrypted with a user-defined passphrase
- * @property {string} [encryptionSalt] - The salt used for PBKDF2 during encryption
- * @property {string} [encryptionIV] - The initialization vector used for encryption (AES-GCM)
+ * @property {string} [encryptionPublicKey] - The user's encryption public key
+ * @property {string} [encryptedPrivateKey] - The user's encrypted private key
+ * @property {string} [encryptionSalt] - The salt used for key derivation (e.g., PBKDF2)
+ * @property {string} [encryptionIV] - The IV used for encrypting the private key
  * @property {Date} [createdAt] - Date when the user was created (added by timestamps)
  * @property {Date} [updatedAt] - Date when the user was last updated (added by timestamps)
  */
@@ -136,26 +136,21 @@ const userSchema = mongoose.Schema(
       type: Boolean,
       default: false,
     },
-    // --- New Fields for E2EE ---
     encryptionPublicKey: {
       type: String,
-      required: false,
-      // Provided by the client after key generation.
+      default: null,
     },
     encryptedPrivateKey: {
       type: String,
-      required: false,
-      // The private key encrypted on the client with the user’s encryption passphrase.
+      default: null,
     },
     encryptionSalt: {
       type: String,
-      required: false,
-      // Salt used for PBKDF2 when encrypting the private key.
+      default: null,
     },
     encryptionIV: {
       type: String,
-      required: false,
-      // IV used for AES-GCM encryption of the private key.
+      default: null,
     },
   },
 

api/server/controllers/AskController.js

@@ -4,6 +4,13 @@ const { sendMessage, createOnProgress } = require('~/server/utils');
 const { saveMessage } = require('~/models');
 const { logger } = require('~/config');
 
+let crypto;
+try {
+  crypto = require('crypto');
+} catch (err) {
+  logger.error('[openidStrategy] crypto support is disabled!', err);
+}
+
 const AskController = async (req, res, next, initializeClient, addTitle) => {
   let {
     text,
@@ -74,14 +81,8 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
 
     res.on('close', () => {
       logger.debug('[AskController] Request closed');
-      if (!abortController) {
-        return;
-      } else if (abortController.signal.aborted) {
-        return;
-      } else if (abortController.requestCompleted) {
-        return;
-      }
-
+      if (!abortController) {return;}
+      if (abortController.signal.aborted || abortController.requestCompleted) {return;}
       abortController.abort();
       logger.debug('[AskController] Request aborted on close');
     });
@@ -95,10 +96,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       onStart,
       abortController,
       progressCallback,
-      progressOptions: {
-        res,
-        // parentMessageId: overrideParentMessageId || userMessageId,
-      },
+      progressOptions: { res },
     };
 
     /** @type {TMessage} */
@@ -115,6 +113,58 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       delete userMessage.image_urls;
     }
 
+    // --- Encryption Branch ---
+    // Only encrypt if the user has set up encryption (i.e. non-empty encryptionPublicKey)
+    if (
+      req.user.encryptionPublicKey &&
+      req.user.encryptionPublicKey.trim() !== '' &&
+      response.text &&
+      crypto
+    ) {
+      try {
+        // Reconstruct the user's RSA public key in PEM format.
+        const pubKeyBase64 = req.user.encryptionPublicKey;
+        const pemPublicKey = `-----BEGIN PUBLIC KEY-----\n${pubKeyBase64.match(/.{1,64}/g).join('\n')}\n-----END PUBLIC KEY-----`;
+
+        // Generate a random 256-bit AES key and a 12-byte IV.
+        const aesKey = crypto.randomBytes(32);
+        const iv = crypto.randomBytes(12);
+
+        // Encrypt the response text using AES-GCM.
+        const cipher = crypto.createCipheriv('aes-256-gcm', aesKey, iv);
+        let ciphertext = cipher.update(response.text, 'utf8', 'base64');
+        ciphertext += cipher.final('base64');
+        const authTag = cipher.getAuthTag().toString('base64');
+
+        // Encrypt the AES key using the client's RSA public key.
+        let encryptedKey;
+        try {
+          encryptedKey = crypto.publicEncrypt(
+            {
+              key: pemPublicKey,
+              padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
+              oaepHash: 'sha256',
+            },
+            aesKey,
+          ).toString('base64');
+        } catch (err) {
+          logger.error('Error encrypting AES key:', err);
+          throw new Error('Encryption failure');
+        }
+
+        // Replace the plaintext response with the encrypted payload.
+        response.text = ciphertext;
+        response.iv = iv.toString('base64');
+        response.authTag = authTag;
+        response.encryptedKey = encryptedKey;
+        logger.debug('[AskController] Response message encrypted.');
+      } catch (encError) {
+        logger.error('[AskController] Error during response encryption:', encError);
+        // Optionally, you may choose to return plaintext if encryption fails.
+      }
+    }
+    // --- End Encryption Branch ---
+
     if (!abortController.signal.aborted) {
       sendMessage(res, {
         final: true,

api/server/controllers/UserController.js

@@ -167,17 +167,20 @@ const updateUserEncryptionController = async (req, res) => {
   try {
     const { encryptionPublicKey, encryptedPrivateKey, encryptionSalt, encryptionIV } = req.body;
 
-    // Validate required parameters
-    if (!encryptionPublicKey || !encryptedPrivateKey || !encryptionSalt || !encryptionIV) {
+    // Allow disabling encryption by passing null for all fields.
+    const allNull = encryptionPublicKey === null && encryptedPrivateKey === null && encryptionSalt === null && encryptionIV === null;
+    const allPresent = encryptionPublicKey && encryptedPrivateKey && encryptionSalt && encryptionIV;
+
+    if (!allNull && !allPresent) {
       return res.status(400).json({ message: 'Missing encryption parameters.' });
     }
 
-    // Use the helper function to update the user.
+    // Update the user record with the provided encryption parameters (or null to disable)
     const updatedUser = await updateUser(req.user.id, {
-      encryptionPublicKey,
-      encryptedPrivateKey,
-      encryptionSalt,
-      encryptionIV,
+      encryptionPublicKey: encryptionPublicKey || null,
+      encryptedPrivateKey: encryptedPrivateKey || null,
+      encryptionSalt: encryptionSalt || null,
+      encryptionIV: encryptionIV || null,
     });
 
     if (!updatedUser) {