package internal import ( "bufio" "encoding/json" "fmt" "io" "net/http" "os" fhttp "github.com/bogdanfinn/fhttp" "path/filepath" "strings" "sync" "sync/atomic" "time" "github.com/fsnotify/fsnotify" ) // TokenInfo 存储单个 token 的信息 type TokenInfo struct { Token string `json:"token"` Email string `json:"email"` UserID string `json:"user_id"` Valid bool `json:"valid"` LastChecked time.Time `json:"last_checked"` UseCount int64 `json:"use_count"` } // TokenManager 管理所有用户 token type TokenManager struct { mu sync.RWMutex tokens map[string]*TokenInfo // token -> TokenInfo validTokens []string // 有效 token 列表 currentIndex int // 轮询索引 dataDir string watcher *fsnotify.Watcher checkInterval time.Duration stopChan chan struct{} multimodalCount int64 // 多模态请求计数 totalCalls int64 // 累计调用次数 successCalls int64 // 成功调用次数 } var ( tokenManager *TokenManager tokenOnce sync.Once ) // GetTokenManager 获取单例 TokenManager func GetTokenManager() *TokenManager { tokenOnce.Do(func() { tokenManager = &TokenManager{ tokens: make(map[string]*TokenInfo), validTokens: make([]string, 0), dataDir: "data", checkInterval: 5 * time.Minute, // 每5分钟检查一次 stopChan: make(chan struct{}), } }) return tokenManager } // Start 启动 token 管理器 func (tm *TokenManager) Start() error { // 确保 data 目录存在 if err := os.MkdirAll(tm.dataDir, 0755); err != nil { return fmt.Errorf("创建 data 目录失败: %v", err) } // 初始加载 token if err := tm.loadTokens(); err != nil { LogWarn("初始加载 token 失败: %v", err) } // 启动文件监听 if err := tm.startWatcher(); err != nil { LogWarn("启动文件监听失败: %v", err) } // 启动定期验证 go tm.startValidator() LogInfo("TokenManager 已启动,当前有效 token 数: %d", len(tm.validTokens)) return nil } // Stop 停止 token 管理器 func (tm *TokenManager) Stop() { close(tm.stopChan) if tm.watcher != nil { tm.watcher.Close() } } // loadTokens 从 data 目录加载所有 token func (tm *TokenManager) loadTokens() error { tokenFile := filepath.Join(tm.dataDir, "tokens.txt") file, err := os.Open(tokenFile) if err != nil { if os.IsNotExist(err) { // 创建示例文件 tm.createExampleTokenFile(tokenFile) return nil } return err } defer file.Close() tm.mu.Lock() defer tm.mu.Unlock() // 保留旧的统计数据 oldTokens := tm.tokens tm.tokens = make(map[string]*TokenInfo) tm.validTokens = make([]string, 0) scanner := bufio.NewScanner(file) for scanner.Scan() { line := strings.TrimSpace(scanner.Text()) if line == "" || strings.HasPrefix(line, "#") { continue } token := line // 支持 token=xxx 格式 if strings.HasPrefix(line, "token=") { token = strings.TrimPrefix(line, "token=") } if token == "" { continue } // 复用旧的 TokenInfo 如果存在 if oldInfo, exists := oldTokens[token]; exists { tm.tokens[token] = oldInfo if oldInfo.Valid { tm.validTokens = append(tm.validTokens, token) } } else { // 新 token,解析并标记为待验证 info := &TokenInfo{ Token: token, Valid: true, // 初始假设有效,验证时会更新 } // 尝试解析 JWT 获取信息 if payload, err := DecodeJWTPayload(token); err == nil && payload != nil { info.Email = payload.Email info.UserID = payload.ID } tm.tokens[token] = info tm.validTokens = append(tm.validTokens, token) } } validN := len(tm.validTokens) LogInfo("已加载 %d 个 token", validN) scanErr := scanner.Err() if validN > 0 { invalidateAnonymousPoolSlots() } return scanErr } // createExampleTokenFile 创建示例 token 文件 func (tm *TokenManager) createExampleTokenFile(path string) { content := `# 用户 Token 文件 # 每行一个 token,支持以下格式: # 1. 直接写 token # 2. token=xxx 格式 # 以 # 开头的行为注释 # 示例: # eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.xxxxx ` os.WriteFile(path, []byte(content), 0644) LogInfo("已创建示例 token 文件: %s", path) } // startWatcher 启动文件变化监听 func (tm *TokenManager) startWatcher() error { watcher, err := fsnotify.NewWatcher() if err != nil { return err } tm.watcher = watcher go func() { for { select { case event, ok := <-watcher.Events: if !ok { return } if event.Op&(fsnotify.Write|fsnotify.Create) != 0 { if strings.HasSuffix(event.Name, "tokens.txt") { LogInfo("检测到 token 文件变化,重新加载...") time.Sleep(100 * time.Millisecond) // 等待文件写入完成 tm.loadTokens() } } case err, ok := <-watcher.Errors: if !ok { return } LogError("文件监听错误: %v", err) case <-tm.stopChan: return } } }() return watcher.Add(tm.dataDir) } // startValidator 启动定期验证 func (tm *TokenManager) startValidator() { // 首次延迟验证 time.Sleep(10 * time.Second) tm.validateAllTokens() ticker := time.NewTicker(tm.checkInterval) defer ticker.Stop() for { select { case <-ticker.C: tm.validateAllTokens() case <-tm.stopChan: return } } } // validateAllTokens 验证所有 token func (tm *TokenManager) validateAllTokens() { tm.mu.RLock() tokens := make([]string, 0, len(tm.tokens)) for token := range tm.tokens { tokens = append(tokens, token) } tm.mu.RUnlock() LogInfo("开始验证 %d 个 token...", len(tokens)) invalidCount := 0 for _, token := range tokens { valid := tm.validateToken(token) tm.mu.Lock() if info, exists := tm.tokens[token]; exists { info.Valid = valid info.LastChecked = time.Now() if !valid { invalidCount++ } } tm.mu.Unlock() time.Sleep(500 * time.Millisecond) // 避免请求过快 } // 更新有效 token 列表 tm.rebuildValidTokens() LogInfo("Token 验证完成,失效 %d 个,剩余有效 %d 个", invalidCount, len(tm.validTokens)) // 自动删除失效 token if invalidCount > 0 { tm.removeInvalidTokens() } } // validateToken 验证单个 token func (tm *TokenManager) validateToken(token string) bool { req, err := fhttp.NewRequest("GET", "https://chat.z.ai/api/v1/auths/", nil) if err != nil { return false } ApplyBrowserFingerprintHeaders(req.Header) req.Header.Set("Accept", "application/json") req.Header.Set("Accept-Language", "zh-CN") req.Header.Set("Cache-Control", "no-cache") req.Header.Set("Content-Type", "application/json") req.Header.Set("DNT", "1") req.Header.Set("Pragma", "no-cache") req.Header.Set("Referer", "https://chat.z.ai/") req.Header.Set("Sec-Fetch-Dest", "empty") req.Header.Set("Sec-Fetch-Mode", "cors") req.Header.Set("Sec-Fetch-Site", "same-origin") req.Header.Set("sec-gpc", "1") req.AddCookie(&fhttp.Cookie{Name: "token", Value: token}) client, err := TLSHTTPClient(10 * time.Second) if err != nil { LogDebug("Token 验证 tls client: %v", err) return false } resp, err := client.Do(req) if err != nil { LogDebug("Token 验证请求失败: %v", err) return false } defer resp.Body.Close() // 读取响应 body, _ := io.ReadAll(resp.Body) if resp.StatusCode != http.StatusOK { LogDebug("Token 验证失败,状态码: %d", resp.StatusCode) return false } // 尝试解析响应获取新 token var authResp struct { Token string `json:"token"` Email string `json:"email"` ID string `json:"id"` } if err := json.Unmarshal(body, &authResp); err == nil && authResp.Token != "" { // 更新 token 信息 tm.mu.Lock() if info, exists := tm.tokens[token]; exists { if authResp.Email != "" { info.Email = authResp.Email } if authResp.ID != "" { info.UserID = authResp.ID } } tm.mu.Unlock() } return true } // rebuildValidTokens 重建有效 token 列表 func (tm *TokenManager) rebuildValidTokens() { tm.mu.Lock() defer tm.mu.Unlock() tm.validTokens = make([]string, 0) for token, info := range tm.tokens { if info.Valid { tm.validTokens = append(tm.validTokens, token) } } } // removeInvalidTokens 从文件中移除失效 token func (tm *TokenManager) removeInvalidTokens() { tm.mu.Lock() defer tm.mu.Unlock() tokenFile := filepath.Join(tm.dataDir, "tokens.txt") invalidFile := filepath.Join(tm.dataDir, "tokens_invalid.txt") // 收集失效 token var invalidTokens []string for token, info := range tm.tokens { if !info.Valid { invalidTokens = append(invalidTokens, token) delete(tm.tokens, token) } } if len(invalidTokens) == 0 { return } // 追加到失效文件 f, err := os.OpenFile(invalidFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) if err == nil { defer f.Close() timestamp := time.Now().Format("2006-01-02 15:04:05") for _, token := range invalidTokens { f.WriteString(fmt.Sprintf("# 失效于 %s\n%s\n", timestamp, token)) } } // 重写有效 token 文件 var validTokenLines []string for _, token := range tm.validTokens { validTokenLines = append(validTokenLines, token) } content := "# 用户 Token 文件(自动更新)\n" content += fmt.Sprintf("# 更新时间: %s\n", time.Now().Format("2006-01-02 15:04:05")) content += "# 失效 token 已移至 tokens_invalid.txt\n\n" content += strings.Join(validTokenLines, "\n") if len(validTokenLines) > 0 { content += "\n" } os.WriteFile(tokenFile, []byte(content), 0644) LogInfo("已移除 %d 个失效 token 到 %s", len(invalidTokens), invalidFile) } // HasValidUpstreamTokens 是否存在可用的 z.ai 上游 token(TokenManager 轮询用) func (tm *TokenManager) HasValidUpstreamTokens() bool { tm.mu.RLock() defer tm.mu.RUnlock() return len(tm.validTokens) > 0 } // GetToken 获取一个有效 token(轮询) func (tm *TokenManager) GetToken() string { tm.mu.Lock() defer tm.mu.Unlock() if len(tm.validTokens) == 0 { return "" } token := tm.validTokens[tm.currentIndex%len(tm.validTokens)] tm.currentIndex++ // 增加使用计数 if info, exists := tm.tokens[token]; exists { info.UseCount++ } return token } // RecordCall 记录调用 func (tm *TokenManager) RecordCall(success bool, isMultimodal bool) { atomic.AddInt64(&tm.totalCalls, 1) if success { atomic.AddInt64(&tm.successCalls, 1) } if isMultimodal { atomic.AddInt64(&tm.multimodalCount, 1) } } // GetStats 获取统计数据 func (tm *TokenManager) GetStats() TokenManagerStats { tm.mu.RLock() defer tm.mu.RUnlock() total := atomic.LoadInt64(&tm.totalCalls) success := atomic.LoadInt64(&tm.successCalls) multimodal := atomic.LoadInt64(&tm.multimodalCount) var successRate float64 if total > 0 { successRate = float64(success) / float64(total) * 100 } return TokenManagerStats{ ValidTokenCount: len(tm.validTokens), TotalTokenCount: len(tm.tokens), MultimodalCount: multimodal, TotalCalls: total, SuccessCalls: success, SuccessRate: successRate, } } // TokenManagerStats token 管理器统计数据 type TokenManagerStats struct { ValidTokenCount int `json:"valid_token_count"` TotalTokenCount int `json:"total_token_count"` MultimodalCount int64 `json:"multimodal_count"` TotalCalls int64 `json:"total_calls"` SuccessCalls int64 `json:"success_calls"` SuccessRate float64 `json:"success_rate"` } // GetClientIP 从请求中获取客户端 IP func GetClientIP(r *http.Request) string { // 优先从 X-Forwarded-For 获取 xff := r.Header.Get("X-Forwarded-For") if xff != "" { // X-Forwarded-For 可能包含多个 IP,取第一个 ips := strings.Split(xff, ",") if len(ips) > 0 { ip := strings.TrimSpace(ips[0]) if ip != "" { return ip } } } // 尝试 X-Real-IP xri := r.Header.Get("X-Real-IP") if xri != "" { return xri } // 最后使用 RemoteAddr ip := r.RemoteAddr // 去除端口 if idx := strings.LastIndex(ip, ":"); idx != -1 { ip = ip[:idx] } return ip }