mirror of
https://github.com/BlueSkyXN/AI2API.git
synced 2026-04-22 15:47:19 +08:00
BlueSkyXN
520 lines
14 KiB
Go
520 lines
14 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/json"
|
|
"flag"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"path"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// Config 结构体用于存储命令行参数配置
|
|
type Config struct {
|
|
KeyFile string // API 密钥文件路径
|
|
TargetURL string // 目标 API 基础 URL
|
|
Port string // 代理服务器监听端口
|
|
Address string // 代理服务器监听地址
|
|
Password string // 客户端身份验证密码
|
|
MaxWorkers int // 最大工作协程数
|
|
MaxQueue int // 最大请求队列长度
|
|
}
|
|
|
|
// parseFlags 解析命令行参数并返回 Config 实例
|
|
func parseFlags() *Config {
|
|
cfg := &Config{}
|
|
flag.StringVar(&cfg.KeyFile, "key-file", "", "Path to the API key file")
|
|
flag.StringVar(&cfg.TargetURL, "target-url", "", "Target API base URL")
|
|
flag.StringVar(&cfg.Port, "port", "8080", "Port to listen on")
|
|
flag.StringVar(&cfg.Address, "address", "localhost", "Address to listen on")
|
|
flag.StringVar(&cfg.Password, "password", "", "Password for client authentication")
|
|
|
|
// 添加WorkerPool相关配置
|
|
maxWorkers := flag.Int("max-workers", 50, "Maximum number of worker goroutines")
|
|
maxQueue := flag.Int("max-queue", 500, "Maximum size of request queue")
|
|
|
|
flag.Parse()
|
|
|
|
// 将WorkerPool配置添加到Config结构体
|
|
cfg.MaxWorkers = *maxWorkers
|
|
cfg.MaxQueue = *maxQueue
|
|
|
|
return cfg
|
|
}
|
|
|
|
// KeyPool 管理 API 密钥池
|
|
type KeyPool struct {
|
|
keys []string // 密钥列表
|
|
mu sync.Mutex // 互斥锁，确保线程安全
|
|
currentIndex int // 当前密钥索引，用于循环抽取
|
|
}
|
|
|
|
// NewKeyPool 从文件中加载密钥并创建 KeyPool 实例
|
|
func NewKeyPool(filePath string) (*KeyPool, error) {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to open key file %s: %v", filePath, err)
|
|
return nil, err
|
|
}
|
|
defer file.Close()
|
|
|
|
var keys []string
|
|
scanner := bufio.NewScanner(file)
|
|
for scanner.Scan() {
|
|
key := strings.TrimSpace(scanner.Text())
|
|
if key != "" {
|
|
keys = append(keys, key)
|
|
}
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
log.Printf("[ERROR] Failed to read key file %s: %v", filePath, err)
|
|
return nil, err
|
|
}
|
|
log.Printf("[INFO] Loaded %d keys from file %s", len(keys), filePath)
|
|
return &KeyPool{keys: keys, currentIndex: 0}, nil
|
|
}
|
|
|
|
// GetRandomKey 按顺序循环返回一个密钥
|
|
func (kp *KeyPool) GetRandomKey() string {
|
|
kp.mu.Lock()
|
|
defer kp.mu.Unlock()
|
|
if len(kp.keys) == 0 {
|
|
return ""
|
|
}
|
|
key := kp.keys[kp.currentIndex]
|
|
kp.currentIndex = (kp.currentIndex + 1) % len(kp.keys) // 循环到下一个索引
|
|
return key
|
|
}
|
|
|
|
// 定义请求结构体
|
|
type ProxyRequest struct {
|
|
Request *http.Request
|
|
Response http.ResponseWriter
|
|
Done chan bool // 用于通知请求处理完成
|
|
}
|
|
|
|
// Worker结构体，表示一个工作协程
|
|
type Worker struct {
|
|
ID int
|
|
TaskQueue chan *ProxyRequest // 任务队列
|
|
Quit chan bool // 退出信号
|
|
WorkerPool *WorkerPool // 所属工作池
|
|
}
|
|
|
|
// 创建新的Worker
|
|
func NewWorker(id int, workerPool *WorkerPool) *Worker {
|
|
return &Worker{
|
|
ID: id,
|
|
TaskQueue: make(chan *ProxyRequest),
|
|
Quit: make(chan bool),
|
|
WorkerPool: workerPool,
|
|
}
|
|
}
|
|
|
|
// Worker开始工作
|
|
func (w *Worker) Start() {
|
|
go func() {
|
|
for {
|
|
// 将worker注册到工作池的空闲队列
|
|
w.WorkerPool.WorkerQueue <- w.TaskQueue
|
|
|
|
select {
|
|
case task := <-w.TaskQueue:
|
|
// 处理请求
|
|
w.WorkerPool.HandleFunc(task.Response, task.Request)
|
|
task.Done <- true
|
|
case <-w.Quit:
|
|
// 收到退出信号
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Worker停止工作
|
|
func (w *Worker) Stop() {
|
|
go func() {
|
|
w.Quit <- true
|
|
}()
|
|
}
|
|
|
|
// WorkerPool结构体，管理工作协程池
|
|
type WorkerPool struct {
|
|
WorkerQueue chan chan *ProxyRequest // 空闲Worker队列
|
|
TaskQueue chan *ProxyRequest // 任务队列
|
|
MaxWorkers int // 最大Worker数量
|
|
MaxQueue int // 最大队列长度
|
|
HandleFunc func(http.ResponseWriter, *http.Request) // 请求处理函数
|
|
}
|
|
|
|
// 创建新的WorkerPool
|
|
func NewWorkerPool(maxWorkers int, maxQueue int, handleFunc func(http.ResponseWriter, *http.Request)) *WorkerPool {
|
|
pool := &WorkerPool{
|
|
WorkerQueue: make(chan chan *ProxyRequest, maxWorkers),
|
|
TaskQueue: make(chan *ProxyRequest, maxQueue),
|
|
MaxWorkers: maxWorkers,
|
|
MaxQueue: maxQueue,
|
|
HandleFunc: handleFunc,
|
|
}
|
|
return pool
|
|
}
|
|
|
|
// 启动WorkerPool
|
|
func (wp *WorkerPool) Start() {
|
|
// 创建并启动workers
|
|
for i := 0; i < wp.MaxWorkers; i++ {
|
|
worker := NewWorker(i, wp)
|
|
worker.Start()
|
|
log.Printf("[INFO] Started worker %d", i)
|
|
}
|
|
|
|
// 启动任务分发协程
|
|
go wp.dispatch()
|
|
}
|
|
|
|
// 停止WorkerPool
|
|
func (wp *WorkerPool) Stop() {
|
|
// TODO: 实现停止逻辑
|
|
}
|
|
|
|
// 将任务分发给空闲worker
|
|
func (wp *WorkerPool) dispatch() {
|
|
for {
|
|
select {
|
|
case task := <-wp.TaskQueue:
|
|
// 等待空闲worker
|
|
workerTaskQueue := <-wp.WorkerQueue
|
|
// 将任务发送给worker
|
|
workerTaskQueue <- task
|
|
}
|
|
}
|
|
}
|
|
|
|
// 将请求提交到WorkerPool
|
|
func (wp *WorkerPool) Submit(response http.ResponseWriter, request *http.Request) bool {
|
|
task := &ProxyRequest{
|
|
Request: request,
|
|
Response: response,
|
|
Done: make(chan bool, 1),
|
|
}
|
|
|
|
select {
|
|
case wp.TaskQueue <- task:
|
|
// 请求成功加入队列
|
|
<-task.Done // 等待任务完成
|
|
return true
|
|
default:
|
|
// 队列已满，实现背压
|
|
log.Println("[WARN] Task queue is full, rejecting request")
|
|
http.Error(response, "Server is busy, please try again later", http.StatusServiceUnavailable)
|
|
return false
|
|
}
|
|
}
|
|
|
|
// ProxyHandler 处理 HTTP 代理请求
|
|
type ProxyHandler struct {
|
|
cfg *Config // 配置信息
|
|
keyPool *KeyPool // 密钥池
|
|
client *http.Client // HTTP 客户端
|
|
workerPool *WorkerPool // 工作协程池
|
|
}
|
|
|
|
// NewProxyHandler 创建 ProxyHandler 实例
|
|
func NewProxyHandler(cfg *Config, keyPool *KeyPool) *ProxyHandler {
|
|
handler := &ProxyHandler{
|
|
cfg: cfg,
|
|
keyPool: keyPool,
|
|
client: &http.Client{},
|
|
}
|
|
return handler
|
|
}
|
|
|
|
// InitWorkerPool 初始化工作协程池
|
|
func (ph *ProxyHandler) InitWorkerPool(maxWorkers int, maxQueue int) {
|
|
ph.workerPool = NewWorkerPool(maxWorkers, maxQueue, ph.HandleRequest)
|
|
ph.workerPool.Start()
|
|
log.Printf("[INFO] Started worker pool with %d workers and queue size %d", maxWorkers, maxQueue)
|
|
}
|
|
|
|
// ServeHTTP 实现 HTTP 处理逻辑
|
|
func (ph *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
// 记录接收到的请求
|
|
log.Printf("[INFO] Received request: %s %s", r.Method, r.URL.String())
|
|
|
|
// 将请求提交到工作池处理
|
|
ph.workerPool.Submit(w, r)
|
|
}
|
|
|
|
// 简化的模型提取函数
|
|
func extractModelInfo(bodyData []byte) (string, error) {
|
|
var data map[string]interface{}
|
|
if err := json.Unmarshal(bodyData, &data); err != nil {
|
|
return "", err
|
|
}
|
|
if model, ok := data["model"].(string); ok {
|
|
return model, nil
|
|
}
|
|
return "", nil
|
|
}
|
|
|
|
// 检测是否为流式请求
|
|
func isStreamRequest(r *http.Request, bodyData []byte) bool {
|
|
// 检查URL参数
|
|
if strings.Contains(r.URL.RawQuery, "stream=true") {
|
|
return true
|
|
}
|
|
|
|
// 检查请求体
|
|
var data map[string]interface{}
|
|
if err := json.Unmarshal(bodyData, &data); err == nil {
|
|
if stream, ok := data["stream"].(bool); ok && stream {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// HandleRequest 处理请求的方法，由Worker调用
|
|
func (ph *ProxyHandler) HandleRequest(w http.ResponseWriter, r *http.Request) {
|
|
// 验证客户端身份
|
|
if !ph.authenticate(r) {
|
|
log.Println("[WARN] Unauthorized access attempt")
|
|
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
log.Println("[INFO] Authentication successful")
|
|
|
|
// 读取请求体 - 只读取一次
|
|
var bodyBytes []byte
|
|
var err error
|
|
if r.Body != nil {
|
|
bodyBytes, err = io.ReadAll(r.Body)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to read request body: %v", err)
|
|
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
// 异步提取模型信息
|
|
go func() {
|
|
if model, err := extractModelInfo(bodyBytes); err == nil && model != "" {
|
|
log.Printf("[INFO] Model specified in request: %s", model)
|
|
}
|
|
}()
|
|
|
|
// 检测是否为流式请求
|
|
if isStreamRequest(r, bodyBytes) {
|
|
log.Println("[INFO] Streaming request detected")
|
|
}
|
|
}
|
|
|
|
// 构建目标 URL
|
|
targetURL, err := ph.buildTargetURL(r)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to build target URL: %v", err)
|
|
http.Error(w, "Bad Request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
log.Printf("[INFO] Target URL: %s", targetURL)
|
|
|
|
// 重试逻辑
|
|
maxRetries := len(ph.keyPool.keys)
|
|
attemptedKeys := make(map[string]bool)
|
|
log.Printf("[INFO] Starting key selection process, total keys available: %d", maxRetries)
|
|
|
|
for i := 0; i < maxRetries; i++ {
|
|
key := ph.getUnusedKey(attemptedKeys)
|
|
if key == "" {
|
|
log.Printf("[ERROR] No unused keys remaining after %d attempts", i)
|
|
break
|
|
}
|
|
attemptedKeys[key] = true
|
|
maskedKey := maskKey(key)
|
|
log.Printf("[INFO] Attempt %d/%d: Selecting key %s", i+1, maxRetries, maskedKey)
|
|
|
|
// 创建请求 - 为每次尝试创建新的请求体副本
|
|
req, err := ph.createRequest(r, targetURL, key, bodyBytes)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to create request with key %s: %v", maskedKey, err)
|
|
log.Printf("[INFO] Switching to another key due to request creation failure")
|
|
continue
|
|
}
|
|
|
|
// 发送请求
|
|
log.Printf("[INFO] Sending request to target API with key %s", maskedKey)
|
|
resp, err := ph.client.Do(req)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to send request with key %s: %v", maskedKey, err)
|
|
log.Printf("[INFO] Switching to another key due to network error")
|
|
continue
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// 处理响应
|
|
log.Printf("[INFO] Received response with status code %d", resp.StatusCode)
|
|
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
|
log.Println("[INFO] Request successful, forwarding response")
|
|
ph.forwardResponse(w, resp)
|
|
return
|
|
} else if resp.StatusCode == 403 || resp.StatusCode == 429 {
|
|
log.Printf("[WARN] Received %d status code with key %s", resp.StatusCode, maskedKey)
|
|
log.Printf("[INFO] Switching to another key due to status code %d", resp.StatusCode)
|
|
continue
|
|
} else {
|
|
log.Printf("[INFO] Forwarding response with status code %d", resp.StatusCode)
|
|
ph.forwardResponse(w, resp)
|
|
return
|
|
}
|
|
}
|
|
|
|
// 所有密钥尝试后仍失败
|
|
log.Printf("[ERROR] All %d keys failed after retries", maxRetries)
|
|
http.Error(w, "Failed to get response from API after all retries", http.StatusBadGateway)
|
|
}
|
|
|
|
// getUnusedKey 获取一个未使用过的密钥
|
|
func (ph *ProxyHandler) getUnusedKey(attempted map[string]bool) string {
|
|
key := ph.keyPool.GetRandomKey()
|
|
// 如果获取到的密钥已使用过，则尝试其他密钥
|
|
for attempted[key] && len(attempted) < len(ph.keyPool.keys) {
|
|
key = ph.keyPool.GetRandomKey()
|
|
}
|
|
// 如果所有密钥都已尝试过，返回空字符串
|
|
if attempted[key] {
|
|
return ""
|
|
}
|
|
return key
|
|
}
|
|
|
|
// authenticate 验证客户端身份
|
|
func (ph *ProxyHandler) authenticate(r *http.Request) bool {
|
|
authHeader := r.Header.Get("Authorization")
|
|
if authHeader == "" {
|
|
return false
|
|
}
|
|
parts := strings.Split(authHeader, " ")
|
|
if len(parts) != 2 || parts[0] != "Bearer" {
|
|
return false
|
|
}
|
|
return parts[1] == ph.cfg.Password
|
|
}
|
|
|
|
// buildTargetURL 构建目标 API 的完整 URL
|
|
func (ph *ProxyHandler) buildTargetURL(r *http.Request) (string, error) {
|
|
u, err := url.Parse(ph.cfg.TargetURL)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
u.Path = path.Join(u.Path, r.URL.Path)
|
|
u.RawQuery = r.URL.RawQuery
|
|
return u.String(), nil
|
|
}
|
|
|
|
// createRequest 创建转发请求 - 只修改Authorization头，为每次请求创建新的请求体
|
|
func (ph *ProxyHandler) createRequest(r *http.Request, targetURL, key string, bodyBytes []byte) (*http.Request, error) {
|
|
// 为每次请求创建新的请求体
|
|
var bodyReader io.Reader
|
|
if len(bodyBytes) > 0 {
|
|
bodyReader = bytes.NewReader(bodyBytes)
|
|
}
|
|
|
|
// 创建请求
|
|
req, err := http.NewRequest(r.Method, targetURL, bodyReader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 复制所有原始请求头
|
|
for k, v := range r.Header {
|
|
req.Header[k] = v
|
|
}
|
|
|
|
// 仅替换Authorization头
|
|
req.Header.Set("Authorization", "Bearer "+key)
|
|
return req, nil
|
|
}
|
|
|
|
// forwardResponse 将响应转发给客户端，支持流式和非流式
|
|
func (ph *ProxyHandler) forwardResponse(w http.ResponseWriter, resp *http.Response) {
|
|
// 设置响应头
|
|
for k, v := range resp.Header {
|
|
w.Header()[k] = v
|
|
}
|
|
w.WriteHeader(resp.StatusCode)
|
|
|
|
// 处理流式响应
|
|
if strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") || resp.Header.Get("Transfer-Encoding") == "chunked" {
|
|
log.Println("[INFO] Handling streaming response")
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
log.Println("[ERROR] Streaming unsupported by server")
|
|
http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
reader := bufio.NewReader(resp.Body)
|
|
for {
|
|
line, err := reader.ReadBytes('\n')
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
log.Println("[INFO] Stream ended")
|
|
break
|
|
}
|
|
log.Printf("[ERROR] Error reading stream: %v", err)
|
|
http.Error(w, "Error reading stream", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Write(line)
|
|
flusher.Flush()
|
|
}
|
|
} else {
|
|
// 非流式响应，直接复制
|
|
_, err := io.Copy(w, resp.Body)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to forward response: %v", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// maskKey 直接返回原始密钥，不再进行掩码处理
|
|
func maskKey(key string) string {
|
|
return key
|
|
}
|
|
|
|
// main 函数，启动代理服务器
|
|
func main() {
|
|
// 解析配置
|
|
cfg := parseFlags()
|
|
if cfg.KeyFile == "" || cfg.TargetURL == "" || cfg.Password == "" {
|
|
log.Println("[ERROR] Missing required flags: --key-file, --target-url, --password")
|
|
os.Exit(1)
|
|
}
|
|
log.Printf("[INFO] Configuration loaded: KeyFile=%s, TargetURL=%s, Address=%s, Port=%s, MaxWorkers=%d, MaxQueue=%d",
|
|
cfg.KeyFile, cfg.TargetURL, cfg.Address, cfg.Port, cfg.MaxWorkers, cfg.MaxQueue)
|
|
|
|
// 初始化密钥池
|
|
keyPool, err := NewKeyPool(cfg.KeyFile)
|
|
if err != nil {
|
|
log.Printf("[ERROR] Failed to initialize key pool: %v", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
// 创建代理处理器
|
|
proxyHandler := NewProxyHandler(cfg, keyPool)
|
|
|
|
// 初始化并启动工作池
|
|
proxyHandler.InitWorkerPool(cfg.MaxWorkers, cfg.MaxQueue)
|
|
|
|
// 启动服务器
|
|
addr := cfg.Address + ":" + cfg.Port
|
|
log.Printf("[INFO] Starting proxy server on %s", addr)
|
|
if err := http.ListenAndServe(addr, proxyHandler); err != nil {
|
|
log.Printf("[ERROR] Failed to start server: %v", err)
|
|
os.Exit(1)
|
|
}
|
|
} |