Skip to content

Commit

Permalink
added encryption of sensitive data
Browse files Browse the repository at this point in the history
  • Loading branch information
neven4 committed Sep 23, 2024
1 parent 2280c25 commit ea01acd
Show file tree
Hide file tree
Showing 6 changed files with 78 additions and 6 deletions.
4 changes: 3 additions & 1 deletion .env
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 5 additions & 1 deletion src/lib/server/embeddingEndpoints/hfApi/embeddingHfApi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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 }) => {
Expand All @@ -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: {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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 }) => {
Expand All @@ -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 }),
Expand Down
5 changes: 4 additions & 1 deletion src/lib/server/embeddingEndpoints/tei/embeddingEndpoints.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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,
Expand All @@ -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 }),
});
Expand Down
27 changes: 25 additions & 2 deletions src/lib/server/embeddingModels.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -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,
Expand All @@ -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;
Expand All @@ -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":
Expand Down
37 changes: 37 additions & 0 deletions src/lib/utils/encryption.ts
Original file line number Diff line number Diff line change
@@ -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");
}
}

0 comments on commit ea01acd

Please sign in to comment.