From d37cc1cf4d0197a3761149121553810f465ec6d2 Mon Sep 17 00:00:00 2001 From: Ruben Talstra Date: Sun, 16 Feb 2025 16:58:59 +0100 Subject: [PATCH] refactor: works. now fixing to decrypt the text in the UI. --- api/app/clients/BaseClient.js | 107 +++++++++++++++++++++--- api/server/controllers/AskController.js | 16 ++-- 2 files changed, 107 insertions(+), 16 deletions(-) diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index ebf3ca12d9..880559cfca 100644 --- a/api/app/clients/BaseClient.js +++ b/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, - { - ...message, - endpoint: this.options.endpoint, - unfinished: false, - user, - }, + updateParams, { context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveMessage' }, ); @@ -1121,4 +1208,4 @@ class BaseClient { } } -module.exports = BaseClient; +module.exports = BaseClient; \ No newline at end of file diff --git a/api/server/controllers/AskController.js b/api/server/controllers/AskController.js index 83177555d6..d0cde5f5d2 100644 --- a/api/server/controllers/AskController.js +++ b/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() !== '') { - const pubKeyBase64 = dbUser.encryptionPublicKey; + if (clientOptions.encryptionPublicKey && clientOptions.encryptionPublicKey.trim() !== '') { + 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 ---