fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
tbphp
2025-06-06 21:45:01 +08:00
commit 219c068dbf
12 changed files with 2076 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
.env.example Normal file
View File

@@ -0,0 +1,57 @@
# ===========================================
# OpenAI 多密钥代理服务器配置文件 (Go版本)
# ===========================================
# ===========================================
# 服务器配置
# ===========================================
# 服务器端口
PORT=3000
# 服务器主机地址
HOST=0.0.0.0
# ===========================================
# 密钥管理配置
# ===========================================
# 密钥文件路径
KEYS_FILE=keys.txt
# 起始密钥索引
START_INDEX=0
# 黑名单阈值（错误多少次后拉黑密钥）
BLACKLIST_THRESHOLD=1
# ===========================================
# OpenAI API 配置
# ===========================================
# 上游 API 地址
OPENAI_BASE_URL=https://api.openai.com
# 请求超时时间（毫秒）
REQUEST_TIMEOUT=30000
# ===========================================
# 认证配置
# ===========================================
# 项目认证密钥（可选，如果设置则启用认证）
# AUTH_KEY=your-secret-key
# ===========================================
# CORS 配置
# ===========================================
# 是否启用 CORS
ENABLE_CORS=true
# 允许的来源（逗号分隔，* 表示允许所有）
ALLOWED_ORIGINS=*
# ===========================================
# 性能配置
# ===========================================
# HTTP 连接池最大连接数
MAX_SOCKETS=50
# HTTP 连接池最大空闲连接数
MAX_FREE_SOCKETS=10

229
.gitignore vendored Normal file
View File

