added encryption of sensitive data

Author: neven4

.env

@@ -176,4 +176,6 @@ HF_ORG_ADMIN=
 HF_ORG_EARLY_ACCESS=
 
 PUBLIC_SMOOTH_UPDATES=false
-COMMUNITY_TOOLS=false
+COMMUNITY_TOOLS=false
+
+ENCRYPTION_KEY=#your encryption key here

src/lib/server/embeddingEndpoints/hfApi/embeddingHfApi.ts

@@ -4,6 +4,7 @@ import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
 import { logger } from "$lib/server/logger";
 import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointHfApiSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -23,6 +24,9 @@ export async function embeddingEndpointHfApi(
 ): Promise<EmbeddingEndpoint> {
 	const { model } = input;
 	const { authorization } = embeddingEndpointHfApiSchema.parse(input);
+
+	const decryptedAuthorization = authorization && decrypt(authorization);
+
 	const url = "https://api-inference.huggingface.co/models/" + model.name;
 
 	return async ({ inputs }) => {
@@ -35,7 +39,7 @@ export async function embeddingEndpointHfApi(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(authorization ? { Authorization: authorization } : {}),
+						...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
 					},
 					body: JSON.stringify({
 						inputs: {

src/lib/server/embeddingEndpoints/openai/embeddingEndpoints.ts

@@ -3,6 +3,7 @@ import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
 import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
 import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointOpenAIParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -22,6 +23,8 @@ export async function embeddingEndpointOpenAI(
 	const { model } = input;
 	const { url, apiKey, defaultHeaders } = embeddingEndpointOpenAIParametersSchema.parse(input);
 
+	const decryptedApiKey = decrypt(apiKey);
+
 	const maxBatchSize = model.maxBatchSize || 100;
 
 	return async ({ inputs }) => {
@@ -36,7 +39,7 @@ export async function embeddingEndpointOpenAI(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
+						...(decryptedApiKey ? { Authorization: `Bearer ${decryptedApiKey}` } : {}),
 						...defaultHeaders,
 					},
 					body: JSON.stringify({ input: batchInputs, model: model.name }),

src/lib/server/embeddingEndpoints/tei/embeddingEndpoints.ts

@@ -4,6 +4,7 @@ import { chunk } from "$lib/utils/chunk";
 import { env } from "$env/dynamic/private";
 import { logger } from "$lib/server/logger";
 import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
+import { decrypt } from "$lib/utils/encryption";
 
 export const embeddingEndpointTeiParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -45,6 +46,8 @@ export async function embeddingEndpointTei(
 	const { model } = input;
 	const { url, authorization } = embeddingEndpointTeiParametersSchema.parse(input);
 
+	const decryptedAuthorization = authorization && decrypt(authorization);
+
 	const { max_client_batch_size, max_batch_tokens } = await getModelInfoByUrl(url);
 	const maxBatchSize = Math.min(
 		max_client_batch_size,
@@ -63,7 +66,7 @@ export async function embeddingEndpointTei(
 					headers: {
 						Accept: "application/json",
 						"Content-Type": "application/json",
-						...(authorization ? { Authorization: authorization } : {}),
+						...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
 					},
 					body: JSON.stringify({ inputs: batchInputs, normalize: true, truncate: true }),
 				});

src/lib/server/embeddingModels.ts

@@ -12,6 +12,7 @@ import JSON5 from "json5";
 import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
 import { collections } from "./database";
 import { ObjectId } from "mongodb";
+import { encrypt } from "$lib/utils/encryption";
 
 const modelConfig = z.object({
 	/** Used as an identifier in DB */
@@ -44,6 +45,29 @@ const rawEmbeddingModelJSON =
 
 const embeddingModelsRaw = z.array(modelConfig).parse(JSON5.parse(rawEmbeddingModelJSON));
 
+const encryptEndpoints = (endpoints: z.infer<typeof modelConfig>["endpoints"]) =>
+	endpoints.map((endpoint) => {
+		switch (endpoint.type) {
+			case "openai":
+				return {
+					...endpoint,
+					apiKey: encrypt(endpoint.apiKey),
+				};
+			case "tei":
+				return {
+					...endpoint,
+					authorization: endpoint.authorization && encrypt(endpoint.authorization),
+				};
+			case "hfapi":
+				return {
+					...endpoint,
+					authorization: endpoint.authorization && encrypt(endpoint.authorization),
+				};
+			default:
+				return endpoint;
+		}
+	});
+
 const embeddingModels = embeddingModelsRaw.map((rawEmbeddingModel) => {
 	const embeddingModel: EmbeddingModel = {
 		name: rawEmbeddingModel.name,
@@ -57,7 +81,7 @@ const embeddingModels = embeddingModelsRaw.map((rawEmbeddingModel) => {
 		_id: new ObjectId(),
 		createdAt: new Date(),
 		updatedAt: new Date(),
-		endpoints: rawEmbeddingModel.endpoints,
+		endpoints: encryptEndpoints(rawEmbeddingModel.endpoints),
 	};
 
 	return embeddingModel;
@@ -79,7 +103,6 @@ export const getEmbeddingEndpoint = async (embeddingModel: EmbeddingModel) => {
 	for (const endpoint of embeddingModel.endpoints) {
 		if (random < endpoint.weight) {
 			const args = { ...endpoint, model: embeddingModel };
-			console.log(args.type);
 
 			switch (args.type) {
 				case "tei":

src/lib/utils/encryption.ts

@@ -0,0 +1,37 @@
+import crypto from "crypto";
+import { env } from "$env/dynamic/private";
+
+if (!env.ENCRYPTION_KEY) {
+	throw new Error(
+		"Encryption key is not set, please set the ENCRYPTION_KEY environment variable inside .env.local"
+	);
+}
+
+const algorithm = "aes-256-cbc";
+const key = crypto.scryptSync(env.ENCRYPTION_KEY, "salt", 32);
+
+export function encrypt(text: string): string {
+	try {
+		const iv = crypto.randomBytes(16);
+		const cipher = crypto.createCipheriv(algorithm, key, iv);
+		const encrypted = Buffer.concat([cipher.update(text, "utf8"), cipher.final()]);
+		return `${iv.toString("hex")}:${encrypted.toString("hex")}`;
+	} catch (error) {
+		console.error("Encryption failed:", error);
+		throw new Error("Encryption failed");
+	}
+}
+
+export function decrypt(text: string): string {
+	try {
+		const [ivHex, encryptedHex] = text.split(":");
+		const iv = Buffer.from(ivHex, "hex");
+		const encryptedText = Buffer.from(encryptedHex, "hex");
+		const decipher = crypto.createDecipheriv(algorithm, key, iv);
+		const decrypted = Buffer.concat([decipher.update(encryptedText), decipher.final()]);
+		return decrypted.toString("utf8");
+	} catch (error) {
+		console.error("Decryption failed:", error);
+		throw new Error("Decryption failed");
+	}
+}