From 5ac651ad8e259934d183e46d3d284ffc925dfc75 Mon Sep 17 00:00:00 2001
From: glay <glay.guo@gmail.com>
Date: Sat, 7 Dec 2024 20:34:05 +0800
Subject: [PATCH] Enhance encryption security

---
 app/api/bedrock.ts              |  8 ++---
 app/client/platforms/bedrock.ts | 18 ++++------
 app/utils/aws.ts                | 63 +++++++++++++++++++++++++++++----
 3 files changed, 67 insertions(+), 22 deletions(-)

diff --git a/app/api/bedrock.ts b/app/api/bedrock.ts
index 7da14a17b..3183a1162 100644
--- a/app/api/bedrock.ts
+++ b/app/api/bedrock.ts
@@ -103,10 +103,10 @@ async function requestBedrock(req: NextRequest) {
     });
 
     // Make request to AWS Bedrock
-    // console.log(
-    //   "[Bedrock Request] Final Body:",
-    //   JSON.stringify(requestBody, null, 2),
-    // );
+    console.log(
+      "[Bedrock Request] Final Body:",
+      JSON.stringify(requestBody, null, 2),
+    );
     const res = await fetch(endpoint, {
       method: "POST",
       headers,
diff --git a/app/client/platforms/bedrock.ts b/app/client/platforms/bedrock.ts
index cd08cb3bb..7311c8a66 100644
--- a/app/client/platforms/bedrock.ts
+++ b/app/client/platforms/bedrock.ts
@@ -137,7 +137,8 @@ export class BedrockApi implements LLMApi {
               },
             },
           })),
-          // toolChoice: { auto: {} }
+          toolChoice: { auto: {} },
+          //  toolChoice: { any: {} }
         };
       }
 
@@ -339,16 +340,11 @@ export class BedrockApi implements LLMApi {
     const controller = new AbortController();
     options.onController?.(controller);
 
-    if (!accessStore.isValidBedrock()) {
-      throw new Error(
-        "Invalid AWS credentials. Please check your configuration and ensure ENCRYPTION_KEY is set.",
-      );
-    }
-
     let finalRequestBody = this.formatRequestBody(messages, modelConfig);
 
     try {
       const isApp = !!getClientConfig()?.isApp;
+      // const isApp = true;
       const bedrockAPIPath = `${BEDROCK_BASE_URL}/model/${
         modelConfig.model
       }/invoke${shouldStream ? "-with-response-stream" : ""}`;
@@ -699,10 +695,10 @@ function bedrockStream(
       responseRes = res;
 
       const contentType = res.headers.get("content-type");
-      console.log(
-        "[Bedrock Stream Request] response content type: ",
-        contentType,
-      );
+      // console.log(
+      //   "[Bedrock Stream Request] response content type: ",
+      //   contentType,
+      // );
 
       if (contentType?.startsWith("text/plain")) {
         responseText = await res.text();
diff --git a/app/utils/aws.ts b/app/utils/aws.ts
index dffff338d..22c6d0baa 100644
--- a/app/utils/aws.ts
+++ b/app/utils/aws.ts
@@ -2,7 +2,7 @@ import SHA256 from "crypto-js/sha256";
 import HmacSHA256 from "crypto-js/hmac-sha256";
 import Hex from "crypto-js/enc-hex";
 import Utf8 from "crypto-js/enc-utf8";
-import { AES, enc } from "crypto-js";
+import { AES, enc, lib, PBKDF2, mode, pad, algo } from "crypto-js";
 
 // Types and Interfaces
 export interface BedrockCredentials {
@@ -29,13 +29,51 @@ type ParsedEvent = Record<string, any>;
 type EventResult = ParsedEvent[];
 
 // Encryption utilities
+function generateSalt(): string {
+  const salt = lib.WordArray.random(128 / 8);
+  return salt.toString(enc.Base64);
+}
+
+function generateIV(): string {
+  const iv = lib.WordArray.random(128 / 8);
+  return iv.toString(enc.Base64);
+}
+
+function deriveKey(password: string, salt: string): lib.WordArray {
+  // Use PBKDF2 with SHA256 for key derivation
+  return PBKDF2(password, salt, {
+    keySize: 256 / 32,
+    iterations: 10000,
+    hasher: algo.SHA256,
+  });
+}
+
+// Using a dot as separator since it's not used in Base64
+const SEPARATOR = ".";
+
 export function encrypt(data: string, encryptionKey: string): string {
   if (!data) return "";
   if (!encryptionKey) {
     throw new Error("Encryption key is required for AWS credential encryption");
   }
   try {
-    return AES.encrypt(data, encryptionKey).toString();
+    // Generate salt and IV
+    const salt = generateSalt();
+    const iv = generateIV();
+
+    // Derive key using PBKDF2
+    const key = deriveKey(encryptionKey, salt);
+
+    // Encrypt the data
+    const encrypted = AES.encrypt(data, key, {
+      iv: enc.Base64.parse(iv),
+      mode: mode.CBC,
+      padding: pad.Pkcs7,
+    });
+
+    // Combine salt, IV, and encrypted data
+    // Format: salt.iv.encryptedData
+    return [salt, iv, encrypted.toString()].join(SEPARATOR);
   } catch (error) {
     throw new Error("Failed to encrypt AWS credentials");
   }
@@ -47,12 +85,21 @@ export function decrypt(encryptedData: string, encryptionKey: string): string {
     throw new Error("Encryption key is required for AWS credential decryption");
   }
   try {
-    const bytes = AES.decrypt(encryptedData, encryptionKey);
-    const decrypted = bytes.toString(enc.Utf8);
-    if (!decrypted && encryptedData) {
+    let components = encryptedData.split(SEPARATOR);
+    const [salt, iv, data] = components;
+    // For new format, use the provided salt and IV
+    const key = deriveKey(encryptionKey, salt);
+    const decrypted = AES.decrypt(data, key, {
+      iv: enc.Base64.parse(iv),
+      mode: mode.CBC,
+      padding: pad.Pkcs7,
+    });
+
+    const result = decrypted.toString(enc.Utf8);
+    if (!result) {
       throw new Error("Failed to decrypt AWS credentials");
     }
-    return decrypted;
+    return result;
   } catch (error) {
     throw new Error("Failed to decrypt AWS credentials");
   }
@@ -61,7 +108,9 @@ export function decrypt(encryptedData: string, encryptionKey: string): string {
 export function maskSensitiveValue(value: string): string {
   if (!value) return "";
   if (value.length <= 4) return value;
-  return "*".repeat(value.length - 4) + value.slice(-4);
+  // Use constant-time operations to prevent timing attacks
+  const masked = Buffer.alloc(value.length - 4, "*").toString();
+  return value.slice(0, 2) + masked + value.slice(-2);
 }
 
 // AWS Signing