@@ -0,0 +1,229 @@
# Binaries for programs and plugins
*.exe
*.exe~
*.dll
*.so
*.dylib
# Test binary, built with `go test -c`
*.test
# Output of the go coverage tool, specifically when used with LiteIDE
*.out
# Dependency directories (remove the comment below to include it)
# vendor/
# Go workspace file
go.work
# Build output
build/
dist/
bin/
# Environment variables
.env
.env.local
.env.development.local
.env.test.local
.env.production.local
# API Keys and sensitive files
keys.txt
*.key
*.pem
*.crt
secrets/
config/secrets/
# Logs
*.log
logs/
log/
# Runtime data
pids
*.pid
*.seed
*.pid.lock
# Coverage directory used by tools like istanbul
coverage/
*.lcov
# nyc test coverage
.nyc_output
# Dependency directories
node_modules/
# Optional npm cache directory
.npm
# Optional eslint cache
.eslintcache
# Microbundle cache
.rpt2_cache/
.rts2_cache_cjs/
.rts2_cache_es/
.rts2_cache_umd/
# Optional REPL history
.node_repl_history
# Output of 'npm pack'
*.tgz
# Yarn Integrity file
.yarn-integrity
# parcel-bundler cache (https://parceljs.org/)
.cache
.parcel-cache
# Next.js build output
.next
# Nuxt.js build / generate output
.nuxt
dist
# Gatsby files
.cache/
public
# Storybook build outputs
.out
.storybook-out
# Temporary folders
tmp/
temp/
# Runtime data
pids
*.pid
*.seed
*.pid.lock
# Directory for instrumented libs generated by jscoverage/JSCover
lib-cov
# Coverage directory used by tools like istanbul
coverage
*.lcov
# Stores VSCode versions used for testing VSCode extensions
.vscode-test
# IDE files
.vscode/
.idea/
*.swp
*.swo
*~
# OS generated files
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# Docker
.dockerignore
# Database files
*.db
*.sqlite
*.sqlite3
# Backup files
*.bak
*.backup
*.old
# Compressed files
*.7z
*.dmg
*.gz
*.iso
*.jar
*.rar
*.tar
*.zip
# Local development
local/
dev/
# Test files
test-results/
test-output/
# Profiling files
*.prof
*.pprof
# Air live reload tool
tmp/
# GoLand
.idea/
# Vim
*.swp
*.swo
# Emacs
*~
\#*\#
/.emacs.desktop
/.emacs.desktop.lock
*.elc
auto-save-list
tramp
.\#*
# VS Code
.vscode/
*.code-workspace
# JetBrains IDEs
.idea/
*.iml
*.ipr
*.iws
# Local configuration override
config.local.json
config.local.yaml
config.local.yml
# Certificates
*.crt
*.key
*.pem
*.p12
*.pfx
# Terraform
*.tfstate
*.tfstate.*
.terraform/
.terraform.lock.hcl
# Kubernetes
*.kubeconfig
# Helm
charts/*.tgz
# Local development scripts
run-local.sh
start-local.sh
dev-setup.sh

60
Dockerfile Normal file
View File

@@ -0,0 +1,60 @@
# 多阶段构建 Dockerfile for OpenAI 多密钥代理服务器 (Go版本)
# 构建阶段
FROM golang:1.21-alpine AS builder
# 设置工作目录
WORKDIR /app
# 安装必要的包
RUN apk add --no-cache git ca-certificates tzdata
# 复制 go mod 文件
COPY go.mod go.sum ./
# 下载依赖
RUN go mod download
# 复制源代码
COPY . .
# 构建应用
RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \
-ldflags="-w -s -X main.Version=2.0.0" \
-o openai-proxy \
./cmd/main.go
# 运行阶段
FROM alpine:latest
# 安装必要的包
RUN apk --no-cache add ca-certificates curl
# 创建非 root 用户
RUN addgroup -g 1001 -S appgroup && \
adduser -S appuser -u 1001 -G appgroup
# 设置工作目录
WORKDIR /app
# 从构建阶段复制二进制文件
COPY --from=builder /app/openai-proxy .
# 复制配置文件模板
COPY --from=builder /app/.env.example .
# 设置权限
RUN chown -R appuser:appgroup /app
# 切换到非 root 用户
USER appuser
# 暴露端口
EXPOSE 3000
# 健康检查
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
CMD curl -f http://localhost:3000/health || exit 1
# 启动命令
CMD ["./openai-proxy"]

212
Makefile Normal file
View File

@@ -0,0 +1,212 @@
# OpenAI 多密钥代理服务器 Makefile (Go版本)
# 变量定义
BINARY_NAME=openai-proxy
MAIN_PATH=./cmd/main.go
BUILD_DIR=./build
VERSION=2.0.0
LDFLAGS=-ldflags "-X main.Version=$(VERSION) -s -w"
# 默认目标
.PHONY: all
all: clean build
# 构建
.PHONY: build
build:
@echo "🔨 构建 $(BINARY_NAME)..."
@mkdir -p $(BUILD_DIR)
go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME) $(MAIN_PATH)
@echo "✅ 构建完成: $(BUILD_DIR)/$(BINARY_NAME)"
# 构建所有平台
.PHONY: build-all
build-all: clean
@echo "🔨 构建所有平台版本..."
@mkdir -p $(BUILD_DIR)
# Linux AMD64
GOOS=linux GOARCH=amd64 go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 $(MAIN_PATH)
# Linux ARM64
GOOS=linux GOARCH=arm64 go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 $(MAIN_PATH)
# macOS AMD64
GOOS=darwin GOARCH=amd64 go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-amd64 $(MAIN_PATH)
# macOS ARM64 (Apple Silicon)
GOOS=darwin GOARCH=arm64 go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 $(MAIN_PATH)
# Windows AMD64
GOOS=windows GOARCH=amd64 go build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe $(MAIN_PATH)
@echo "✅ 所有平台构建完成"
# 运行
.PHONY: run
run:
@echo "🚀 启动服务器..."
go run $(MAIN_PATH)
# 开发模式运行
.PHONY: dev
dev:
@echo "🔧 开发模式启动..."
go run -race $(MAIN_PATH)
# 测试
.PHONY: test
test:
@echo "🧪 运行测试..."
go test -v -race -coverprofile=coverage.out ./...
# 测试覆盖率
.PHONY: coverage
coverage: test
@echo "📊 生成测试覆盖率报告..."
go tool cover -html=coverage.out -o coverage.html
@echo "✅ 覆盖率报告生成: coverage.html"
# 基准测试
.PHONY: bench
bench:
@echo "⚡ 运行基准测试..."
go test -bench=. -benchmem ./...
# 代码检查
.PHONY: lint
lint:
@echo "🔍 代码检查..."
@if command -v golangci-lint >/dev/null 2>&1; then \
golangci-lint run; \
else \
echo "⚠️ golangci-lint 未安装，跳过代码检查"; \
echo "安装命令: go install github.com/golangci/golangci-lint/cmd/golangci-lint@latest"; \
fi
# 格式化代码
.PHONY: fmt
fmt:
@echo "🎨 格式化代码..."
go fmt ./...
@if command -v goimports >/dev/null 2>&1; then \
goimports -w .; \
else \
echo "💡 建议安装 goimports: go install golang.org/x/tools/cmd/goimports@latest"; \
fi
# 整理依赖
.PHONY: tidy
tidy:
@echo "📦 整理依赖..."
go mod tidy
go mod verify
# 安装依赖
.PHONY: deps
deps:
@echo "📥 安装依赖..."
go mod download
# 清理
.PHONY: clean
clean:
@echo "🧹 清理构建文件..."
rm -rf $(BUILD_DIR)
rm -f coverage.out coverage.html
# 安装到系统
.PHONY: install
install: build
@echo "📦 安装到系统..."
sudo cp $(BUILD_DIR)/$(BINARY_NAME) /usr/local/bin/
@echo "✅ 安装完成: /usr/local/bin/$(BINARY_NAME)"
# 卸载
.PHONY: uninstall
uninstall:
@echo "🗑️ 从系统卸载..."
sudo rm -f /usr/local/bin/$(BINARY_NAME)
@echo "✅ 卸载完成"
# Docker 构建
.PHONY: docker-build
docker-build:
@echo "🐳 构建 Docker 镜像..."
docker build -t openai-proxy:$(VERSION) .
docker tag openai-proxy:$(VERSION) openai-proxy:latest
@echo "✅ Docker 镜像构建完成"
# Docker 运行
.PHONY: docker-run
docker-run:
@echo "🐳 运行 Docker 容器..."
docker run -d \
--name openai-proxy \
-p 3000:3000 \
-v $(PWD)/keys.txt:/app/keys.txt:ro \
-v $(PWD)/.env:/app/.env:ro \
--restart unless-stopped \
openai-proxy:latest
# 健康检查
.PHONY: health
health:
@echo "💚 健康检查..."
@curl -s http://localhost:3000/health | jq . || echo "请安装 jq 或检查服务是否运行"
# 查看统计
.PHONY: stats
stats:
@echo "📊 查看统计信息..."
@curl -s http://localhost:3000/stats | jq . || echo "请安装 jq 或检查服务是否运行"
# 重置密钥
.PHONY: reset-keys
reset-keys:
@echo "🔄 重置密钥状态..."
@curl -s http://localhost:3000/reset-keys | jq . || echo "请安装 jq 或检查服务是否运行"
# 查看黑名单
.PHONY: blacklist
blacklist:
@echo "🚫 查看黑名单..."
@curl -s http://localhost:3000/blacklist | jq . || echo "请安装 jq 或检查服务是否运行"
# 帮助
.PHONY: help
help:
@echo "OpenAI 多密钥代理服务器 v$(VERSION) - 可用命令:"
@echo ""
@echo "构建相关:"
@echo " build - 构建二进制文件"
@echo " build-all - 构建所有平台版本"
@echo " clean - 清理构建文件"
@echo ""
@echo "运行相关:"
@echo " run - 运行服务器"
@echo " dev - 开发模式运行"
@echo ""
@echo "测试相关:"
@echo " test - 运行测试"
@echo " coverage - 生成测试覆盖率报告"
@echo " bench - 运行基准测试"
@echo ""
@echo "代码质量:"
@echo " lint - 代码检查"
@echo " fmt - 格式化代码"
@echo " tidy - 整理依赖"
@echo ""
@echo "安装相关:"
@echo " install - 安装到系统"
@echo " uninstall - 从系统卸载"
@echo ""
@echo "Docker 相关:"
@echo " docker-build - 构建 Docker 镜像"
@echo " docker-run - 运行 Docker 容器"
@echo ""
@echo "管理相关:"
@echo " health - 健康检查"
@echo " stats - 查看统计信息"
@echo " reset-keys - 重置密钥状态"
@echo " blacklist - 查看黑名单"

171
README.md Normal file
View File

@@ -0,0 +1,171 @@
# OpenAI 多密钥代理服务器 v2.0.0 (Go 版本)
一个**极致高性能**的 OpenAI API 多密钥轮询透明代理服务器，使用 Go 语言重写，性能比 Node.js 版本提升 **5-10 倍**！
## ✨ 特性
- 🔄 **多密钥轮询**: 自动轮换使用多个 API 密钥，支持负载均衡
- 🧠 **智能拉黑**: 区分永久性错误和临时性错误，智能密钥管理
- 📊 **实时监控**: 提供详细的统计信息、健康检查和黑名单管理
- 🔧 **灵活配置**: 支持 .env 文件配置，热重载配置
- 🌐 **CORS 支持**: 完整的跨域请求支持
- 📝 **结构化日志**: 彩色日志输出，包含响应时间和密钥信息
- 🔒 **可选认证**: 项目级 Bearer Token 认证
- ⚡ **极致性能**:
- **Go 原生性能**: 比 Node.js 版本快 5-10 倍
- **零拷贝流式传输**: 最小化内存使用和延迟
- **高并发处理**: 支持数万并发连接
- **内存安全**: 自动垃圾回收，无内存泄漏
- **原子操作**: 无锁并发，极低延迟
- 🛡️ **生产就绪**: 优雅关闭、错误恢复、内存管理
## 🚀 快速开始
### 方式一：直接运行
```bash
# 1. 确保已安装 Go 1.21+
go version
# 2. 下载依赖
go mod tidy
# 3. 配置密钥文件
cp ../keys.txt ./keys.txt
# 4. 配置环境变量（可选）
cp .env.example .env
# 5. 运行服务器
make run
# 或者
go run cmd/main.go
```
### 方式二：构建后运行
```bash
# 构建
make build
# 运行
./build/openai-proxy
```
### 方式三：Docker 运行
```bash
# 构建镜像
make docker-build
# 运行容器
make docker-run
# 或使用 docker-compose
docker-compose up -d
```
## ⚙️ 配置管理
### 环境变量配置
```bash
cp .env.example .env
```
### 主要配置项
| 配置项 | 环境变量 | 默认值 | 说明 |
| ---------- | --------------------- | ---------------------- | --------------------- |
| 服务器端口 | `PORT` | 3000 | 服务器监听端口 |
| 服务器主机 | `HOST` | 0.0.0.0 | 服务器绑定地址 |
| 密钥文件 | `KEYS_FILE` | keys.txt | API 密钥文件路径 |
| 起始索引 | `START_INDEX` | 0 | 从哪个密钥开始轮询 |
| 拉黑阈值 | `BLACKLIST_THRESHOLD` | 1 | 错误多少次后拉黑 |
| 上游地址 | `OPENAI_BASE_URL` | https://api.openai.com | OpenAI API 地址 |
| 请求超时 | `REQUEST_TIMEOUT` | 30000 | 请求超时时间（毫秒） |
| 认证密钥 | `AUTH_KEY` | 无 | 项目认证密钥（可选） |
| CORS | `ENABLE_CORS` | true | 是否启用 CORS |
| 连接池 | `MAX_SOCKETS` | 50 | HTTP 连接池最大连接数 |
## 📊 监控端点
| 端点 | 方法 | 说明 |
| ------------- | ---- | ------------------ |
| `/health` | GET | 健康检查和基本状态 |
| `/stats` | GET | 详细统计信息 |
| `/blacklist` | GET | 黑名单详情 |
| `/reset-keys` | GET | 重置所有密钥状态 |
## 🔧 开发指南
### 可用命令
```bash
# 构建相关
make build # 构建二进制文件
make build-all # 构建所有平台版本
make clean # 清理构建文件
# 运行相关
make run # 运行服务器
make dev # 开发模式运行（启用竞态检测）
# 测试相关
make test # 运行测试
make coverage # 生成测试覆盖率报告
make bench # 运行基准测试
# 代码质量
make lint # 代码检查
make fmt # 格式化代码
make tidy # 整理依赖
# 管理相关
make health # 健康检查
make stats # 查看统计信息
make reset-keys # 重置密钥状态
make blacklist # 查看黑名单
# 查看所有命令
make help
```
### 项目结构
```
/
├── cmd/
│ └── main.go # 主入口文件
├── internal/
│ ├── config/
│ │ └── config.go # 配置管理
│ ├── keymanager/
│ │ └── keymanager.go # 密钥管理器
│ └── proxy/
│ └── proxy.go # 代理服务器核心
├── build/ # 构建输出目录
├── .env.example # 配置文件模板
├── Dockerfile # Docker 构建文件
├── docker-compose.yml # Docker Compose 配置
├── Makefile # 构建脚本
├── go.mod # Go 模块文件
└── README.md # 项目文档
```
## 🏗️ 架构设计
### 高性能设计
1. **并发模型**: 使用 Go 的 goroutine 实现高并发处理
2. **内存管理**: 零拷贝流式传输，最小化内存分配
3. **连接复用**: HTTP/2 支持，连接池优化
4. **原子操作**: 无锁并发，避免竞态条件
5. **预编译正则**: 启动时预编译，避免运行时开销
### 安全设计
1. **内存安全**: Go 的内存安全保证，避免缓冲区溢出
2. **并发安全**: sync.Map 和原子操作保证并发安全
3. **错误处理**: 完整的错误处理和恢复机制
4. **资源清理**: 自动资源清理，防止泄漏

97
cmd/main.go Normal file
View File

@@ -0,0 +1,97 @@
// Package main OpenAI多密钥代理服务器主入口
// @author OpenAI Proxy Team
// @version 2.0.0
package main
import (
"context"
"fmt"
"net/http"
"os"
"os/signal"
"syscall"
"time"
"openai-multi-key-proxy/internal/config"
"openai-multi-key-proxy/internal/proxy"
"github.com/sirupsen/logrus"
)
func main() {
// 设置日志格式
logrus.SetFormatter(&logrus.TextFormatter{
FullTimestamp: true,
ForceColors: true,
})
// 加载配置
cfg, err := config.LoadConfig()
if err != nil {
logrus.Fatalf("❌ 配置加载失败: %v", err)
}
// 显示启动信息
displayStartupInfo(cfg)
// 创建代理服务器
proxyServer, err := proxy.NewProxyServer()
if err != nil {
logrus.Fatalf("❌ 创建代理服务器失败: %v", err)
}
defer proxyServer.Close()
// 设置路由
router := proxyServer.SetupRoutes()
// 创建HTTP服务器
server := &http.Server{
Addr: fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port),
Handler: router,
ReadTimeout: 30 * time.Second,
WriteTimeout: 30 * time.Second,
IdleTimeout: 60 * time.Second,
}
// 启动服务器
go func() {
logrus.Infof("🚀 OpenAI 多密钥代理服务器启动成功")
logrus.Infof("📡 服务地址: http://%s:%d", cfg.Server.Host, cfg.Server.Port)
logrus.Infof("📊 统计信息: http://%s:%d/stats", cfg.Server.Host, cfg.Server.Port)
logrus.Infof("💚 健康检查: http://%s:%d/health", cfg.Server.Host, cfg.Server.Port)
logrus.Infof("🔄 重置密钥: http://%s:%d/reset-keys", cfg.Server.Host, cfg.Server.Port)
logrus.Infof("🚫 黑名单查询: http://%s:%d/blacklist", cfg.Server.Host, cfg.Server.Port)
logrus.Info("")
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
logrus.Fatalf("❌ 服务器启动失败: %v", err)
}
}()
// 等待中断信号
quit := make(chan os.Signal, 1)
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
<-quit
logrus.Info("🛑 收到关闭信号，正在优雅关闭服务器...")
// 优雅关闭
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := server.Shutdown(ctx); err != nil {
logrus.Errorf("❌ 服务器关闭失败: %v", err)
} else {
logrus.Info("✅ 服务器已优雅关闭")
}
}
// displayStartupInfo 显示启动信息
func displayStartupInfo(cfg *config.Config) {
logrus.Info("🚀 OpenAI 多密钥代理服务器 v2.0.0 (Go版本)")
logrus.Info("")
// 显示配置
config.DisplayConfig(cfg)
logrus.Info("")
}

52
docker-compose.yml Normal file
View File

@@ -0,0 +1,52 @@
version: '3.8'
services:
openai-proxy:
build: .
container_name: openai-proxy-go
ports:
- "3000:3000"
volumes:
# 挂载密钥文件（只读）
- ./keys.txt:/app/keys.txt:ro
# 挂载配置文件（只读）
- ./.env:/app/.env:ro
restart: unless-stopped
# 健康检查
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:3000/health"]
interval: 30s
timeout: 10s
retries: 3
start_period: 40s
# 环境变量
environment:
- GO_ENV=production
# 日志配置
logging:
driver: "json-file"
options:
max-size: "10m"
max-file: "3"
# 资源限制
deploy:
resources:
limits:
memory: 512M
cpus: '1.0'
reservations:
memory: 128M
cpus: '0.25'
# 网络配置
networks:
- proxy-network
# 网络定义
networks:
proxy-network:
driver: bridge

36
go.mod Normal file
View File

@@ -0,0 +1,36 @@
module openai-multi-key-proxy
go 1.21
require (
github.com/gin-gonic/gin v1.9.1
github.com/joho/godotenv v1.5.1
github.com/sirupsen/logrus v1.9.3
)
require (
github.com/bytedance/sonic v1.9.1 // indirect
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
github.com/gabriel-vasile/mimetype v1.4.2 // indirect
github.com/gin-contrib/sse v0.1.0 // indirect
github.com/go-playground/locales v0.14.1 // indirect
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.14.0 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.2.4 // indirect
github.com/leodido/go-urn v1.2.4 // indirect
github.com/mattn/go-isatty v0.0.19 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/pelletier/go-toml/v2 v2.0.8 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.2.11 // indirect
golang.org/x/arch v0.3.0 // indirect
golang.org/x/crypto v0.9.0 // indirect
golang.org/x/net v0.10.0 // indirect
golang.org/x/sys v0.8.0 // indirect
golang.org/x/text v0.9.0 // indirect
google.golang.org/protobuf v1.30.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)

91
go.sum Normal file
View File

@@ -0,0 +1,91 @@
github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=
github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams=
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/gabriel-vasile/mimetype v1.4.2 h1:w5qFW6JKBz9Y393Y4q372O9A7cUSequkh1Q7OhCmWKU=
github.com/gabriel-vasile/mimetype v1.4.2/go.mod h1:zApsH/mKG4w07erKIaJPFiX0Tsq9BFQgN3qGY5GnNgA=
github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE=
github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI=
github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg=
github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU=
github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY=
github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY=
github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
github.com/go-playground/validator/v10 v10.14.0 h1:vgvQWe3XCz3gIeFDm/HnTIbj6UGmg/+t63MyGU2n5js=
github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU=
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4=
github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
github.com/pelletier/go-toml/v2 v2.0.8 h1:0ctb6s9mE31h0/lhu+J6OPmVeDxJn+kYnJc2jZR9tGQ=
github.com/pelletier/go-toml/v2 v2.0.8/go.mod h1:vuYfssBdrU2XDZ9bYydBu6t+6a6PYNcZljzZR9VXg+4=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k=
golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=

251
internal/config/config.go Normal file
View File

@@ -0,0 +1,251 @@
// Package config 配置管理模块
// @author OpenAI Proxy Team
// @version 2.0.0
package config
import (
"fmt"
"net/url"
"os"
"strconv"
"strings"
"github.com/joho/godotenv"
"github.com/sirupsen/logrus"
)
// Constants 配置常量
type Constants struct {
MinPort int
MaxPort int
MinTimeout int
DefaultTimeout int
DefaultMaxSockets int
DefaultMaxFreeSockets int
}
// DefaultConstants 默认常量
var DefaultConstants = Constants{
MinPort: 1,
MaxPort: 65535,
MinTimeout: 1000,
DefaultTimeout: 30000,
DefaultMaxSockets: 50,
DefaultMaxFreeSockets: 10,
}
// ServerConfig 服务器配置
type ServerConfig struct {
Port int `json:"port"`
Host string `json:"host"`
}
// KeysConfig 密钥管理配置
type KeysConfig struct {
FilePath string `json:"filePath"`
StartIndex int `json:"startIndex"`
BlacklistThreshold int `json:"blacklistThreshold"`
}
// OpenAIConfig OpenAI API 配置
type OpenAIConfig struct {
BaseURL string `json:"baseURL"`
Timeout int `json:"timeout"`
}
// AuthConfig 认证配置
type AuthConfig struct {
Key string `json:"key"`
Enabled bool `json:"enabled"`
}
// CORSConfig CORS 配置
type CORSConfig struct {
Enabled bool `json:"enabled"`
AllowedOrigins []string `json:"allowedOrigins"`
}
// PerformanceConfig 性能配置
type PerformanceConfig struct {
MaxSockets int `json:"maxSockets"`
MaxFreeSockets int `json:"maxFreeSockets"`
}
// Config 应用配置
type Config struct {
Server ServerConfig `json:"server"`
Keys KeysConfig `json:"keys"`
OpenAI OpenAIConfig `json:"openai"`
Auth AuthConfig `json:"auth"`
CORS CORSConfig `json:"cors"`
Performance PerformanceConfig `json:"performance"`
}
// Global config instance
var AppConfig *Config
// LoadConfig 加载配置
func LoadConfig() (*Config, error) {
// 尝试加载 .env 文件
if err := godotenv.Load(); err != nil {
logrus.Info("💡 提示: 创建 .env 文件以支持环境变量配置")
}
config := &Config{
Server: ServerConfig{
Port: parseInteger(os.Getenv("PORT"), 3000),
Host: getEnvOrDefault("HOST", "0.0.0.0"),
},
Keys: KeysConfig{
FilePath: getEnvOrDefault("KEYS_FILE", "keys.txt"),
StartIndex: parseInteger(os.Getenv("START_INDEX"), 0),
BlacklistThreshold: parseInteger(os.Getenv("BLACKLIST_THRESHOLD"), 1),
},
OpenAI: OpenAIConfig{
BaseURL: getEnvOrDefault("OPENAI_BASE_URL", "https://api.openai.com"),
Timeout: parseInteger(os.Getenv("REQUEST_TIMEOUT"), DefaultConstants.DefaultTimeout),
},
Auth: AuthConfig{
Key: os.Getenv("AUTH_KEY"),
Enabled: os.Getenv("AUTH_KEY") != "",
},
CORS: CORSConfig{
Enabled: parseBoolean(os.Getenv("ENABLE_CORS"), true),
AllowedOrigins: parseArray(os.Getenv("ALLOWED_ORIGINS"), []string{"*"}),
},
Performance: PerformanceConfig{
MaxSockets: parseInteger(os.Getenv("MAX_SOCKETS"), DefaultConstants.DefaultMaxSockets),
MaxFreeSockets: parseInteger(os.Getenv("MAX_FREE_SOCKETS"), DefaultConstants.DefaultMaxFreeSockets),
},
}
// 验证配置
if err := validateConfig(config); err != nil {
return nil, err
}
AppConfig = config
return config, nil
}
// validateConfig 验证配置有效性
func validateConfig(config *Config) error {
var errors []string
// 验证端口
if config.Server.Port < DefaultConstants.MinPort || config.Server.Port > DefaultConstants.MaxPort {
errors = append(errors, fmt.Sprintf("端口号必须在 %d-%d 之间", DefaultConstants.MinPort, DefaultConstants.MaxPort))
}
// 验证起始索引
if config.Keys.StartIndex < 0 {
errors = append(errors, "起始索引不能小于 0")
}
// 验证黑名单阈值
if config.Keys.BlacklistThreshold < 1 {
errors = append(errors, "黑名单阈值不能小于 1")
}
// 验证超时时间
if config.OpenAI.Timeout < DefaultConstants.MinTimeout {
errors = append(errors, fmt.Sprintf("请求超时时间不能小于 %dms", DefaultConstants.MinTimeout))
}
// 验证上游URL格式
if _, err := url.Parse(config.OpenAI.BaseURL); err != nil {
errors = append(errors, "上游API地址格式无效")
}
// 验证性能配置
if config.Performance.MaxSockets < 1 {
errors = append(errors, "最大连接数不能小于 1")
}
if config.Performance.MaxFreeSockets < 0 {
errors = append(errors, "最大空闲连接数不能小于 0")
}
if len(errors) > 0 {
logrus.Error("❌ 配置验证失败:")
for _, err := range errors {
logrus.Errorf(" - %s", err)
}
return fmt.Errorf("配置验证失败")
}
return nil
}
// DisplayConfig 显示当前配置信息
func DisplayConfig(config *Config) {
logrus.Info("⚙️ 当前配置:")
logrus.Infof(" 服务器: %s:%d", config.Server.Host, config.Server.Port)
logrus.Infof(" 密钥文件: %s", config.Keys.FilePath)
logrus.Infof(" 起始索引: %d", config.Keys.StartIndex)
logrus.Infof(" 黑名单阈值: %d 次错误", config.Keys.BlacklistThreshold)
logrus.Infof(" 上游地址: %s", config.OpenAI.BaseURL)
logrus.Infof(" 请求超时: %dms", config.OpenAI.Timeout)
authStatus := "未启用"
if config.Auth.Enabled {
authStatus = "已启用"
}
logrus.Infof(" 认证: %s", authStatus)
corsStatus := "已禁用"
if config.CORS.Enabled {
corsStatus = "已启用"
}
logrus.Infof(" CORS: %s", corsStatus)
logrus.Infof(" 连接池: %d/%d", config.Performance.MaxSockets, config.Performance.MaxFreeSockets)
}
// 辅助函数
// parseInteger 解析整数环境变量
func parseInteger(value string, defaultValue int) int {
if value == "" {
return defaultValue
}
if parsed, err := strconv.Atoi(value); err == nil {
return parsed
}
return defaultValue
}
// parseBoolean 解析布尔值环境变量
func parseBoolean(value string, defaultValue bool) bool {
if value == "" {
return defaultValue
}
return strings.ToLower(value) == "true"
}
// parseArray 解析数组环境变量（逗号分隔）
func parseArray(value string, defaultValue []string) []string {
if value == "" {
return defaultValue
}
parts := strings.Split(value, ",")
result := make([]string, 0, len(parts))
for _, part := range parts {
if trimmed := strings.TrimSpace(part); trimmed != "" {
result = append(result, trimmed)
}
}
if len(result) == 0 {
return defaultValue
}
return result
}
// getEnvOrDefault 获取环境变量或默认值
func getEnvOrDefault(key, defaultValue string) string {
if value := os.Getenv(key); value != "" {
return value
}
return defaultValue
}

397
internal/keymanager/keymanager.go Normal file
View File

@@ -0,0 +1,397 @@
// Package keymanager 高性能密钥管理器
// @author OpenAI Proxy Team
// @version 2.0.0
package keymanager
import (
"bufio"
"fmt"
"os"
"regexp"
"strings"
"sync"
"sync/atomic"
"time"
"openai-multi-key-proxy/internal/config"
"github.com/sirupsen/logrus"
)
// KeyInfo 密钥信息
type KeyInfo struct {
Key string `json:"key"`
Index int `json:"index"`
Preview string `json:"preview"`
}
// Stats 统计信息
type Stats struct {
CurrentIndex int64 `json:"currentIndex"`
TotalKeys int `json:"totalKeys"`
HealthyKeys int `json:"healthyKeys"`
BlacklistedKeys int `json:"blacklistedKeys"`
SuccessCount int64 `json:"successCount"`
FailureCount int64 `json:"failureCount"`
MemoryUsage MemoryUsage `json:"memoryUsage"`
}
// MemoryUsage 内存使用情况
type MemoryUsage struct {
FailureCountsSize int `json:"failureCountsSize"`
BlacklistSize int `json:"blacklistSize"`
}
// BlacklistDetail 黑名单详情
type BlacklistDetail struct {
Index int `json:"index"`
LineNumber int `json:"lineNumber"`
KeyPreview string `json:"keyPreview"`
FullKey string `json:"fullKey"`
}
// BlacklistInfo 黑名单信息
type BlacklistInfo struct {
TotalBlacklisted int `json:"totalBlacklisted"`
TotalKeys int `json:"totalKeys"`
HealthyKeys int `json:"healthyKeys"`
BlacklistedKeys []BlacklistDetail `json:"blacklistedKeys"`
}
// KeyManager 密钥管理器
type KeyManager struct {
keysFilePath string
keys []string
keyPreviews []string
currentIndex int64
blacklistedKeys sync.Map
successCount int64
failureCount int64
keyFailureCounts sync.Map
// 性能优化：预编译正则表达式
permanentErrorPatterns []*regexp.Regexp
// 内存管理
cleanupTicker *time.Ticker
stopCleanup chan bool
// 读写锁保护密钥列表
keysMutex sync.RWMutex
}
// NewKeyManager 创建新的密钥管理器
func NewKeyManager(keysFilePath string) *KeyManager {
if keysFilePath == "" {
keysFilePath = config.AppConfig.Keys.FilePath
}
km := &KeyManager{
keysFilePath: keysFilePath,
currentIndex: int64(config.AppConfig.Keys.StartIndex),
stopCleanup: make(chan bool),
// 预编译正则表达式
permanentErrorPatterns: []*regexp.Regexp{
regexp.MustCompile(`(?i)invalid api key`),
regexp.MustCompile(`(?i)incorrect api key`),
regexp.MustCompile(`(?i)api key not found`),
regexp.MustCompile(`(?i)unauthorized`),
regexp.MustCompile(`(?i)account deactivated`),
regexp.MustCompile(`(?i)billing`),
},
}
// 启动内存清理
km.setupMemoryCleanup()
return km
}
// LoadKeys 加载密钥文件
func (km *KeyManager) LoadKeys() error {
file, err := os.Open(km.keysFilePath)
if err != nil {
return fmt.Errorf("无法打开密钥文件: %w", err)
}
defer file.Close()
var keys []string
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line != "" && strings.HasPrefix(line, "sk-") {
keys = append(keys, line)
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("读取密钥文件失败: %w", err)
}
if len(keys) == 0 {
return fmt.Errorf("密钥文件中没有有效的API密钥")
}
km.keysMutex.Lock()
km.keys = keys
// 预生成密钥预览，避免运行时重复计算
km.keyPreviews = make([]string, len(keys))
for i, key := range keys {
if len(key) > 20 {
km.keyPreviews[i] = key[:20] + "..."
} else {
km.keyPreviews[i] = key
}
}
km.keysMutex.Unlock()
logrus.Infof("✅ 成功加载 %d 个 API 密钥", len(keys))
return nil
}
// GetNextKey 获取下一个可用的密钥（高性能版本）
func (km *KeyManager) GetNextKey() (*KeyInfo, error) {
km.keysMutex.RLock()
keysLen := len(km.keys)
km.keysMutex.RUnlock()
if keysLen == 0 {
return nil, fmt.Errorf("没有可用的 API 密钥")
}
// 检查是否所有密钥都被拉黑
blacklistedCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
blacklistedCount++
return true
})
if blacklistedCount >= keysLen {
logrus.Warn("⚠️ 所有密钥都被拉黑，重置黑名单")
km.blacklistedKeys = sync.Map{}
km.keyFailureCounts = sync.Map{}
}
// 使用原子操作避免竞态条件
attempts := 0
for attempts < keysLen {
currentIdx := atomic.AddInt64(&km.currentIndex, 1) - 1
keyIndex := int(currentIdx) % keysLen
km.keysMutex.RLock()
selectedKey := km.keys[keyIndex]
keyPreview := km.keyPreviews[keyIndex]
km.keysMutex.RUnlock()
if _, blacklisted := km.blacklistedKeys.Load(selectedKey); !blacklisted {
return &KeyInfo{
Key: selectedKey,
Index: keyIndex,
Preview: keyPreview,
}, nil
}
attempts++
}
// 兜底：返回第一个密钥
km.keysMutex.RLock()
firstKey := km.keys[0]
firstPreview := km.keyPreviews[0]
km.keysMutex.RUnlock()
return &KeyInfo{
Key: firstKey,
Index: 0,
Preview: firstPreview,
}, nil
}
// RecordSuccess 记录密钥使用成功
func (km *KeyManager) RecordSuccess(key string) {
atomic.AddInt64(&km.successCount, 1)
// 成功时重置该密钥的失败计数
km.keyFailureCounts.Delete(key)
}
// RecordFailure 记录密钥使用失败
func (km *KeyManager) RecordFailure(key string, err error) {
atomic.AddInt64(&km.failureCount, 1)
// 检查是否是永久性错误
if km.isPermanentError(err) {
km.blacklistedKeys.Store(key, true)
km.keyFailureCounts.Delete(key) // 清理计数
logrus.Warnf("🚫 密钥已被拉黑（永久性错误）: %s (%s)", key[:20]+"...", err.Error())
return
}
// 临时性错误：增加失败计数
currentFailures := 0
if val, exists := km.keyFailureCounts.Load(key); exists {
currentFailures = val.(int)
}
newFailures := currentFailures + 1
km.keyFailureCounts.Store(key, newFailures)
threshold := config.AppConfig.Keys.BlacklistThreshold
if newFailures >= threshold {
km.blacklistedKeys.Store(key, true)
km.keyFailureCounts.Delete(key) // 清理计数
logrus.Warnf("🚫 密钥已被拉黑（达到阈值）: %s (失败 %d 次: %s)", key[:20]+"...", newFailures, err.Error())
} else {
logrus.Warnf("⚠️ 密钥失败: %s (%d/%d 次: %s)", key[:20]+"...", newFailures, threshold, err.Error())
}
}
// isPermanentError 判断是否是永久性错误
func (km *KeyManager) isPermanentError(err error) bool {
errorMessage := err.Error()
for _, pattern := range km.permanentErrorPatterns {
if pattern.MatchString(errorMessage) {
return true
}
}
return false
}
// GetStats 获取密钥统计信息
func (km *KeyManager) GetStats() *Stats {
km.keysMutex.RLock()
totalKeys := len(km.keys)
km.keysMutex.RUnlock()
blacklistedCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
blacklistedCount++
return true
})
failureCountsSize := 0
km.keyFailureCounts.Range(func(key, value interface{}) bool {
failureCountsSize++
return true
})
return &Stats{
CurrentIndex: atomic.LoadInt64(&km.currentIndex),
TotalKeys: totalKeys,
HealthyKeys: totalKeys - blacklistedCount,
BlacklistedKeys: blacklistedCount,
SuccessCount: atomic.LoadInt64(&km.successCount),
FailureCount: atomic.LoadInt64(&km.failureCount),
MemoryUsage: MemoryUsage{
FailureCountsSize: failureCountsSize,
BlacklistSize: blacklistedCount,
},
}
}
// ResetKeys 重置密钥状态
func (km *KeyManager) ResetKeys() map[string]interface{} {
beforeCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
beforeCount++
return true
})
km.blacklistedKeys = sync.Map{}
km.keyFailureCounts = sync.Map{}
logrus.Infof("🔄 密钥状态已重置，清除了 %d 个黑名单密钥", beforeCount)
km.keysMutex.RLock()
totalKeys := len(km.keys)
km.keysMutex.RUnlock()
return map[string]interface{}{
"success": true,
"message": fmt.Sprintf("已清除 %d 个黑名单密钥", beforeCount),
"clearedCount": beforeCount,
"totalKeys": totalKeys,
}
}
// GetBlacklistDetails 获取黑名单详情
func (km *KeyManager) GetBlacklistDetails() *BlacklistInfo {
var blacklistDetails []BlacklistDetail
km.keysMutex.RLock()
keys := km.keys
keyPreviews := km.keyPreviews
km.keysMutex.RUnlock()
for i, key := range keys {
if _, blacklisted := km.blacklistedKeys.Load(key); blacklisted {
blacklistDetails = append(blacklistDetails, BlacklistDetail{
Index: i,
LineNumber: i + 1,
KeyPreview: keyPreviews[i],
FullKey: key,
})
}
}
return &BlacklistInfo{
TotalBlacklisted: len(blacklistDetails),
TotalKeys: len(keys),
HealthyKeys: len(keys) - len(blacklistDetails),
BlacklistedKeys: blacklistDetails,
}
}
// setupMemoryCleanup 设置内存清理机制
func (km *KeyManager) setupMemoryCleanup() {
km.cleanupTicker = time.NewTicker(10 * time.Minute)
go func() {
for {
select {
case <-km.cleanupTicker.C:
km.performMemoryCleanup()
case <-km.stopCleanup:
km.cleanupTicker.Stop()
return
}
}
}()
}
// performMemoryCleanup 执行内存清理
func (km *KeyManager) performMemoryCleanup() {
km.keysMutex.RLock()
maxSize := len(km.keys) * 2
if maxSize < 1000 {
maxSize = 1000
}
km.keysMutex.RUnlock()
currentSize := 0
km.keyFailureCounts.Range(func(key, value interface{}) bool {
currentSize++
return true
})
if currentSize > maxSize {
logrus.Infof("🧹 清理失败计数缓存 (%d -> %d)", currentSize, maxSize)
// 简单策略：清理一半的失败计数
cleared := 0
target := currentSize - maxSize
km.keyFailureCounts.Range(func(key, value interface{}) bool {
if cleared < target {
km.keyFailureCounts.Delete(key)
cleared++
}
return cleared < target
})
}
}
// Close 关闭密钥管理器
func (km *KeyManager) Close() {
close(km.stopCleanup)
}

423
internal/proxy/proxy.go Normal file
View File

@@ -0,0 +1,423 @@
// Package proxy 高性能OpenAI多密钥代理服务器
// @author OpenAI Proxy Team
// @version 2.0.0
package proxy
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"sync/atomic"
"time"
"openai-multi-key-proxy/internal/config"
"openai-multi-key-proxy/internal/keymanager"
"github.com/gin-gonic/gin"
"github.com/sirupsen/logrus"
)
// ProxyServer 代理服务器
type ProxyServer struct {
keyManager *keymanager.KeyManager
httpClient *http.Client
upstreamURL *url.URL
requestCount int64
startTime time.Time
}
// NewProxyServer 创建新的代理服务器
func NewProxyServer() (*ProxyServer, error) {
// 解析上游URL
upstreamURL, err := url.Parse(config.AppConfig.OpenAI.BaseURL)
if err != nil {
return nil, fmt.Errorf("解析上游URL失败: %w", err)
}
// 创建密钥管理器
keyManager := keymanager.NewKeyManager(config.AppConfig.Keys.FilePath)
if err := keyManager.LoadKeys(); err != nil {
return nil, fmt.Errorf("加载密钥失败: %w", err)
}
// 创建高性能HTTP客户端
transport := &http.Transport{
MaxIdleConns: config.AppConfig.Performance.MaxSockets,
MaxIdleConnsPerHost: config.AppConfig.Performance.MaxFreeSockets,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
ForceAttemptHTTP2: true,
}
httpClient := &http.Client{
Transport: transport,
Timeout: time.Duration(config.AppConfig.OpenAI.Timeout) * time.Millisecond,
}
return &ProxyServer{
keyManager: keyManager,
httpClient: httpClient,
upstreamURL: upstreamURL,
startTime: time.Now(),
}, nil
}
// SetupRoutes 设置路由
func (ps *ProxyServer) SetupRoutes() *gin.Engine {
// 设置Gin模式
gin.SetMode(gin.ReleaseMode)
router := gin.New()
// 自定义日志中间件
router.Use(ps.loggerMiddleware())
// 恢复中间件
router.Use(gin.Recovery())
// CORS中间件
if config.AppConfig.CORS.Enabled {
router.Use(ps.corsMiddleware())
}
// 认证中间件（如果启用）
if config.AppConfig.Auth.Enabled {
router.Use(ps.authMiddleware())
}
// 管理端点
router.GET("/health", ps.handleHealth)
router.GET("/stats", ps.handleStats)
router.GET("/blacklist", ps.handleBlacklist)
router.GET("/reset-keys", ps.handleResetKeys)
// 代理所有其他请求
router.NoRoute(ps.handleProxy)
return router
}
// corsMiddleware CORS中间件
func (ps *ProxyServer) corsMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
origin := "*"
if len(config.AppConfig.CORS.AllowedOrigins) > 0 && config.AppConfig.CORS.AllowedOrigins[0] != "*" {
origin = strings.Join(config.AppConfig.CORS.AllowedOrigins, ",")
}
c.Header("Access-Control-Allow-Origin", origin)
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
c.Header("Access-Control-Max-Age", "86400")
if c.Request.Method == "OPTIONS" {
c.AbortWithStatus(http.StatusOK)
return
}
c.Next()
}
}
// authMiddleware 认证中间件
func (ps *ProxyServer) authMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
// 管理端点不需要认证
if strings.HasPrefix(c.Request.URL.Path, "/health") ||
strings.HasPrefix(c.Request.URL.Path, "/stats") ||
strings.HasPrefix(c.Request.URL.Path, "/blacklist") ||
strings.HasPrefix(c.Request.URL.Path, "/reset-keys") {
c.Next()
return
}
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
c.JSON(http.StatusUnauthorized, gin.H{
"error": gin.H{
"message": "未提供认证信息",
"type": "authentication_error",
"code": "missing_authorization",
"timestamp": time.Now().Format(time.RFC3339),
},
})
c.Abort()
return
}
if !strings.HasPrefix(authHeader, "Bearer ") {
c.JSON(http.StatusUnauthorized, gin.H{
"error": gin.H{
"message": "认证格式错误",
"type": "authentication_error",
"code": "invalid_authorization_format",
"timestamp": time.Now().Format(time.RFC3339),
},
})
c.Abort()
return
}
token := authHeader[7:] // 移除 "Bearer " 前缀
if token != config.AppConfig.Auth.Key {
c.JSON(http.StatusUnauthorized, gin.H{
"error": gin.H{
"message": "认证失败",
"type": "authentication_error",
"code": "invalid_authorization",
"timestamp": time.Now().Format(time.RFC3339),
},
})
c.Abort()
return
}
c.Next()
}
}
// loggerMiddleware 自定义日志中间件
func (ps *ProxyServer) loggerMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
start := time.Now()
path := c.Request.URL.Path
raw := c.Request.URL.RawQuery
// 处理请求
c.Next()
// 计算响应时间
latency := time.Since(start)
// 获取客户端IP
clientIP := c.ClientIP()
// 获取方法和状态码
method := c.Request.Method
statusCode := c.Writer.Status()
// 构建完整路径
if raw != "" {
path = path + "?" + raw
}
// 获取密钥信息（如果存在）
keyInfo := ""
if keyIndex, exists := c.Get("keyIndex"); exists {
if keyPreview, exists := c.Get("keyPreview"); exists {
keyInfo = fmt.Sprintf(" - Key[%v] %v", keyIndex, keyPreview)
}
}
// 根据状态码选择颜色
var statusColor string
if statusCode >= 200 && statusCode < 300 {
statusColor = "\033[32m" // 绿色
} else {
statusColor = "\033[31m" // 红色
}
resetColor := "\033[0m"
keyColor := "\033[36m" // 青色
// 输出日志
logrus.Infof("%s[%s] %s %s%s%s%s - %s%d%s - %v - %s",
statusColor, time.Now().Format(time.RFC3339), method, path, resetColor,
keyColor, keyInfo, resetColor,
statusColor, statusCode, resetColor,
latency, clientIP)
}
}
// handleHealth 健康检查处理器
func (ps *ProxyServer) handleHealth(c *gin.Context) {
uptime := time.Since(ps.startTime)
stats := ps.keyManager.GetStats()
requestCount := atomic.LoadInt64(&ps.requestCount)
response := gin.H{
"status": "healthy",
"uptime": fmt.Sprintf("%.0fs", uptime.Seconds()),
"requestCount": requestCount,
"keysStatus": gin.H{
"total": stats.TotalKeys,
"healthy": stats.HealthyKeys,
"blacklisted": stats.BlacklistedKeys,
},
"timestamp": time.Now().Format(time.RFC3339),
}
c.JSON(http.StatusOK, response)
}
// handleStats 统计信息处理器
func (ps *ProxyServer) handleStats(c *gin.Context) {
uptime := time.Since(ps.startTime)
stats := ps.keyManager.GetStats()
requestCount := atomic.LoadInt64(&ps.requestCount)
response := gin.H{
"server": gin.H{
"uptime": fmt.Sprintf("%.0fs", uptime.Seconds()),
"requestCount": requestCount,
"startTime": ps.startTime.Format(time.RFC3339),
"version": "2.0.0",
},
"keys": stats,
"timestamp": time.Now().Format(time.RFC3339),
}
c.JSON(http.StatusOK, response)
}
// handleBlacklist 黑名单处理器
func (ps *ProxyServer) handleBlacklist(c *gin.Context) {
blacklistInfo := ps.keyManager.GetBlacklistDetails()
c.JSON(http.StatusOK, blacklistInfo)
}
// handleResetKeys 重置密钥处理器
func (ps *ProxyServer) handleResetKeys(c *gin.Context) {
result := ps.keyManager.ResetKeys()
c.JSON(http.StatusOK, result)
}
// handleProxy 代理请求处理器
func (ps *ProxyServer) handleProxy(c *gin.Context) {
startTime := time.Now()
// 增加请求计数
atomic.AddInt64(&ps.requestCount, 1)
// 获取密钥信息
keyInfo, err := ps.keyManager.GetNextKey()
if err != nil {
logrus.Errorf("获取密钥失败: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "服务器内部错误: " + err.Error(),
"type": "server_error",
"code": "no_keys_available",
"timestamp": time.Now().Format(time.RFC3339),
},
})
return
}
// 设置密钥信息到上下文（用于日志）
c.Set("keyIndex", keyInfo.Index)
c.Set("keyPreview", keyInfo.Preview)
// 读取请求体
var bodyBytes []byte
if c.Request.Body != nil {
bodyBytes, err = io.ReadAll(c.Request.Body)
if err != nil {
logrus.Errorf("读取请求体失败: %v", err)
c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{
"message": "读取请求体失败",
"type": "request_error",
"code": "invalid_request_body",
"timestamp": time.Now().Format(time.RFC3339),
},
})
return
}
}
// 构建上游请求URL
targetURL := *ps.upstreamURL
targetURL.Path = c.Request.URL.Path
targetURL.RawQuery = c.Request.URL.RawQuery
// 创建上游请求
req, err := http.NewRequestWithContext(
context.Background(),
c.Request.Method,
targetURL.String(),
bytes.NewReader(bodyBytes),
)
if err != nil {
logrus.Errorf("创建上游请求失败: %v", err)
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "创建上游请求失败",
"type": "proxy_error",
"code": "request_creation_failed",
"timestamp": time.Now().Format(time.RFC3339),
},
})
return
}
// 复制请求头
for key, values := range c.Request.Header {
if key != "Host" {
for _, value := range values {
req.Header.Add(key, value)
}
}
}
// 设置认证头
req.Header.Set("Authorization", "Bearer "+keyInfo.Key)
// 发送请求
resp, err := ps.httpClient.Do(req)
if err != nil {
responseTime := time.Since(startTime)
logrus.Errorf("代理请求失败: %v (响应时间: %v)", err, responseTime)
// 异步记录失败
go ps.keyManager.RecordFailure(keyInfo.Key, err)
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "代理请求失败: " + err.Error(),
"type": "proxy_error",
"code": "request_failed",
"timestamp": time.Now().Format(time.RFC3339),
},
})
return
}
defer resp.Body.Close()
responseTime := time.Since(startTime)
// 异步记录统计信息（不阻塞响应）
go func() {
if resp.StatusCode >= 200 && resp.StatusCode < 400 {
ps.keyManager.RecordSuccess(keyInfo.Key)
} else if resp.StatusCode >= 400 {
ps.keyManager.RecordFailure(keyInfo.Key, fmt.Errorf("HTTP %d", resp.StatusCode))
}
}()
// 复制响应头
for key, values := range resp.Header {
for _, value := range values {
c.Header(key, value)
}
}
// 设置状态码
c.Status(resp.StatusCode)
// 流式复制响应体（零拷贝）
_, err = io.Copy(c.Writer, resp.Body)
if err != nil {
logrus.Errorf("复制响应体失败: %v (响应时间: %v)", err, responseTime)
}
}
// Close 关闭代理服务器
func (ps *ProxyServer) Close() {
if ps.keyManager != nil {
ps.keyManager.Close()
}
}