From 2f01a2125f3a06ccee635e9c19f46aef10791f00 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Mon, 30 Dec 2024 17:24:19 +0800
Subject: [PATCH] feat: enhance environment variable handling and security
 features

---
 README.md           | 5 +++++
 common/constants.go | 1 +
 common/crypto.go    | 2 +-
 common/init.go      | 7 ++++++-
 main.go             | 4 +++-
 5 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 7c6417c9..704921a5 100644
--- a/README.md
+++ b/README.md
@@ -89,6 +89,7 @@
 - `COHERE_SAFETY_SETTING`：Cohere模型[安全设置](https://docs.cohere.com/docs/safety-modes#overview)，可选值为 `NONE`, `CONTEXTUAL`，`STRICT`，默认为 `NONE`。
 - `GEMINI_VISION_MAX_IMAGE_NUM`：Gemini模型最大图片数量，默认为 `16`，设置为 `-1` 则不限制。
 - `MAX_FILE_DOWNLOAD_MB`: 最大文件下载大小，单位 MB，默认为 `20`。
+- `CRYPTO_SECRET`：加密密钥，用于加密数据库内容。
 ## 部署
 > [!TIP]
 > 最新版Docker镜像：`calciumion/new-api:latest`  
@@ -98,6 +99,10 @@
 > docker run --rm -v /var/run/docker.sock:/var/run/docker.sock containrrr/watchtower -cR
 > ```
 
+### 多机部署
+- 必须设置环境变量 `SESSION_SECRET`，否则会导致多机部署时登录状态不一致。
+- 如果公用Redis，必须设置 `CRYPTO_SECRET`，否则会导致多机部署时Redis内容无法解密。
+
 ### 部署要求
 - 本地数据库（默认）：SQLite（Docker 部署默认使用 SQLite，必须挂载 `/data` 目录到宿主机）
 - 远程数据库：MySQL 版本 >= 5.7.8，PgSQL 版本 >= 9.6
diff --git a/common/constants.go b/common/constants.go
index 4d8999c5..e2acf83b 100644
--- a/common/constants.go
+++ b/common/constants.go
@@ -30,6 +30,7 @@ var DefaultCollapseSidebar = false // default value of collapse sidebar
 // Any options with "Secret", "Token" in its key won't be return by GetOptions
 
 var SessionSecret = uuid.New().String()
+var CryptoSecret = uuid.New().String()
 
 var OptionMap map[string]string
 var OptionMapRWMutex sync.RWMutex
diff --git a/common/crypto.go b/common/crypto.go
index e17f4b01..c353188a 100644
--- a/common/crypto.go
+++ b/common/crypto.go
@@ -14,7 +14,7 @@ func GenerateHMACWithKey(key []byte, data string) string {
 }
 
 func GenerateHMAC(data string) string {
-	h := hmac.New(sha256.New, []byte(SessionSecret))
+	h := hmac.New(sha256.New, []byte(CryptoSecret))
 	h.Write([]byte(data))
 	return hex.EncodeToString(h.Sum(nil))
 }
diff --git a/common/init.go b/common/init.go
index 9de786d6..694e603e 100644
--- a/common/init.go
+++ b/common/init.go
@@ -22,7 +22,7 @@ func printHelp() {
 	fmt.Println("Usage: one-api [--port <port>] [--log-dir <log directory>] [--version] [--help]")
 }
 
-func init() {
+func LoadEnv() {
 	flag.Parse()
 
 	if *PrintVersion {
@@ -45,6 +45,11 @@ func init() {
 			SessionSecret = ss
 		}
 	}
+	if os.Getenv("CRYPTO_SECRET") != "" {
+		CryptoSecret = os.Getenv("CRYPTO_SECRET")
+	} else {
+		CryptoSecret = SessionSecret
+	}
 	if os.Getenv("SQLITE_PATH") != "" {
 		SQLitePath = os.Getenv("SQLITE_PATH")
 	}
diff --git a/main.go b/main.go
index f6fb8cfc..cf1b8be3 100644
--- a/main.go
+++ b/main.go
@@ -33,9 +33,11 @@ var indexPage []byte
 func main() {
 	err := godotenv.Load(".env")
 	if err != nil {
-		common.SysError("failed to load .env file: " + err.Error())
+		common.SysLog("Support for .env file is disabled")
 	}
 
+	common.LoadEnv()
+
 	common.SetupLogger()
 	common.SysLog("New API " + common.Version + " started")
 	if os.Getenv("GIN_MODE") != "debug" {