refactor: works. now fixing to decrypt the text in the UI.

api/app/clients/BaseClient.js

@@ -1,4 +1,3 @@
-const crypto = require('crypto');
 const fetch = require('node-fetch');
 const {
   supportsBalanceCheck,
@@ -8,7 +7,7 @@ const {
   ErrorTypes,
   Constants,
 } = require('librechat-data-provider');
-const { getMessages, saveMessage, updateMessage, saveConvo } = require('~/models');
+const { getMessages, saveMessage, updateMessage, saveConvo, getUserById } = require('~/models');
 const { addSpaceIfNeeded, isEnabled } = require('~/server/utils');
 const { truncateToolCallOutputs } = require('./prompts');
 const checkBalance = require('~/models/checkBalance');
@@ -16,6 +15,48 @@ const { getFiles } = require('~/models/File');
 const TextStream = require('./TextStream');
 const { logger } = require('~/config');
 
+let crypto;
+try {
+  crypto = require('crypto');
+} catch (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,
+  };
+}
+
 class BaseClient {
   constructor(apiKey, options = {}) {
     this.apiKey = apiKey;
@@ -844,18 +885,64 @@ class BaseClient {
    * @param {string | null} user
    */
   async saveMessageToDatabase(message, endpointOptions, user = null) {
-    if (this.user && user !== this.user) {
+    // Normalize the user information:
+    // If "user" is an object, use it; otherwise, if a string is passed use req.user (if available)
+    const currentUser =
+      user && typeof user === 'object'
+        ? user
+        : (this.options.req && this.options.req.user
+          ? this.options.req.user
+          : { id: user });
+    const currentUserId = currentUser.id || currentUser;
+
+    // Check if the client’s stored user matches the current user.
+    // (this.user might have been set earlier in setMessageOptions)
+    const storedUserId =
+      this.user && typeof this.user === 'object' ? this.user.id : this.user;
+    if (storedUserId && currentUserId && storedUserId !== currentUserId) {
       throw new Error('User mismatch.');
     }
 
+    // console.log('User ID:', currentUserId);
+
+    const dbUser = await getUserById(currentUserId, 'encryptionPublicKey');
+
+    // --- NEW ENCRYPTION BLOCK: Encrypt AI response if encryptionPublicKey exists ---
+    if (dbUser.encryptionPublicKey && message && message.text) {
+      try {
+        // Rebuild the PEM format if necessary.
+        const pemPublicKey = `-----BEGIN PUBLIC KEY-----\n${dbUser.encryptionPublicKey
+          .match(/.{1,64}/g)
+          .join('\n')}\n-----END PUBLIC KEY-----`;
+        const { ciphertext, iv, authTag, encryptedKey } = encryptText(
+          message.text,
+          pemPublicKey,
+        );
+        message.text = ciphertext;
+        message.iv = iv;
+        message.authTag = authTag;
+        message.encryptedKey = encryptedKey;
+        logger.debug('[BaseClient.saveMessageToDatabase] Encrypted message text');
+      } catch (err) {
+        logger.error('[BaseClient.saveMessageToDatabase] Error encrypting message text', err);
+      }
+    }
+    // --- End Encryption Block ---
+
+    // Build update parameters including encryption fields.
+    const updateParams = {
+      ...message,
+      endpoint: this.options.endpoint,
+      unfinished: false,
+      user: currentUserId, // store the user id (ensured to be a string)
+      iv: message.iv ?? null,
+      authTag: message.authTag ?? null,
+      encryptedKey: message.encryptedKey ?? null,
+    };
+
     const savedMessage = await saveMessage(
       this.options.req,
-      {
+      updateParams,
-        ...message,
-        endpoint: this.options.endpoint,
-        unfinished: false,
-        user,
-      },
       { context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveMessage' },
     );
 
@@ -1121,4 +1208,4 @@ class BaseClient {
   }
 }
 
-module.exports = BaseClient;
+module.exports = BaseClient;

api/server/controllers/AskController.js

@@ -87,10 +87,15 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
   // 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.
+  // Build clientOptions including the encryptionPublicKey (if available)
+  const clientOptions = {
+    encryptionPublicKey: dbUser?.encryptionPublicKey,
+  };
+
+  // Rebuild PEM format if encryptionPublicKey is available
   let pemPublicKey = null;
-  if (dbUser?.encryptionPublicKey && dbUser.encryptionPublicKey.trim() !== '') {
+  if (clientOptions.encryptionPublicKey && clientOptions.encryptionPublicKey.trim() !== '') {
-    const pubKeyBase64 = dbUser.encryptionPublicKey;
+    const pubKeyBase64 = clientOptions.encryptionPublicKey;
     pemPublicKey = `-----BEGIN PUBLIC KEY-----\n${pubKeyBase64.match(/.{1,64}/g).join('\n')}\n-----END PUBLIC KEY-----`;
   }
 
@@ -113,7 +118,8 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
 
   let getText;
   try {
-    const { client } = await initializeClient({ req, res, endpointOption });
+    // Pass clientOptions (which includes encryptionPublicKey) along with other parameters to initializeClient
+    const { client } = await initializeClient({ req, res, endpointOption, ...clientOptions });
     const { onProgress: progressCallback, getPartialText } = createOnProgress();
     getText = client.getStreamText != null ? client.getStreamText.bind(client) : getPartialText;
 
@@ -176,7 +182,6 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
         logger.debug('[AskController] User message encrypted.');
       } catch (encError) {
         logger.error('[AskController] Error encrypting user message:', encError);
-        // Optionally, you could choose to throw an error or fallback.
       }
     }
 
@@ -191,7 +196,6 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
         logger.debug('[AskController] Response message encrypted.');
       } catch (encError) {
         logger.error('[AskController] Error encrypting response message:', encError);
-        // Optionally, you can choose to send plaintext or handle the error.
       }
     }
     // --- End Encryption Branch ---