Upload 4 files

Dockerfile

@@ -0,0 +1,39 @@
+# --- 第一阶段：构建阶段 (Builder Stage) ---
+# 使用官方 Go 镜像进行编译
+FROM golang:1.22-alpine AS builder
+
+# 设置工作目录
+WORKDIR /build
+
+# 复制 Go 源代码文件
+COPY api-pool.go .
+
+# 编译 Go 应用
+# CGO_ENABLED=0 尝试静态链接，减少依赖
+# -ldflags="-w -s" 减小二进制文件大小
+# -o /app/api-pool 指定输出路径和名称
+RUN CGO_ENABLED=0 go build -ldflags="-w -s" -o /app/api-pool api-pool.go
+
+# --- 第二阶段：运行阶段 (Final Stage) ---
+# 使用轻量的 Alpine 镜像作为最终运行环境
+FROM alpine:latest
+
+# 设置工作目录
+WORKDIR /app
+
+# 从构建阶段复制编译好的二进制文件
+COPY --from=builder /app/api-pool /app/api-pool
+
+# 复制启动脚本
+COPY entrypoint.sh /app/entrypoint.sh
+
+# 赋予执行权限
+RUN chmod +x /app/api-pool /app/entrypoint.sh
+
+# 暴露应用程序监听的端口 (根据您的参数是 6969)
+EXPOSE 6969
+
+# 设置容器的入口点为启动脚本
+ENTRYPOINT ["/app/entrypoint.sh"]
+
+# 注意：CMD 指令现在由 entrypoint.sh 脚本通过 exec 来执行

README.md

@@ -1,11 +1,11 @@
----
-title: Api Pool
-emoji: 🏃
-colorFrom: pink
-colorTo: purple
-sdk: docker
-pinned: false
-license: gpl-3.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: API Key Pool # 您可以修改标题
+emoji: 🔑 # 您可以修改 Emoji
+colorFrom: green # 您可以修改颜色
+colorTo: blue # 您可以修改颜色
+sdk: docker
+app_port: 6969 # 必须与您的 --port 参数和 EXPOSE 端口一致
+pinned: false
+---
+
+(在此添加您的 Space 描述)

api-pool.go

@@ -0,0 +1,488 @@
+package main
+
+import (
+	"bufio"
+	"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)
+}
+
+// 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")
+
+	// 尝试解析请求体中的模型信息
+	model, err := ph.extractModelFromRequest(r)
+	if err != nil {
+		log.Printf("[WARN] Failed to extract model from request: %v", err)
+	} else if model != "" {
+		log.Printf("[INFO] Model specified in request: %s", model)
+	}
+
+	// 构建目标 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)
+		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 创建转发请求
+func (ph *ProxyHandler) createRequest(r *http.Request, targetURL, key string) (*http.Request, error) {
+	req, err := http.NewRequest(r.Method, targetURL, r.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	// 复制并修改请求头
+	for k, v := range r.Header {
+		if k != "Host" && k != "Connection" && k != "Proxy-Connection" && k != "Authorization" {
+			req.Header[k] = v
+		}
+	}
+	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)
+		}
+	}
+}
+
+// extractModelFromRequest 尝试从请求体中提取模型名称
+func (ph *ProxyHandler) extractModelFromRequest(r *http.Request) (string, error) {
+	if r.Body == nil {
+		return "", nil
+	}
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		return "", err
+	}
+	r.Body = io.NopCloser(strings.NewReader(string(body)))
+
+	var data map[string]interface{}
+	if err := json.Unmarshal(body, &data); err != nil {
+		return "", err
+	}
+	if model, ok := data["model"].(string); ok {
+		return model, nil
+	}
+	return "", nil
+}
+
+// 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)
+	}
+}

entrypoint.sh

@@ -0,0 +1,56 @@
+#!/bin/sh
+# 如果命令失败则立即退出
+set -e
+
+# 定义临时密钥文件的路径
+KEY_FILE_PATH="/tmp/keys.txt"
+
+echo "--- 正在检查 Secrets ---"
+
+# 检查 API_PASSWORD 是否已设置
+if [ -z "${API_PASSWORD}" ]; then
+  echo "[错误] 必须在 Hugging Face Secrets 中设置 API_PASSWORD !"
+  exit 1
+fi
+
+# 检查 key_list 是否已设置
+if [ -z "${key_list}" ]; then
+  echo "[错误] 必须在 Hugging Face Secrets 中设置 key_list !"
+  exit 1
+fi
+
+echo "--- 正在从 Secret 'key_list' 创建临时密钥文件 (${KEY_FILE_PATH}) ---"
+
+# 从环境变量 key_list 读取内容，并写入临时文件
+# 使用 'echo -e' 来解释可能存在的 '\n' 换行符
+# 将标准错误重定向到 /dev/null 以避免打印潜在的密钥内容（尽管通常 echo 不会）
+echo -e "${key_list}" > "${KEY_FILE_PATH}" 2>/dev/null
+
+# 验证文件是否创建成功且非空
+if [ ! -s "${KEY_FILE_PATH}" ]; then
+    echo "[错误] 创建密钥文件失败或文件为空！请检查 'key_list' Secret 的内容。"
+    exit 1
+fi
+
+echo "--- 密钥文件已生成 ---"
+
+# !!! 【重要】生产环境中不要取消下面这行的注释，避免日志泄露 !!!
+# echo "密钥文件内容预览 (前几行):"
+# head -n 3 "${KEY_FILE_PATH}"
+
+echo "--- 正在启动 api-pool 服务 ---"
+
+# 使用 exec 执行 Go 程序，将脚本进程替换为 Go 程序进程
+# 将临时文件路径传给 --key-file
+# 将从 Secret 读取的密码传给 --password
+# 将地址设为 0.0.0.0 以便容器外访问
+# 传入其他您指定的参数
+exec /app/api-pool \
+    --key-file "${KEY_FILE_PATH}" \
+    --target-url "https://api.siliconflow.cn" \
+    --port "6969" \
+    --address "0.0.0.0" \
+    --password "${API_PASSWORD}" \
+    --max-workers=1000 \
+    --max-queue=2000
+    # 注意：--max-workers 和 --max-queue 值较高，请关注 Space 资源使用情况