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()
	}
}

// getQueryWithAlias 读 URL 参数，主名未传时回退到别名（兼容 README 历史参数名）
func getQueryWithAlias(r *http.Request, name, alias string) string {
	if v := r.URL.Query().Get(name); v != "" {
		return v
	}
	return r.URL.Query().Get(alias)
}

// 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)
}

// statusWriter 捕获状态码；status>=400 时截断留存响应体前 ErrMsgMax 字节用于排障
type statusWriter struct {
	http.ResponseWriter
	status int
	errBuf []byte
}

func (w *statusWriter) WriteHeader(code int) {
	w.status = code
	w.ResponseWriter.WriteHeader(code)
}

func (w *statusWriter) Write(p []byte) (int, error) {
	if w.status >= 400 && len(w.errBuf) < fqapi.ErrMsgMax {
		n := min(len(p), fqapi.ErrMsgMax-len(w.errBuf))
		w.errBuf = append(w.errBuf, p[:n]...)
	}
	return w.ResponseWriter.Write(p)
}

// LogRequestMiddleware 异步记录请求元数据（时间/IP/方法/路径/状态/耗时）到数据库
// ponytail: 配置启动时读取一次，热更新不生效；未启用时零开销直通
func LogRequestMiddleware(handler http.HandlerFunc) http.HandlerFunc {
	cfg := config.GetAppConfig()
	if cfg == nil || !cfg.Server.RequestLog.Enabled {
		return handler
	}
	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		sw := &statusWriter{ResponseWriter: w, status: http.StatusOK}
		handler(sw, r)

		ua := r.UserAgent()
		if len(ua) > 512 { // 防恶意超长 UA 撑爆日志
			ua = ua[:512]
		}
		fqapi.LogRequest(fqapi.RequestLogEntry{
			Ts:         start.Unix(),
			IP:         clientIP(r),
			Method:     r.Method,
			Path:       r.URL.Path,
			Query:      decodedQuery(r),
			Status:     sw.status,
			DurationMs: time.Since(start).Milliseconds(),
			ErrorMsg:   string(sw.errBuf),
			UserAgent:  ua,
		})
	}
}

// decodedQuery 重组解码后的查询串（keyword=%E6%B5%8B → keyword=测），日志可读
// ponytail: map 遍历乱序、%26 解码后与分隔符同名 —— 展示用途可接受
func decodedQuery(r *http.Request) string {
	vals := r.URL.Query() // 已解码
	if len(vals) == 0 {
		return ""
	}
	parts := make([]string, 0, len(vals))
	for k, vs := range vals {
		for _, v := range vs {
			parts = append(parts, k+"="+v)
		}
	}
	return strings.Join(parts, "&")
}

// 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)
}