Skip to content

Commit ea01acd

Browse files
committed
added encryption of sensitive data
1 parent 2280c25 commit ea01acd

File tree

6 files changed

+78
-6
lines changed

6 files changed

+78
-6
lines changed

.env

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,4 +176,6 @@ HF_ORG_ADMIN=
176176
HF_ORG_EARLY_ACCESS=
177177

178178
PUBLIC_SMOOTH_UPDATES=false
179-
COMMUNITY_TOOLS=false
179+
COMMUNITY_TOOLS=false
180+
181+
ENCRYPTION_KEY=#your encryption key here

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { chunk } from "$lib/utils/chunk";
44
import { env } from "$env/dynamic/private";
55
import { logger } from "$lib/server/logger";
66
import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
7+
import { decrypt } from "$lib/utils/encryption";
78

89
export const embeddingEndpointHfApiSchema = z.object({
910
weight: z.number().int().positive().default(1),
@@ -23,6 +24,9 @@ export async function embeddingEndpointHfApi(
2324
): Promise<EmbeddingEndpoint> {
2425
const { model } = input;
2526
const { authorization } = embeddingEndpointHfApiSchema.parse(input);
27+
28+
const decryptedAuthorization = authorization && decrypt(authorization);
29+
2630
const url = "https://api-inference.huggingface.co/models/" + model.name;
2731

2832
return async ({ inputs }) => {
@@ -35,7 +39,7 @@ export async function embeddingEndpointHfApi(
3539
headers: {
3640
Accept: "application/json",
3741
"Content-Type": "application/json",
38-
...(authorization ? { Authorization: authorization } : {}),
42+
...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
3943
},
4044
body: JSON.stringify({
4145
inputs: {

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { EmbeddingEndpoint, Embedding } from "../embeddingEndpoints";
33
import { chunk } from "$lib/utils/chunk";
44
import { env } from "$env/dynamic/private";
55
import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
6+
import { decrypt } from "$lib/utils/encryption";
67

78
export const embeddingEndpointOpenAIParametersSchema = z.object({
89
weight: z.number().int().positive().default(1),
@@ -22,6 +23,8 @@ export async function embeddingEndpointOpenAI(
2223
const { model } = input;
2324
const { url, apiKey, defaultHeaders } = embeddingEndpointOpenAIParametersSchema.parse(input);
2425

26+
const decryptedApiKey = decrypt(apiKey);
27+
2528
const maxBatchSize = model.maxBatchSize || 100;
2629

2730
return async ({ inputs }) => {
@@ -36,7 +39,7 @@ export async function embeddingEndpointOpenAI(
3639
headers: {
3740
Accept: "application/json",
3841
"Content-Type": "application/json",
39-
...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
42+
...(decryptedApiKey ? { Authorization: `Bearer ${decryptedApiKey}` } : {}),
4043
...defaultHeaders,
4144
},
4245
body: JSON.stringify({ input: batchInputs, model: model.name }),

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { chunk } from "$lib/utils/chunk";
44
import { env } from "$env/dynamic/private";
55
import { logger } from "$lib/server/logger";
66
import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
7+
import { decrypt } from "$lib/utils/encryption";
78

89
export const embeddingEndpointTeiParametersSchema = z.object({
910
weight: z.number().int().positive().default(1),
@@ -45,6 +46,8 @@ export async function embeddingEndpointTei(
4546
const { model } = input;
4647
const { url, authorization } = embeddingEndpointTeiParametersSchema.parse(input);
4748

49+
const decryptedAuthorization = authorization && decrypt(authorization);
50+
4851
const { max_client_batch_size, max_batch_tokens } = await getModelInfoByUrl(url);
4952
const maxBatchSize = Math.min(
5053
max_client_batch_size,
@@ -63,7 +66,7 @@ export async function embeddingEndpointTei(
6366
headers: {
6467
Accept: "application/json",
6568
"Content-Type": "application/json",
66-
...(authorization ? { Authorization: authorization } : {}),
69+
...(decryptedAuthorization ? { Authorization: decryptedAuthorization } : {}),
6770
},
6871
body: JSON.stringify({ inputs: batchInputs, normalize: true, truncate: true }),
6972
});

src/lib/server/embeddingModels.ts

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import JSON5 from "json5";
1212
import type { EmbeddingModel } from "$lib/types/EmbeddingModel";
1313
import { collections } from "./database";
1414
import { ObjectId } from "mongodb";
15+
import { encrypt } from "$lib/utils/encryption";
1516

1617
const modelConfig = z.object({
1718
/** Used as an identifier in DB */
@@ -44,6 +45,29 @@ const rawEmbeddingModelJSON =
4445

4546
const embeddingModelsRaw = z.array(modelConfig).parse(JSON5.parse(rawEmbeddingModelJSON));
4647

48+
const encryptEndpoints = (endpoints: z.infer<typeof modelConfig>["endpoints"]) =>
49+
endpoints.map((endpoint) => {
50+
switch (endpoint.type) {
51+
case "openai":
52+
return {
53+
...endpoint,
54+
apiKey: encrypt(endpoint.apiKey),
55+
};
56+
case "tei":
57+
return {
58+
...endpoint,
59+
authorization: endpoint.authorization && encrypt(endpoint.authorization),
60+
};
61+
case "hfapi":
62+
return {
63+
...endpoint,
64+
authorization: endpoint.authorization && encrypt(endpoint.authorization),
65+
};
66+
default:
67+
return endpoint;
68+
}
69+
});
70+
4771
const embeddingModels = embeddingModelsRaw.map((rawEmbeddingModel) => {
4872
const embeddingModel: EmbeddingModel = {
4973
name: rawEmbeddingModel.name,
@@ -57,7 +81,7 @@ const embeddingModels = embeddingModelsRaw.map((rawEmbeddingModel) => {
5781
_id: new ObjectId(),
5882
createdAt: new Date(),
5983
updatedAt: new Date(),
60-
endpoints: rawEmbeddingModel.endpoints,
84+
endpoints: encryptEndpoints(rawEmbeddingModel.endpoints),
6185
};
6286

6387
return embeddingModel;
@@ -79,7 +103,6 @@ export const getEmbeddingEndpoint = async (embeddingModel: EmbeddingModel) => {
79103
for (const endpoint of embeddingModel.endpoints) {
80104
if (random < endpoint.weight) {
81105
const args = { ...endpoint, model: embeddingModel };
82-
console.log(args.type);
83106

84107
switch (args.type) {
85108
case "tei":

src/lib/utils/encryption.ts

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import crypto from "crypto";
2+
import { env } from "$env/dynamic/private";
3+
4+
if (!env.ENCRYPTION_KEY) {
5+
throw new Error(
6+
"Encryption key is not set, please set the ENCRYPTION_KEY environment variable inside .env.local"
7+
);
8+
}
9+
10+
const algorithm = "aes-256-cbc";
11+
const key = crypto.scryptSync(env.ENCRYPTION_KEY, "salt", 32);
12+
13+
export function encrypt(text: string): string {
14+
try {
15+
const iv = crypto.randomBytes(16);
16+
const cipher = crypto.createCipheriv(algorithm, key, iv);
17+
const encrypted = Buffer.concat([cipher.update(text, "utf8"), cipher.final()]);
18+
return `${iv.toString("hex")}:${encrypted.toString("hex")}`;
19+
} catch (error) {
20+
console.error("Encryption failed:", error);
21+
throw new Error("Encryption failed");
22+
}
23+
}
24+
25+
export function decrypt(text: string): string {
26+
try {
27+
const [ivHex, encryptedHex] = text.split(":");
28+
const iv = Buffer.from(ivHex, "hex");
29+
const encryptedText = Buffer.from(encryptedHex, "hex");
30+
const decipher = crypto.createDecipheriv(algorithm, key, iv);
31+
const decrypted = Buffer.concat([decipher.update(encryptedText), decipher.final()]);
32+
return decrypted.toString("utf8");
33+
} catch (error) {
34+
console.error("Decryption failed:", error);
35+
throw new Error("Decryption failed");
36+
}
37+
}

0 commit comments

Comments
 (0)