Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP E2EE encryption [ignore just a thought idea] #5856

Closed
wants to merge 2 commits into from
Closed

WIP E2EE encryption [ignore just a thought idea] #5856

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9989f97
started. but not the right implementation.
rubentalstra Feb 13, 2025
9a7db62
Merge branch 'main' into feat/end-to-end-encryption-conversations
rubentalstra Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions api/app/clients/BaseClient.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -793,7 +793,12 @@ class BaseClient {
async loadHistory(conversationId, parentMessageId = null) {
logger.debug('[BaseClient] Loading history:', { conversationId, parentMessageId });

const messages = (await getMessages({ conversationId })) ?? [];
const messages =
(await getMessages({
conversationId,
encryptionKey: this.options.req?.headers['x-encryption-key'],
isEncrypted: this.options.req?.headers['x-encryption-enabled'] === 'true',
})) ?? [];

if (messages.length === 0) {
return [];
Expand Down Expand Up @@ -856,7 +861,11 @@ class BaseClient {
unfinished: false,
user,
},
{ context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveMessage' },
{
context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveMessage',
encryptionKey: this.options.req?.headers['x-encryption-key'],
isEncrypted: this.options.req?.headers['x-encryption-enabled'] === 'true',
},
);

if (this.skipSaveConvo) {
Expand All @@ -882,7 +891,12 @@ class BaseClient {
* @param {Partial<TMessage>} message
*/
async updateMessageInDatabase(message) {
await updateMessage(this.options.req, message);
await updateMessage(this.options.req, {
...message,
// Pass encryption headers to the update operation
encryptionKey: this.options.req?.headers['x-encryption-key'],
isEncrypted: this.options.req?.headers['x-encryption-enabled'] === 'true',
});
}

/**
Expand Down Expand Up @@ -1121,4 +1135,4 @@ class BaseClient {
}
}

module.exports = BaseClient;
module.exports = BaseClient;
84 changes: 65 additions & 19 deletions api/models/Message.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,13 @@ const { logger } = require('~/config');

const idSchema = z.string().uuid();

function getEncryptionOptions(req) {
if (req?.headers?.['x-encryption-enabled'] === 'true' && req?.headers?.['x-encryption-key']) {
return { encryptionKey: req.headers['x-encryption-key'] };
}
return null;
}

/**
* Saves a message in the database.
*
Expand Down Expand Up @@ -61,12 +68,18 @@ async function saveMessage(req, params, metadata) {
update.expiredAt = null;
}

const message = await Message.findOneAndUpdate(
const query = Message.findOneAndUpdate(
{ messageId: params.messageId, user: req.user.id },
update,
{ upsert: true, new: true },
);

const encryptionOptions = getEncryptionOptions(req);
if (encryptionOptions) {
query.setOptions(encryptionOptions);
}

const message = await query;
return message.toObject();
} catch (err) {
logger.error('Error saving message:', err);
Expand Down Expand Up @@ -128,7 +141,6 @@ async function recordMessage({
...rest
}) {
try {
// No parsing of convoId as may use threadId
const message = {
user,
endpoint,
Expand All @@ -138,10 +150,16 @@ async function recordMessage({
...rest,
};

return await Message.findOneAndUpdate({ user, messageId }, message, {
const query = Message.findOneAndUpdate({ user, messageId }, message, {
upsert: true,
new: true,
});

if (rest.encryptionKey && rest.isEncrypted) {
query.setOptions({ encryptionKey: rest.encryptionKey });
}

return await query;
} catch (err) {
logger.error('Error recording message:', err);
throw err;
Expand All @@ -162,7 +180,14 @@ async function recordMessage({
*/
async function updateMessageText(req, { messageId, text }) {
try {
await Message.updateOne({ messageId, user: req.user.id }, { text });
const query = Message.updateOne({ messageId, user: req.user.id }, { text });

const encryptionOptions = getEncryptionOptions(req);
if (encryptionOptions) {
query.setOptions(encryptionOptions);
}

await query;
} catch (err) {
logger.error('Error updating message text:', err);
throw err;
Expand Down Expand Up @@ -190,13 +215,14 @@ async function updateMessageText(req, { messageId, text }) {
async function updateMessage(req, message, metadata) {
try {
const { messageId, ...update } = message;
const updatedMessage = await Message.findOneAndUpdate(
{ messageId, user: req.user.id },
update,
{
new: true,
},
);
const query = Message.findOneAndUpdate({ messageId, user: req.user.id }, update, { new: true });

const encryptionOptions = getEncryptionOptions(req);
if (encryptionOptions) {
query.setOptions(encryptionOptions);
}

const updatedMessage = await query;

if (!updatedMessage) {
throw new Error('Message not found or user not authorized.');
Expand Down Expand Up @@ -234,11 +260,20 @@ async function updateMessage(req, message, metadata) {
*/
async function deleteMessagesSince(req, { messageId, conversationId }) {
try {
const message = await Message.findOne({ messageId, user: req.user.id }).lean();
const query = Message.findOne({ messageId, user: req.user.id });
const encryptionOptions = getEncryptionOptions(req);
if (encryptionOptions) {
query.setOptions(encryptionOptions);
}

const message = await query.lean();

if (message) {
const query = Message.find({ conversationId, user: req.user.id });
return await query.deleteMany({
const deleteQuery = Message.find({ conversationId, user: req.user.id });
if (encryptionOptions) {
deleteQuery.setOptions(encryptionOptions);
}
return await deleteQuery.deleteMany({
createdAt: { $gt: message.createdAt },
});
}
Expand All @@ -260,11 +295,20 @@ async function deleteMessagesSince(req, { messageId, conversationId }) {
*/
async function getMessages(filter, select) {
try {
const query = Message.find(filter);

if (filter.encryptionKey) {
query.setOptions({
encryptionKey: filter.encryptionKey,
});
delete filter.encryptionKey; // Remove from filter after setting options
}

if (select) {
return await Message.find(filter).select(select).sort({ createdAt: 1 }).lean();
query.select(select);
}

return await Message.find(filter).sort({ createdAt: 1 }).lean();
return await query.sort({ createdAt: 1 }).lean();
} catch (err) {
logger.error('Error getting messages:', err);
throw err;
Expand All @@ -281,10 +325,12 @@ async function getMessages(filter, select) {
*/
async function getMessage({ user, messageId }) {
try {
return await Message.findOne({
const query = Message.findOne({
user,
messageId,
}).lean();
});

return await query.lean();
} catch (err) {
logger.error('Error getting message:', err);
throw err;
Expand Down Expand Up @@ -320,4 +366,4 @@ module.exports = {
getMessages,
getMessage,
deleteMessages,
};
};
67 changes: 59 additions & 8 deletions api/models/Share.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ const { nanoid } = require('nanoid');
const { Constants } = require('librechat-data-provider');
const { Conversation } = require('~/models/Conversation');
const SharedLink = require('./schema/shareSchema');
const { getMessages } = require('./Message');
const logger = require('~/config/winston');
const { getMessages } = require('./Message');
const { decrypt } = require('~/server/utils/encryptionUtil');

class ShareServiceError extends Error {
constructor(message, code) {
Expand Down Expand Up @@ -63,7 +64,15 @@ function anonymizeMessages(messages, newConvoId) {
});
}

async function getSharedMessages(shareId) {
const isEncrypted = (text) => {
if (!text || typeof text !== 'string') {
return false;
}
const parts = text.split(':');
return parts.length === 2 && parts[0].length === 32;
};

async function getSharedMessages(shareId, encryptionKey) {
try {
const share = await SharedLink.findOne({ shareId, isPublic: true })
.populate({
Expand All @@ -77,6 +86,27 @@ async function getSharedMessages(shareId) {
return null;
}

// Decrypt messages if encryption key provided
if (encryptionKey) {
share.messages = share.messages.map((message) => {
if (message.text && isEncrypted(message.text)) {
message.text = decrypt(message.text, encryptionKey);
}
if (message.content) {
message.content = message.content.map((item) => {
if (item.text && isEncrypted(item.text)) {
return {
...item,
text: decrypt(item.text, encryptionKey),
};
}
return item;
});
}
return message;
});
}

const newConvoId = anonymizeConvoId(share.conversationId);
const result = {
...share,
Expand All @@ -86,7 +116,7 @@ async function getSharedMessages(shareId) {

return result;
} catch (error) {
logger.error('[getShare] Error getting share link', {
logger.error('[getShare] Error getting share link:', {
error: error.message,
shareId,
});
Expand Down Expand Up @@ -187,13 +217,13 @@ async function deleteAllSharedLinks(user) {
}
}

async function createSharedLink(user, conversationId) {
async function createSharedLink(user, conversationId, encryptionKey) {
if (!user || !conversationId) {
throw new ShareServiceError('Missing required parameters', 'INVALID_PARAMS');
}

try {
const [existingShare, conversationMessages] = await Promise.all([
const [existingShare, messages] = await Promise.all([
SharedLink.findOne({ conversationId, isPublic: true }).select('-_id -__v -user').lean(),
getMessages({ conversationId }),
]);
Expand All @@ -207,11 +237,31 @@ async function createSharedLink(user, conversationId) {
const conversation = await Conversation.findOne({ conversationId }).lean();
const title = conversation?.title || 'Untitled';

// Ensure messages are decrypted before sharing
let sharableMessages = messages;
if (encryptionKey) {
sharableMessages = messages.map((message) => {
const decryptedMessage = { ...message };
if (message.text && isEncrypted(message.text)) {
decryptedMessage.text = decrypt(message.text, encryptionKey);
}
if (message.content) {
decryptedMessage.content = message.content.map((item) => {
if (item.text && isEncrypted(item.text)) {
return { ...item, text: decrypt(item.text, encryptionKey) };
}
return item;
});
}
return decryptedMessage;
});
}

const shareId = nanoid();
await SharedLink.create({
shareId,
conversationId,
messages: conversationMessages,
messages: sharableMessages,
title,
user,
});
Expand Down Expand Up @@ -243,7 +293,7 @@ async function getSharedLink(user, conversationId) {

return { shareId: share.shareId, success: true };
} catch (error) {
logger.error('[getSharedLink] Error getting shared link', {
logger.error('[getSharedLink] Error getting shared link:', {
error: error.message,
user,
conversationId,
Expand Down Expand Up @@ -289,7 +339,7 @@ async function updateSharedLink(user, shareId) {

return { shareId: newShareId, conversationId: updatedShare.conversationId };
} catch (error) {
logger.error('[updateSharedLink] Error updating shared link', {
logger.error('[updateSharedLink] Error updating shared link:', {
error: error.message,
user,
shareId,
Expand Down Expand Up @@ -337,4 +387,5 @@ module.exports = {
deleteSharedLink,
getSharedMessages,
deleteAllSharedLinks,
isEncrypted,
};
Loading
Loading