修改密钥加密逻辑

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 { auth } from "./auth";
import { sign } from "../utils/aws";
import { sign, decrypt } from "../utils/aws";
import { getServerSideConfig } from "../config/server";
import { ModelProvider } from "@/app/constant";
import { prettyObject } from "@/app/utils/format";
import { ModelProvider } from "../constant";
import { prettyObject } from "../utils/format";
const ALLOWED_PATH = new Set(["chat", "models"]);
function parseEventData(chunk: Uint8Array): any {
@ -79,8 +80,6 @@ async function* transformBedrockStream(
const parsed = parseEventData(value);
if (!parsed) continue;
// console.log("Parsed response:", JSON.stringify(parsed, null, 2));
// Handle Titan models
if (modelId.startsWith("amazon.titan")) {
const text = parsed.outputText || "";
@ -191,7 +190,6 @@ async function requestBedrock(req: NextRequest) {
let awsRegion = config.awsRegion;
let awsAccessKey = config.awsAccessKey;
let awsSecretKey = config.awsSecretKey;
let modelId = req.headers.get("ModelID");
// If server-side credentials are not available, parse from Authorization header
if (!awsRegion || !awsAccessKey || !awsSecretKey) {
@ -201,17 +199,25 @@ async function requestBedrock(req: NextRequest) {
}
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");
}
awsRegion = region;
awsAccessKey = accessKey;
awsSecretKey = secretKey;
// Decrypt the credentials
awsRegion = decrypt(encryptedRegion);
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) {
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 { ChatGLMApi } from "./platforms/glm";
import { BedrockApi } from "./platforms/bedrock";
import { encrypt } from "../utils/aws";
export const ROLES = ["system", "user", "assistant"] as const;
export type MessageRole = (typeof ROLES)[number];
@ -279,11 +280,11 @@ export function getHeaders(ignoreHeaders: boolean = false) {
? accessStore.awsRegion &&
accessStore.awsAccessKey &&
accessStore.awsSecretKey
? accessStore.awsRegion +
? encrypt(accessStore.awsRegion) +
":" +
accessStore.awsAccessKey +
encrypt(accessStore.awsAccessKey) +
":" +
accessStore.awsSecretKey
encrypt(accessStore.awsSecretKey)
: ""
: accessStore.openaiApiKey;
return {

View File

@ -980,10 +980,6 @@ export function Settings() {
onChange={(e) =>
accessStore.update((access) => {
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;
})
}
@ -1001,10 +997,6 @@ export function Settings() {
onChange={(e) => {
accessStore.update((access) => {
const accessKey = e.currentTarget.value;
if (accessKey && accessKey.length !== 20) {
showToast(Locale.Settings.Access.Bedrock.AccessKey.Invalid);
return;
}
access.awsAccessKey = accessKey;
});
}}
@ -1023,10 +1015,6 @@ export function Settings() {
onChange={(e) => {
accessStore.update((access) => {
const secretKey = e.currentTarget.value;
if (secretKey && secretKey.length !== 40) {
showToast(Locale.Settings.Access.Bedrock.SecretKey.Invalid);
return;
}
access.awsSecretKey = secretKey;
});
}}

View File

@ -23,35 +23,24 @@ import { createPersistStore } from "../utils/store";
import { ensure } from "../utils/clone";
import { DEFAULT_CONFIG } from "./config";
import { getModelProvider } from "../utils/model";
import { encrypt, decrypt } from "../utils/aws";
let fetchState = 0; // 0 not fetch, 1 fetching, 2 done
const isApp = getClientConfig()?.buildMode === "export";
const DEFAULT_OPENAI_URL = isApp ? OPENAI_BASE_URL : ApiPath.OpenAI;
const DEFAULT_GOOGLE_URL = isApp ? GEMINI_BASE_URL : ApiPath.Google;
const DEFAULT_ANTHROPIC_URL = isApp ? ANTHROPIC_BASE_URL : ApiPath.Anthropic;
const DEFAULT_BAIDU_URL = isApp ? BAIDU_BASE_URL : ApiPath.Baidu;
const DEFAULT_BYTEDANCE_URL = isApp ? BYTEDANCE_BASE_URL : ApiPath.ByteDance;
const DEFAULT_ALIBABA_URL = isApp ? ALIBABA_BASE_URL : ApiPath.Alibaba;
const DEFAULT_TENCENT_URL = isApp ? TENCENT_BASE_URL : ApiPath.Tencent;
const DEFAULT_MOONSHOT_URL = isApp ? MOONSHOT_BASE_URL : ApiPath.Moonshot;
const DEFAULT_STABILITY_URL = isApp ? STABILITY_BASE_URL : ApiPath.Stability;
const DEFAULT_IFLYTEK_URL = isApp ? IFLYTEK_BASE_URL : ApiPath.Iflytek;
const DEFAULT_XAI_URL = isApp ? XAI_BASE_URL : ApiPath.XAI;
const DEFAULT_CHATGLM_URL = isApp ? CHATGLM_BASE_URL : ApiPath.ChatGLM;
const DEFAULT_BEDROCK_URL = isApp ? BEDROCK_BASE_URL : ApiPath.Bedrock;
const DEFAULT_ACCESS_STATE = {
@ -141,17 +130,14 @@ const DEFAULT_ACCESS_STATE = {
export const useAccessStore = createPersistStore(
{ ...DEFAULT_ACCESS_STATE },
(set, get) => ({
enabledAccessControl() {
this.fetch();
return get().needCode;
},
edgeVoiceName() {
this.fetch();
return get().edgeTTSVoiceName;
},
@ -253,7 +239,6 @@ export const useAccessStore = createPersistStore(
DEFAULT_CONFIG.modelConfig.model = model;
DEFAULT_CONFIG.modelConfig.providerName = providerName as any;
}
return res;
})
.then((res: DangerConfig) => {
@ -267,6 +252,31 @@ export const useAccessStore = createPersistStore(
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,

View File

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