refactor: request is encrypted. response from AI is still saved in plaintext but from the stream the final response is encrypted.

api/server/controllers/AskController.js

@@ -1,16 +1,59 @@
 const { getResponseSender, Constants } = require('librechat-data-provider');
 const { createAbortController, handleAbortError } = require('~/server/middleware');
 const { sendMessage, createOnProgress } = require('~/server/utils');
-const { saveMessage } = require('~/models');
+const { saveMessage, getUserById } = require('~/models');
 const { logger } = require('~/config');
 
 let crypto;
 try {
   crypto = require('crypto');
 } catch (err) {
-  logger.error('[openidStrategy] crypto support is disabled!', err);
+  logger.error('[AskController] crypto support is disabled!', err);
 }
 
+/**
+ * Helper function to encrypt plaintext using AES-256-GCM and then RSA-encrypt the AES key.
+ * @param {string} plainText - The plaintext to encrypt.
+ * @param {string} pemPublicKey - The RSA public key in PEM format.
+ * @returns {Object} An object containing the ciphertext, iv, authTag, and encryptedKey.
+ */
+function encryptText(plainText, pemPublicKey) {
+  // Generate a random 256-bit AES key and a 12-byte IV.
+  const aesKey = crypto.randomBytes(32);
+  const iv = crypto.randomBytes(12);
+
+  // Encrypt the plaintext using AES-256-GCM.
+  const cipher = crypto.createCipheriv('aes-256-gcm', aesKey, iv);
+  let ciphertext = cipher.update(plainText, 'utf8', 'base64');
+  ciphertext += cipher.final('base64');
+  const authTag = cipher.getAuthTag().toString('base64');
+
+  // Encrypt the AES key using the user's RSA public key.
+  const encryptedKey = crypto.publicEncrypt(
+    {
+      key: pemPublicKey,
+      padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
+      oaepHash: 'sha256',
+    },
+    aesKey,
+  ).toString('base64');
+
+  return {
+    ciphertext,
+    iv: iv.toString('base64'),
+    authTag,
+    encryptedKey,
+  };
+}
+
+/**
+ * AskController
+ * - Initializes the client.
+ * - Obtains the response from the language model.
+ * - Retrieves the full user record (to get encryption parameters).
+ * - If the user has encryption enabled (i.e. encryptionPublicKey is provided),
+ *   encrypts both the request (userMessage) and the response before saving.
+ */
 const AskController = async (req, res, next, initializeClient, addTitle) => {
   let {
     text,
@@ -39,7 +82,17 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
     modelDisplayLabel,
   });
   const newConvo = !conversationId;
-  const user = req.user.id;
+  const userId = req.user.id; // User ID from authentication
+
+  // Retrieve full user record from DB (including encryption parameters)
+  const dbUser = await getUserById(userId, 'encryptionPublicKey encryptedPrivateKey encryptionSalt encryptionIV');
+
+  // If the user has provided an encryption public key, rebuild the PEM format.
+  let pemPublicKey = null;
+  if (dbUser?.encryptionPublicKey && dbUser.encryptionPublicKey.trim() !== '') {
+    const pubKeyBase64 = dbUser.encryptionPublicKey;
+    pemPublicKey = `-----BEGIN PUBLIC KEY-----\n${pubKeyBase64.match(/.{1,64}/g).join('\n')}\n-----END PUBLIC KEY-----`;
+  }
 
   const getReqData = (data = {}) => {
     for (let key in data) {
@@ -59,11 +112,9 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
   };
 
   let getText;
-
   try {
     const { client } = await initializeClient({ req, res, endpointOption });
     const { onProgress: progressCallback, getPartialText } = createOnProgress();
-
     getText = client.getStreamText != null ? client.getStreamText.bind(client) : getPartialText;
 
     const getAbortData = () => ({
@@ -88,7 +139,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
     });
 
     const messageOptions = {
-      user,
+      user: userId,
       parentMessageId,
       conversationId,
       overrideParentMessageId,
@@ -99,10 +150,11 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       progressOptions: { res },
     };
 
-    /** @type {TMessage} */
+    // Get the response from the language model client.
     let response = await client.sendMessage(text, messageOptions);
     response.endpoint = endpointOption.endpoint;
 
+    // Ensure the conversation has a title.
     const { conversation = {} } = await client.responsePromise;
     conversation.title =
       conversation && !conversation.title ? null : conversation?.title || 'New Chat';
@@ -113,54 +165,33 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       delete userMessage.image_urls;
     }
 
-    // --- Encryption Branch ---
+    // --- Encrypt the user message if encryption is enabled ---
-    // Only encrypt if the user has set up encryption (i.e. non-empty encryptionPublicKey)
+    if (pemPublicKey && userMessage && userMessage.text) {
-    if (
-      req.user.encryptionPublicKey &&
-      req.user.encryptionPublicKey.trim() !== '' &&
-      response.text &&
-      crypto
-    ) {
       try {
-        // Reconstruct the user's RSA public key in PEM format.
+        const { ciphertext, iv, authTag, encryptedKey } = encryptText(userMessage.text, pemPublicKey);
-        const pubKeyBase64 = req.user.encryptionPublicKey;
+        userMessage.text = ciphertext;
-        const pemPublicKey = `-----BEGIN PUBLIC KEY-----\n${pubKeyBase64.match(/.{1,64}/g).join('\n')}\n-----END PUBLIC KEY-----`;
+        userMessage.iv = iv;
-
+        userMessage.authTag = authTag;
-        // Generate a random 256-bit AES key and a 12-byte IV.
+        userMessage.encryptedKey = encryptedKey;
-        const aesKey = crypto.randomBytes(32);
+        logger.debug('[AskController] User message encrypted.');
-        const iv = crypto.randomBytes(12);
+      } catch (encError) {
-
+        logger.error('[AskController] Error encrypting user message:', encError);
-        // Encrypt the response text using AES-GCM.
+        // Optionally, you could choose to throw an error or fallback.
-        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.
+    // --- Encrypt the AI response if encryption is enabled ---
+    if (pemPublicKey && response.text) {
+      try {
+        const { ciphertext, iv, authTag, encryptedKey } = encryptText(response.text, pemPublicKey);
         response.text = ciphertext;
-        response.iv = iv.toString('base64');
+        response.iv = iv;
         response.authTag = authTag;
         response.encryptedKey = encryptedKey;
         logger.debug('[AskController] Response message encrypted.');
       } catch (encError) {
-        logger.error('[AskController] Error during response encryption:', encError);
+        logger.error('[AskController] Error encrypting response message:', encError);
-        // Optionally, you may choose to return plaintext if encryption fails.
+        // Optionally, you can choose to send plaintext or handle the error.
       }
     }
     // --- End Encryption Branch ---
@@ -178,15 +209,15 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       if (!client.savedMessageIds.has(response.messageId)) {
         await saveMessage(
           req,
-          { ...response, user },
+          { ...response, user: userId },
-          { context: 'api/server/controllers/AskController.js - response end' },
+          { context: 'AskController - response end' },
         );
       }
     }
 
     if (!client.skipSaveUserMessage) {
       await saveMessage(req, userMessage, {
-        context: 'api/server/controllers/AskController.js - don\'t skip saving user message',
+        context: 'AskController - save user message',
       });
     }
 
@@ -206,7 +237,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       messageId: responseMessageId,
       parentMessageId: overrideParentMessageId ?? userMessageId ?? parentMessageId,
     }).catch((err) => {
-      logger.error('[AskController] Error in `handleAbortError`', err);
+      logger.error('[AskController] Error in handleAbortError', err);
     });
   }
 };