OpenFlare/internal/apps/risk_control/middleware.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

// Package risk_control 提供风险控制中间件
package risk_control

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"github.com/Rain-kl/Wavelet/internal/apps/oauth"
	"github.com/Rain-kl/Wavelet/internal/common/response"
	"github.com/Rain-kl/Wavelet/internal/config"
	"github.com/Rain-kl/Wavelet/internal/db/idgen"
	"github.com/Rain-kl/Wavelet/internal/model"
	"github.com/Rain-kl/Wavelet/internal/model/analytics"
	"github.com/gin-gonic/gin"
)

const maxAuditLogHeadersBytes = 2 * 1024

var auditLogHeaderAllowlist = map[string]struct{}{
	"Authorization":   {},
	"Cookie":          {},
	"X-Forwarded-For": {},
	"X-Real-Ip":       {},
	"User-Agent":      {},
	"Content-Type":    {},
}

func marshalAuditLogHeaders(headers http.Header) string {
	if headers == nil {
		return ""
	}

	filtered := make(http.Header)
	for key, values := range headers {
		if _, ok := auditLogHeaderAllowlist[key]; !ok {
			continue
		}
		filtered[key] = redactAuditLogHeaderValues(key, values)
	}

	headersBytes, err := json.Marshal(filtered)
	if err != nil {
		return ""
	}
	if len(headersBytes) <= maxAuditLogHeadersBytes {
		return string(headersBytes)
	}
	return string(headersBytes[:maxAuditLogHeadersBytes])
}

func redactAuditLogHeaderValues(key string, values []string) []string {
	switch key {
	case "Authorization", "Cookie":
		redacted := make([]string, len(values))
		for i, value := range values {
			redacted[i] = hashAuditLogSensitiveValue(value)
		}
		return redacted
	default:
		return values
	}
}

func hashAuditLogSensitiveValue(value string) string {
	if strings.TrimSpace(value) == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(value))
	return "sha256:" + hex.EncodeToString(sum[:8])
}

// RiskControlMiddleware 全局日志采集中间件
func RiskControlMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 如果未启用 ClickHouse，直接放行
		if !config.Config.ClickHouse.Enabled {
			c.Next()
			return
		}

		// 1. 限流背压检测（检测本地缓冲队列是否已满）
		if IsBufferFull() {
			response.AbortTooManyRequests(c, "系统繁忙，请稍后再试")
			return
		}

		start := time.Now()

		// 2. 执行后续请求（穿过业务处理和认证中间件）
		c.Next()

		// 3. 后置身份检查：仅记录通过认证的请求
		userObj, exists := oauth.GetFromContext[*model.User](c, oauth.UserObjKey)
		if !exists || userObj == nil {
			return
		}

		// 4. 计算耗时并异步推送到缓冲队列
		latency := time.Since(start).Milliseconds()

		headersStr := marshalAuditLogHeaders(c.Request.Header)

		const maxHTTPStatus = 999
		status := c.Writer.Status()
		if status < 0 {
			status = 0
		} else if status > maxHTTPStatus {
			status = maxHTTPStatus
		}

		logItem := &analytics.UserAccessLog{
			ID:        idgen.NextUint64ID(),
			UserID:    userObj.ID, // 直接从 Context 获取已登录用户ID，避免数据库查询
			Path:      c.Request.URL.Path,
			Method:    c.Request.Method,
			IP:        c.ClientIP(),
			UserAgent: c.Request.UserAgent(),
			Headers:   headersStr,
			Status:    int32(status),
			Latency:   latency,
			CreatedAt: time.Now(),
		}

		// 非阻塞地推入缓存队列
		QueueAccessLog(logItem)
	}
}