init
This commit is contained in:
tbphp
423
internal/proxy/proxy.go
Normal file
423
internal/proxy/proxy.go
Normal 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()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user