refactor: creating a starting point for E2EE

This commit is contained in:
Ruben Talstra 2025-02-15 23:04:26 +01:00
parent 18d019d8b3
commit 606fea044a
No known key found for this signature in database
GPG key ID: 2A5A7174A60F3BEA
5 changed files with 301 additions and 188 deletions

View file

@ -15,24 +15,17 @@ import type { TUser } from 'librechat-data-provider';
import { useToastContext } from '~/Providers';
import { useSetUserEncryptionMutation } from '~/data-provider';
// Helper: Convert a Base64 string to Uint8Array.
const base64ToUint8Array = (base64: string): Uint8Array => {
const binaryString = window.atob(base64);
const len = binaryString.length;
const bytes = new Uint8Array(len);
for (let i = 0; i < len; i++) {
bytes[i] = binaryString.charCodeAt(i);
}
return bytes;
};
// Helper: Convert a Uint8Array to a hex string (for debugging).
/**
* Helper: Convert a Uint8Array to a hex string (for debugging).
*/
const uint8ArrayToHex = (array: Uint8Array): string =>
Array.from(array)
.map(b => b.toString(16).padStart(2, '0'))
.map((b) => b.toString(16).padStart(2, '0'))
.join('');
// Derive an AES-GCM key from the passphrase using PBKDF2.
/**
* Derive an AES-GCM key from the passphrase using PBKDF2.
*/
const deriveKey = async (passphrase: string, salt: Uint8Array): Promise<CryptoKey> => {
const encoder = new TextEncoder();
const keyMaterial = await window.crypto.subtle.importKey(
@ -46,7 +39,7 @@ const deriveKey = async (passphrase: string, salt: Uint8Array): Promise<CryptoKe
{
name: 'PBKDF2',
salt,
iterations: 100000, // Adjust as needed for security/performance trade-off
iterations: 100000,
hash: 'SHA-256',
},
keyMaterial,
@ -60,22 +53,6 @@ const deriveKey = async (passphrase: string, salt: Uint8Array): Promise<CryptoKe
return derivedKey;
};
// Decrypt the encrypted private key using the derived key and IV.
const decryptPrivateKey = async (
derivedKey: CryptoKey,
iv: Uint8Array,
ciphertextBase64: string
): Promise<string> => {
const ciphertext = base64ToUint8Array(ciphertextBase64);
const decryptedBuffer = await window.crypto.subtle.decrypt(
{ name: 'AES-GCM', iv },
derivedKey,
ciphertext
);
const decoder = new TextDecoder();
return decoder.decode(decryptedBuffer);
};
const UserKeysSettings: FC = () => {
const localize = useLocalize();
const { user } = useAuthContext();
@ -122,22 +99,23 @@ const UserKeysSettings: FC = () => {
true,
['encrypt', 'decrypt']
);
// Export public and private keys.
const publicKeyBuffer = await window.crypto.subtle.exportKey('spki', keyPair.publicKey);
const privateKeyBuffer = await window.crypto.subtle.exportKey('pkcs8', keyPair.privateKey);
const publicKeyBase64 = btoa(String.fromCharCode(...new Uint8Array(publicKeyBuffer)));
const privateKeyBase64 = btoa(String.fromCharCode(...new Uint8Array(privateKeyBuffer)));
const publicKeyBase64 = window.btoa(String.fromCharCode(...new Uint8Array(publicKeyBuffer)));
const privateKeyBase64 = window.btoa(String.fromCharCode(...new Uint8Array(privateKeyBuffer)));
console.debug('New public key:', publicKeyBase64);
console.debug('New private key (plaintext):', privateKeyBase64);
// Generate a salt and IV.
// Generate a salt and IV for AES-GCM.
const salt = window.crypto.getRandomValues(new Uint8Array(16)); // 16 bytes salt
const iv = window.crypto.getRandomValues(new Uint8Array(12)); // 12 bytes IV
// Derive a symmetric key from the passphrase.
const derivedKey = await deriveKey(passphrase, salt);
// Encrypt the private key.
// Encrypt the private key using AES-GCM.
const encoder = new TextEncoder();
const privateKeyBytes = encoder.encode(privateKeyBase64);
const encryptedPrivateKeyBuffer = await window.crypto.subtle.encrypt(
@ -145,11 +123,13 @@ const UserKeysSettings: FC = () => {
derivedKey,
privateKeyBytes
);
const encryptedPrivateKeyBase64 = btoa(String.fromCharCode(...new Uint8Array(encryptedPrivateKeyBuffer)));
const encryptedPrivateKeyBase64 = window.btoa(
String.fromCharCode(...new Uint8Array(encryptedPrivateKeyBuffer))
);
// Convert salt and IV to Base64.
const saltBase64 = btoa(String.fromCharCode(...salt));
const ivBase64 = btoa(String.fromCharCode(...iv));
// Convert salt and IV to Base64 strings.
const saltBase64 = window.btoa(String.fromCharCode(...salt));
const ivBase64 = window.btoa(String.fromCharCode(...iv));
console.debug('Activation complete:');
console.debug('Encrypted private key:', encryptedPrivateKeyBase64);
@ -168,7 +148,7 @@ const UserKeysSettings: FC = () => {
};
const handleSubmit = async (): Promise<void> => {
// Activate encryption (or re-activate) by generating new keys.
// Activate encryption by generating new keys.
const newEncryption = await activateEncryption();
if (newEncryption) {
try {
@ -176,6 +156,7 @@ const UserKeysSettings: FC = () => {
await setEncryption(newEncryption);
showToast({ message: localize('com_ui_upload_success') });
// Update local user state with the new encryption keys.
// Later, when the user unlocks their keys, store the decrypted private key.
setUser((prev) => ({
...prev,
...newEncryption,
@ -198,9 +179,7 @@ const UserKeysSettings: FC = () => {
<div className="flex items-center justify-between">
<div className="flex items-center space-x-2">
<Key className="flex w-[20px] h-[20px]" />
<span id="user-keys-label">
{localize('com_nav_chat_encryption_settings')}
</span>
<span id="user-keys-label">{localize('com_nav_chat_encryption_settings')}</span>
</div>
<Button
variant="outline"
@ -224,9 +203,7 @@ const UserKeysSettings: FC = () => {
<OGDialog open={dialogOpen} onOpenChange={setDialogOpen}>
<OGDialogContent className="w-11/12 max-w-sm" style={{ borderRadius: '12px' }}>
<OGDialogHeader>
<OGDialogTitle>
{localize('com_nav_chat_enter_your_passphrase')}
</OGDialogTitle>
<OGDialogTitle>{localize('com_nav_chat_enter_your_passphrase')}</OGDialogTitle>
</OGDialogHeader>
<div className="p-4 flex flex-col gap-4">
<Input

View file

@ -0,0 +1,131 @@
import {
createPayload,
isAgentsEndpoint,
isAssistantsEndpoint,
removeNullishValues,
TPayload,
TSubmission,
} from 'librechat-data-provider';
/**
* Convert an ArrayBuffer to a Base64 string.
*/
function arrayBufferToBase64(buffer: ArrayBuffer): string {
const bytes = new Uint8Array(buffer);
let binary = '';
bytes.forEach((b) => (binary += String.fromCharCode(b)));
return window.btoa(binary);
}
/**
* Convert a Base64 string to a Uint8Array.
*/
function base64ToUint8Array(base64: string): Uint8Array {
const binaryString = window.atob(base64);
const len = binaryString.length;
const bytes = new Uint8Array(len);
for (let i = 0; i < len; i++) {
bytes[i] = binaryString.charCodeAt(i);
}
return bytes;
}
/**
* Encrypts a plaintext string using RSA-OAEP.
* The public key must be provided as a Base64-encoded SPKI key.
*/
export async function encryptMessage(
plainText: string,
userEncryptionPublicKey: string
): Promise<string> {
// Convert the Base64 public key into a binary array.
const binaryKey = base64ToUint8Array(userEncryptionPublicKey);
const publicKey = await window.crypto.subtle.importKey(
'spki',
binaryKey.buffer,
{ name: 'RSA-OAEP', hash: 'SHA-256' },
true,
['encrypt']
);
const encoder = new TextEncoder();
const encodedContent = encoder.encode(plainText);
const encryptedBuffer = await window.crypto.subtle.encrypt(
{ name: 'RSA-OAEP' },
publicKey,
encodedContent
);
// Return the encrypted data as a Base64 string.
return arrayBufferToBase64(encryptedBuffer);
}
/**
* Decrypts an encrypted string using RSA-OAEP.
* The private key must be provided as a Base64-encoded PKCS#8 key.
*/
export async function decryptMessage(
encryptedText: string,
userPrivateKey: string
): Promise<string> {
// Convert the Base64-encoded private key to a binary array.
const binaryKey = base64ToUint8Array(userPrivateKey);
const privateKey = await window.crypto.subtle.importKey(
'pkcs8',
binaryKey.buffer,
{ name: 'RSA-OAEP', hash: 'SHA-256' },
true,
['decrypt']
);
// Convert the Base64-encoded encrypted text to a binary array.
const encryptedBuffer = base64ToUint8Array(encryptedText);
const decryptedBuffer = await window.crypto.subtle.decrypt(
{ name: 'RSA-OAEP' },
privateKey,
encryptedBuffer.buffer
);
const decoder = new TextDecoder();
return decoder.decode(decryptedBuffer);
}
/**
* Creates a payload from the submission.
* If encryption is enabled (or if a userEncryptionPublicKey is available),
* it encrypts the text and attaches a random IV (used as a flag).
*/
export async function createPayloadWithEncryption(
submission: TSubmission,
encryptionEnabled: boolean,
userEncryptionPublicKey?: string
): Promise<{ server: string; payload: TPayload }> {
// Create the standard payload.
const payloadData = createPayload(submission);
let { payload } = payloadData;
// Remove nullish values for endpoints that require it.
if (isAssistantsEndpoint(payload.endpoint) || isAgentsEndpoint(payload.endpoint)) {
payload = removeNullishValues(payload) as TPayload;
}
// Force encryption if encryptionEnabled is true OR a public key is available.
if ((encryptionEnabled || userEncryptionPublicKey) && userEncryptionPublicKey) {
const plainText = payload.text;
if (plainText !== undefined) {
// Generate a random IV (12 bytes for AES-GCM). Although RSA-OAEP doesnt need an IV,
// we attach it as a marker that the message is encrypted.
const ivArray = new Uint8Array(12);
window.crypto.getRandomValues(ivArray);
const ivString = arrayBufferToBase64(ivArray.buffer);
// Encrypt the message text.
const encryptedText = await encryptMessage(plainText, userEncryptionPublicKey);
console.log('Encryption successful:', { encryptedText, iv: ivString });
payload.text = encryptedText;
// Attach the IV as a marker.
(payload as TPayload & { messageEncryptionIV?: string }).messageEncryptionIV = ivString;
}
}
return { server: payloadData.server, payload };
}

View file

@ -4,7 +4,6 @@ import { SSE } from 'sse.js';
import { useSetRecoilState } from 'recoil';
import {
request,
/* @ts-ignore */
createPayload,
isAgentsEndpoint,
removeNullishValues,
@ -17,6 +16,7 @@ import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/d
import { useAuthContext } from '~/hooks/AuthContext';
import useEventHandlers from './useEventHandlers';
import store from '~/store';
import { createPayloadWithEncryption } from '~/hooks/SSE/encryptionHelpers';
type ChatHelpers = Pick<
EventHandlerParams,
@ -33,11 +33,12 @@ export default function useSSE(
chatHelpers: ChatHelpers,
isAddedRequest = false,
runIndex = 0,
encryptionEnabled: boolean = false // Flag to enable/disable encryption
) {
const genTitle = useGenTitleMutation();
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
const { token, isAuthenticated } = useAuthContext();
const { token, isAuthenticated, user } = useAuthContext();
const [completed, setCompleted] = useState(new Set());
const setAbortScroll = useSetRecoilState(store.abortScrollFamily(runIndex));
const setShowStopButton = useSetRecoilState(store.showStopButtonByIndex(runIndex));
@ -84,156 +85,158 @@ export default function useSSE(
return;
}
let { userMessage } = submission;
const payloadData = createPayload(submission);
let { payload } = payloadData;
if (isAssistantsEndpoint(payload.endpoint) || isAgentsEndpoint(payload.endpoint)) {
payload = removeNullishValues(payload) as TPayload;
}
let textIndex = null;
const sse = new SSE(payloadData.server, {
payload: JSON.stringify(payload),
headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${token}` },
});
sse.addEventListener('attachment', (e: MessageEvent) => {
try {
const data = JSON.parse(e.data);
attachmentHandler({ data, submission: submission as EventSubmission });
} catch (error) {
console.error(error);
}
});
sse.addEventListener('message', (e: MessageEvent) => {
const data = JSON.parse(e.data);
if (data.final != null) {
const { plugins } = data;
finalHandler(data, { ...submission, plugins } as EventSubmission);
(startupConfig?.checkBalance ?? false) && balanceQuery.refetch();
console.log('final', data);
return;
} else if (data.created != null) {
const runId = v4();
setActiveRunId(runId);
userMessage = {
...userMessage,
...data.message,
overrideParentMessageId: userMessage.overrideParentMessageId,
};
createdHandler(data, { ...submission, userMessage } as EventSubmission);
} else if (data.event != null) {
stepHandler(data, { ...submission, userMessage } as EventSubmission);
} else if (data.sync != null) {
const runId = v4();
setActiveRunId(runId);
/* synchronize messages to Assistants API as well as with real DB ID's */
syncHandler(data, { ...submission, userMessage } as EventSubmission);
} else if (data.type != null) {
const { text, index } = data;
if (text != null && index !== textIndex) {
textIndex = index;
}
contentHandler({ data, submission: submission as EventSubmission });
(async () => {
let payloadData;
// Use encryption if encryptionEnabled is true OR if the user has a public key.
if ((encryptionEnabled || (user && user.encryptionPublicKey)) && user?.encryptionPublicKey) {
payloadData = await createPayloadWithEncryption(
submission,
true,
user.encryptionPublicKey
);
} else {
const text = data.text ?? data.response;
const { plugin, plugins } = data;
const initialResponse = {
...(submission.initialResponse as TMessage),
parentMessageId: data.parentMessageId,
messageId: data.messageId,
};
if (data.message != null) {
messageHandler(text, { ...submission, plugin, plugins, userMessage, initialResponse });
}
}
});
sse.addEventListener('open', () => {
setAbortScroll(false);
console.log('connection is opened');
});
sse.addEventListener('cancel', async () => {
const streamKey = (submission as TSubmission | null)?.['initialResponse']?.messageId;
if (completed.has(streamKey)) {
setIsSubmitting(false);
setCompleted((prev) => {
prev.delete(streamKey);
return new Set(prev);
});
return;
payloadData = createPayload(submission);
}
setCompleted((prev) => new Set(prev.add(streamKey)));
const latestMessages = getMessages();
const conversationId = latestMessages?.[latestMessages.length - 1]?.conversationId;
return await abortConversation(
conversationId ?? userMessage.conversationId ?? submission.conversationId,
submission as EventSubmission,
latestMessages,
);
});
let { payload } = payloadData;
if (isAssistantsEndpoint(payload.endpoint) || isAgentsEndpoint(payload.endpoint)) {
payload = removeNullishValues(payload) as TPayload;
}
sse.addEventListener('error', async (e: MessageEvent) => {
/* @ts-ignore */
if (e.responseCode === 401) {
/* token expired, refresh and retry */
let textIndex: number | null = null;
const sse = new SSE(payloadData.server, {
payload: JSON.stringify(payload),
headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${token}` },
});
sse.addEventListener('attachment', (e: MessageEvent) => {
try {
const refreshResponse = await request.refreshToken();
const token = refreshResponse?.token ?? '';
if (!token) {
throw new Error('Token refresh failed.');
}
sse.headers = {
'Content-Type': 'application/json',
Authorization: `Bearer ${token}`,
const data = JSON.parse(e.data);
attachmentHandler({ data, submission: submission as EventSubmission });
} catch (error) {
console.error(error);
}
});
sse.addEventListener('message', (e: MessageEvent) => {
const data = JSON.parse(e.data);
if (data.final != null) {
const { plugins } = data;
finalHandler(data, { ...submission, plugins } as EventSubmission);
if (startupConfig?.checkBalance) {balanceQuery.refetch();}
console.log('final', data);
return;
} else if (data.created != null) {
const runId = v4();
setActiveRunId(runId);
submission.userMessage = {
...submission.userMessage,
...data.message,
overrideParentMessageId: submission.userMessage.overrideParentMessageId,
};
request.dispatchTokenUpdatedEvent(token);
sse.stream();
return;
} catch (error) {
/* token refresh failed, continue handling the original 401 */
console.log(error);
createdHandler(data, { ...submission } as EventSubmission);
} else if (data.event != null) {
stepHandler(data, { ...submission } as EventSubmission);
} else if (data.sync != null) {
const runId = v4();
setActiveRunId(runId);
syncHandler(data, { ...submission } as EventSubmission);
} else if (data.type != null) {
const { text, index } = data;
if (text != null && index !== textIndex) {
textIndex = index;
}
contentHandler({ data, submission: submission as EventSubmission });
} else {
const text = data.text ?? data.response;
const { plugin, plugins } = data;
const initialResponse = {
...(submission.initialResponse as TMessage),
parentMessageId: data.parentMessageId,
messageId: data.messageId,
};
if (data.message != null) {
messageHandler(text, { ...submission, plugin, plugins, initialResponse });
}
}
}
});
console.log('error in server stream.');
(startupConfig?.checkBalance ?? false) && balanceQuery.refetch();
sse.addEventListener('open', () => {
setAbortScroll(false);
console.log('connection is opened');
});
let data: TResData | undefined = undefined;
try {
data = JSON.parse(e.data) as TResData;
} catch (error) {
console.error(error);
console.log(e);
setIsSubmitting(false);
}
sse.addEventListener('cancel', async () => {
const streamKey = (submission as TSubmission | null)?.['initialResponse']?.messageId;
if (completed.has(streamKey)) {
setIsSubmitting(false);
setCompleted((prev) => {
prev.delete(streamKey);
return new Set(prev);
});
return;
}
errorHandler({ data, submission: { ...submission, userMessage } as EventSubmission });
});
setCompleted((prev) => new Set(prev.add(streamKey)));
const latestMessages = getMessages();
const conversationId =
latestMessages?.[latestMessages.length - 1]?.conversationId;
return await abortConversation(
conversationId ?? submission.userMessage.conversationId ?? submission.conversationId,
submission as EventSubmission,
latestMessages
);
});
setIsSubmitting(true);
sse.stream();
sse.addEventListener('error', async (e: MessageEvent) => {
if ((e as any).responseCode === 401) {
try {
const refreshResponse = await request.refreshToken();
const newToken = refreshResponse?.token ?? '';
if (!newToken) {throw new Error('Token refresh failed.');}
sse.headers = {
'Content-Type': 'application/json',
Authorization: `Bearer ${newToken}`,
};
request.dispatchTokenUpdatedEvent(newToken);
sse.stream();
return;
} catch (error) {
console.log(error);
}
}
return () => {
const isCancelled = sse.readyState <= 1;
sse.close();
if (isCancelled) {
const e = new Event('cancel');
/* @ts-ignore */
sse.dispatchEvent(e);
}
};
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [submission]);
}
console.log('error in server stream.');
if (startupConfig?.checkBalance) {balanceQuery.refetch();}
let data: TResData | undefined = undefined;
try {
data = JSON.parse(e.data) as TResData;
} catch (error) {
console.error(error);
console.log(e);
setIsSubmitting(false);
}
errorHandler({ data, submission: { ...submission } as EventSubmission });
});
setIsSubmitting(true);
sse.stream();
return () => {
const isCancelled = sse.readyState <= 1;
sse.close();
if (isCancelled) {
const event = new Event('cancel');
sse.dispatchEvent(event);
}
};
})();
}, [submission, encryptionEnabled, user?.encryptionPublicKey]);
}

View file

@ -502,6 +502,7 @@ export type TMessage = z.input<typeof tMessageSchema> & {
siblingIndex?: number;
attachments?: TAttachment[];
clientTimestamp?: string;
messageEncryptionIV?: string;
};
export const coerceNumber = z.union([z.number(), z.string()]).transform((val) => {

View file

@ -112,6 +112,7 @@ export type TUser = {
createdAt: string;
updatedAt: string;
encryptionPublicKey?: string;
decryptedPrivateKey?: string;
encryptedPrivateKey?: string; // Encrypted as a Base64 string
encryptionSalt?: string; // Base64 encoded salt used for PBKDF2
encryptionIV?: string; // Base64 encoded IV for AES-GCM