修改密钥加密逻辑

This commit is contained in:
glay 2024-11-23 13:52:14 +08:00
parent a85db21e1f
commit ff88421904
5 changed files with 50 additions and 44 deletions

View File

@ -1,9 +1,10 @@
import { NextRequest, NextResponse } from "next/server"; import { NextRequest, NextResponse } from "next/server";
import { auth } from "./auth"; import { auth } from "./auth";
import { sign } from "../utils/aws"; import { sign, decrypt } from "../utils/aws";
import { getServerSideConfig } from "../config/server"; import { getServerSideConfig } from "../config/server";
import { ModelProvider } from "@/app/constant"; import { ModelProvider } from "../constant";
import { prettyObject } from "@/app/utils/format"; import { prettyObject } from "../utils/format";
const ALLOWED_PATH = new Set(["chat", "models"]); const ALLOWED_PATH = new Set(["chat", "models"]);
function parseEventData(chunk: Uint8Array): any { function parseEventData(chunk: Uint8Array): any {
@ -79,8 +80,6 @@ async function* transformBedrockStream(
const parsed = parseEventData(value); const parsed = parseEventData(value);
if (!parsed) continue; if (!parsed) continue;
// console.log("Parsed response:", JSON.stringify(parsed, null, 2));
// Handle Titan models // Handle Titan models
if (modelId.startsWith("amazon.titan")) { if (modelId.startsWith("amazon.titan")) {
const text = parsed.outputText || ""; const text = parsed.outputText || "";
@ -191,7 +190,6 @@ async function requestBedrock(req: NextRequest) {
let awsRegion = config.awsRegion; let awsRegion = config.awsRegion;
let awsAccessKey = config.awsAccessKey; let awsAccessKey = config.awsAccessKey;
let awsSecretKey = config.awsSecretKey; let awsSecretKey = config.awsSecretKey;
let modelId = req.headers.get("ModelID");
// If server-side credentials are not available, parse from Authorization header // If server-side credentials are not available, parse from Authorization header
if (!awsRegion || !awsAccessKey || !awsSecretKey) { if (!awsRegion || !awsAccessKey || !awsSecretKey) {
@ -201,17 +199,25 @@ async function requestBedrock(req: NextRequest) {
} }
const [_, credentials] = authHeader.split("Bearer "); const [_, credentials] = authHeader.split("Bearer ");
const [region, accessKey, secretKey] = credentials.split(":"); console.log("credentials===============" + credentials);
const [encryptedRegion, encryptedAccessKey, encryptedSecretKey] =
credentials.split(":");
if (!region || !accessKey || !secretKey) { if (!encryptedRegion || !encryptedAccessKey || !encryptedSecretKey) {
throw new Error("Invalid Authorization header format"); throw new Error("Invalid Authorization header format");
} }
awsRegion = region; // Decrypt the credentials
awsAccessKey = accessKey; awsRegion = decrypt(encryptedRegion);
awsSecretKey = secretKey; awsAccessKey = decrypt(encryptedAccessKey);
awsSecretKey = decrypt(encryptedSecretKey);
if (!awsRegion || !awsAccessKey || !awsSecretKey) {
throw new Error("Failed to decrypt AWS credentials");
}
} }
let modelId = req.headers.get("ModelID");
if (!awsRegion || !awsAccessKey || !awsSecretKey || !modelId) { if (!awsRegion || !awsAccessKey || !awsSecretKey || !modelId) {
throw new Error("Missing required AWS credentials or model ID"); throw new Error("Missing required AWS credentials or model ID");
} }

View File

@ -23,6 +23,7 @@ import { SparkApi } from "./platforms/iflytek";
import { XAIApi } from "./platforms/xai"; import { XAIApi } from "./platforms/xai";
import { ChatGLMApi } from "./platforms/glm"; import { ChatGLMApi } from "./platforms/glm";
import { BedrockApi } from "./platforms/bedrock"; import { BedrockApi } from "./platforms/bedrock";
import { encrypt } from "../utils/aws";
export const ROLES = ["system", "user", "assistant"] as const; export const ROLES = ["system", "user", "assistant"] as const;
export type MessageRole = (typeof ROLES)[number]; export type MessageRole = (typeof ROLES)[number];
@ -279,11 +280,11 @@ export function getHeaders(ignoreHeaders: boolean = false) {
? accessStore.awsRegion && ? accessStore.awsRegion &&
accessStore.awsAccessKey && accessStore.awsAccessKey &&
accessStore.awsSecretKey accessStore.awsSecretKey
? accessStore.awsRegion + ? encrypt(accessStore.awsRegion) +
":" + ":" +
accessStore.awsAccessKey + encrypt(accessStore.awsAccessKey) +
":" + ":" +
accessStore.awsSecretKey encrypt(accessStore.awsSecretKey)
: "" : ""
: accessStore.openaiApiKey; : accessStore.openaiApiKey;
return { return {

View File

@ -980,10 +980,6 @@ export function Settings() {
onChange={(e) => onChange={(e) =>
accessStore.update((access) => { accessStore.update((access) => {
const region = e.currentTarget.value; const region = e.currentTarget.value;
if (!/^[a-z]{2}-[a-z]+-\d+$/.test(region)) {
showToast(Locale.Settings.Access.Bedrock.Region.Invalid);
return;
}
access.awsRegion = region; access.awsRegion = region;
}) })
} }
@ -1001,10 +997,6 @@ export function Settings() {
onChange={(e) => { onChange={(e) => {
accessStore.update((access) => { accessStore.update((access) => {
const accessKey = e.currentTarget.value; const accessKey = e.currentTarget.value;
if (accessKey && accessKey.length !== 20) {
showToast(Locale.Settings.Access.Bedrock.AccessKey.Invalid);
return;
}
access.awsAccessKey = accessKey; access.awsAccessKey = accessKey;
}); });
}} }}
@ -1023,10 +1015,6 @@ export function Settings() {
onChange={(e) => { onChange={(e) => {
accessStore.update((access) => { accessStore.update((access) => {
const secretKey = e.currentTarget.value; const secretKey = e.currentTarget.value;
if (secretKey && secretKey.length !== 40) {
showToast(Locale.Settings.Access.Bedrock.SecretKey.Invalid);
return;
}
access.awsSecretKey = secretKey; access.awsSecretKey = secretKey;
}); });
}} }}

View File

@ -23,35 +23,24 @@ import { createPersistStore } from "../utils/store";
import { ensure } from "../utils/clone"; import { ensure } from "../utils/clone";
import { DEFAULT_CONFIG } from "./config"; import { DEFAULT_CONFIG } from "./config";
import { getModelProvider } from "../utils/model"; import { getModelProvider } from "../utils/model";
import { encrypt, decrypt } from "../utils/aws";
let fetchState = 0; // 0 not fetch, 1 fetching, 2 done let fetchState = 0; // 0 not fetch, 1 fetching, 2 done
const isApp = getClientConfig()?.buildMode === "export"; const isApp = getClientConfig()?.buildMode === "export";
const DEFAULT_OPENAI_URL = isApp ? OPENAI_BASE_URL : ApiPath.OpenAI; const DEFAULT_OPENAI_URL = isApp ? OPENAI_BASE_URL : ApiPath.OpenAI;
const DEFAULT_GOOGLE_URL = isApp ? GEMINI_BASE_URL : ApiPath.Google; const DEFAULT_GOOGLE_URL = isApp ? GEMINI_BASE_URL : ApiPath.Google;
const DEFAULT_ANTHROPIC_URL = isApp ? ANTHROPIC_BASE_URL : ApiPath.Anthropic; const DEFAULT_ANTHROPIC_URL = isApp ? ANTHROPIC_BASE_URL : ApiPath.Anthropic;
const DEFAULT_BAIDU_URL = isApp ? BAIDU_BASE_URL : ApiPath.Baidu; const DEFAULT_BAIDU_URL = isApp ? BAIDU_BASE_URL : ApiPath.Baidu;
const DEFAULT_BYTEDANCE_URL = isApp ? BYTEDANCE_BASE_URL : ApiPath.ByteDance; const DEFAULT_BYTEDANCE_URL = isApp ? BYTEDANCE_BASE_URL : ApiPath.ByteDance;
const DEFAULT_ALIBABA_URL = isApp ? ALIBABA_BASE_URL : ApiPath.Alibaba; const DEFAULT_ALIBABA_URL = isApp ? ALIBABA_BASE_URL : ApiPath.Alibaba;
const DEFAULT_TENCENT_URL = isApp ? TENCENT_BASE_URL : ApiPath.Tencent; const DEFAULT_TENCENT_URL = isApp ? TENCENT_BASE_URL : ApiPath.Tencent;
const DEFAULT_MOONSHOT_URL = isApp ? MOONSHOT_BASE_URL : ApiPath.Moonshot; const DEFAULT_MOONSHOT_URL = isApp ? MOONSHOT_BASE_URL : ApiPath.Moonshot;
const DEFAULT_STABILITY_URL = isApp ? STABILITY_BASE_URL : ApiPath.Stability; const DEFAULT_STABILITY_URL = isApp ? STABILITY_BASE_URL : ApiPath.Stability;
const DEFAULT_IFLYTEK_URL = isApp ? IFLYTEK_BASE_URL : ApiPath.Iflytek; const DEFAULT_IFLYTEK_URL = isApp ? IFLYTEK_BASE_URL : ApiPath.Iflytek;
const DEFAULT_XAI_URL = isApp ? XAI_BASE_URL : ApiPath.XAI; const DEFAULT_XAI_URL = isApp ? XAI_BASE_URL : ApiPath.XAI;
const DEFAULT_CHATGLM_URL = isApp ? CHATGLM_BASE_URL : ApiPath.ChatGLM; const DEFAULT_CHATGLM_URL = isApp ? CHATGLM_BASE_URL : ApiPath.ChatGLM;
const DEFAULT_BEDROCK_URL = isApp ? BEDROCK_BASE_URL : ApiPath.Bedrock; const DEFAULT_BEDROCK_URL = isApp ? BEDROCK_BASE_URL : ApiPath.Bedrock;
const DEFAULT_ACCESS_STATE = { const DEFAULT_ACCESS_STATE = {
@ -141,17 +130,14 @@ const DEFAULT_ACCESS_STATE = {
export const useAccessStore = createPersistStore( export const useAccessStore = createPersistStore(
{ ...DEFAULT_ACCESS_STATE }, { ...DEFAULT_ACCESS_STATE },
(set, get) => ({ (set, get) => ({
enabledAccessControl() { enabledAccessControl() {
this.fetch(); this.fetch();
return get().needCode; return get().needCode;
}, },
edgeVoiceName() { edgeVoiceName() {
this.fetch(); this.fetch();
return get().edgeTTSVoiceName; return get().edgeTTSVoiceName;
}, },
@ -253,7 +239,6 @@ export const useAccessStore = createPersistStore(
DEFAULT_CONFIG.modelConfig.model = model; DEFAULT_CONFIG.modelConfig.model = model;
DEFAULT_CONFIG.modelConfig.providerName = providerName as any; DEFAULT_CONFIG.modelConfig.providerName = providerName as any;
} }
return res; return res;
}) })
.then((res: DangerConfig) => { .then((res: DangerConfig) => {
@ -267,6 +252,31 @@ export const useAccessStore = createPersistStore(
fetchState = 2; fetchState = 2;
}); });
}, },
// Override the set method to encrypt AWS credentials before storage
set: (partial: { [key: string]: any }) => {
if (partial.awsAccessKey) {
partial.awsAccessKey = encrypt(partial.awsAccessKey);
}
if (partial.awsSecretKey) {
partial.awsSecretKey = encrypt(partial.awsSecretKey);
}
if (partial.awsRegion) {
partial.awsRegion = encrypt(partial.awsRegion);
}
set(partial);
},
// Add getter to decrypt AWS credentials when needed
get: () => {
const state = get();
return {
...state,
awsRegion: state.awsRegion ? decrypt(state.awsRegion) : "",
awsAccessKey: state.awsAccessKey ? decrypt(state.awsAccessKey) : "",
awsSecretKey: state.awsSecretKey ? decrypt(state.awsSecretKey) : "",
};
},
}), }),
{ {
name: StoreKey.Access, name: StoreKey.Access,

View File

@ -17,7 +17,7 @@ export function encrypt(data: string): string {
return AES.encrypt(data, SECRET_KEY).toString(); return AES.encrypt(data, SECRET_KEY).toString();
} catch (error) { } catch (error) {
console.error("Encryption failed:", error); console.error("Encryption failed:", error);
return data; return "";
} }
} }
@ -31,7 +31,8 @@ export function decrypt(encryptedData: string): string {
} }
return decrypted; return decrypted;
} catch (error) { } catch (error) {
return encryptedData; console.error("Decryption failed:", error);
return "";
} }
} }