fanqie-go/internal/handlers/common.go

package handlers

import (
	"crypto/subtle"
	"encoding/json"
	"fanqie/internal/config"
	"fanqie/internal/fqapi"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/time/rate"
)

var (
	startTime    = time.Now()
	deviceHeader string

	// 请求统计
	requestStats = newRequestStats()

	// 统计保存定时器
	statsSaveTicker *time.Ticker
)

// statsKinds 统计类别
var statsKinds = []string{"total", "search", "detail", "catalog", "content", "manga", "audio", "video", "category"}

// newRequestStats 初始化各类计数器
func newRequestStats() map[string]*atomic.Int64 {
	m := make(map[string]*atomic.Int64, len(statsKinds))
	for _, k := range statsKinds {
		m[k] = &atomic.Int64{}
	}
	return m
}

// incrementStat 原子递增请求计数
func incrementStat(kinds ...string) {
	for _, k := range kinds {
		requestStats[k].Add(1)
	}
}

// getStatsSnapshot 获取所有计数器的快照
func getStatsSnapshot() map[string]int64 {
	snapshot := make(map[string]int64, len(requestStats))
	for k, c := range requestStats {
		snapshot[k] = c.Load()
	}
	return snapshot
}

// SetDeviceHeader 设置设备header
func SetDeviceHeader(header string) {
	deviceHeader = header
}

// LoadStatistics 从数据库加载统计数据
func LoadStatistics() {
	stats, err := fqapi.LoadStatistics()
	if err != nil {
		fmt.Printf("加载统计数据失败: %v\n", err)
		return
	}

	for k, v := range stats {
		if c, ok := requestStats[k]; ok {
			c.Store(v)
		}
	}
}

// SaveStatistics 保存统计数据到数据库
func SaveStatistics() {
	if err := fqapi.SaveStatistics(getStatsSnapshot()); err != nil {
		fmt.Printf("保存统计数据失败: %v\n", err)
	}
}

// StartStatisticsSaver 启动统计数据定期保存
func StartStatisticsSaver() {
	// 每30秒保存一次统计数据
	statsSaveTicker = time.NewTicker(30 * time.Second)
	go func() {
		for range statsSaveTicker.C {
			SaveStatistics()
		}
	}()
}

// StopStatisticsSaver 停止统计数据定期保存
func StopStatisticsSaver() {
	if statsSaveTicker != nil {
		statsSaveTicker.Stop()
		// 最后保存一次
		SaveStatistics()
	}
}

// respondError 统一错误响应（JSON 安全编码；isClientErr 为 true 时返回 400，否则 500）
func respondError(w http.ResponseWriter, msg string, isClientErr bool) {
	status := http.StatusInternalServerError
	if isClientErr {
		status = http.StatusBadRequest
	}
	respondJSON(w, status, map[string]interface{}{"code": -1, "msg": msg})
}

// RateLimit 每 IP 限流（server.rate_limit_qps > 0 时启用）
// ponytail: 配置启动时读取一次，热更新不生效；visitor 表由后台协程清理防泄漏
func RateLimit(handler http.HandlerFunc) http.HandlerFunc {
	cfg := config.GetAppConfig()
	qps := 0.0
	if cfg != nil {
		qps = cfg.Server.RateLimitQPS
	}
	if qps <= 0 {
		return handler
	}

	const burst = 10
	var mu sync.Mutex
	limitters := map[string]*rate.Limiter{}
	lastSeen := map[string]time.Time{}

	go func() {
		for range time.Tick(time.Minute) {
			mu.Lock()
			for ip, t := range lastSeen {
				if time.Since(t) > 10*time.Minute {
					delete(lastSeen, ip)
					delete(limitters, ip)
				}
			}
			mu.Unlock()
		}
	}()

	return func(w http.ResponseWriter, r *http.Request) {
		ip := clientIP(r)
		mu.Lock()
		l, ok := limitters[ip]
		if !ok {
			l = rate.NewLimiter(rate.Limit(qps), burst)
			limitters[ip] = l
		}
		lastSeen[ip] = time.Now()
		mu.Unlock()

		if !l.Allow() {
			respondJSON(w, http.StatusTooManyRequests, map[string]interface{}{
				"code": -1, "msg": "请求过于频繁，请稍后再试",
			})
			return
		}
		handler(w, r)
	}
}

// clientIP 取客户端 IP：优先 X-Forwarded-For 首段（反代场景），否则 RemoteAddr
func clientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		return strings.TrimSpace(strings.Split(xff, ",")[0])
	}
	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return host
	}
	return r.RemoteAddr
}

// RequireAccessKey 验证访问密钥中间件
func RequireAccessKey(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cfg := config.GetConfig()

		// 如果未配置 access_key 或为空，则不启用验证
		if cfg == nil || cfg.Server.AccessKey == "" {
			handler(w, r)
			return
		}

		// 从查询参数或 Header 中获取 key
		key := r.URL.Query().Get("key")
		if key == "" {
			key = r.Header.Get("X-Access-Key")
		}

		// 验证密钥（常量时间比较，防时序攻击）
		if subtle.ConstantTimeCompare([]byte(key), []byte(cfg.Server.AccessKey)) != 1 {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusUnauthorized)
			fmt.Fprintf(w, `{"code": -1, "msg": "访问被拒绝：无效的访问密钥"}`)
			return
		}

		// 验证通过，继续处理请求
		handler(w, r)
	}
}

// OptionalAccessKey 业务接口保护（仅 server.protect_all 开启时校验）
func OptionalAccessKey(handler http.HandlerFunc) http.HandlerFunc {
	cfg := config.GetAppConfig()
	if cfg == nil || !cfg.Server.ProtectAll {
		return handler
	}
	return RequireAccessKey(handler)
}

// HandleHealth 存活探针（供容器/反向代理使用）
func HandleHealth(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "ok")
}

// HandleMain 主页处理
func HandleMain(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	response := map[string]interface{}{
		"code":    0,
		"message": "success",
		"data":    nil,
	}
	json.NewEncoder(w).Encode(response)
}

// HandleSelfInfo 设备信息处理
func HandleSelfInfo(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	// 获取设备池统计信息
	pool := fqapi.GetDeviceKeyPool()
	poolStats := pool.GetStats()

	response := map[string]interface{}{
		"code":    0,
		"message": "SUCCESS",
		"data": map[string]interface{}{
			"device_pool_enabled": pool.IsEnabled(),
			"device_pool_stats":   poolStats,
		},
	}
	json.NewEncoder(w).Encode(response)
}