OpenFlare/backend/plugins/domain/risk_control/middleware.go

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

// Package risk_control provides the access control, IP rate limiting, and telemetry risk analysis domain plugin for Cordis.
package risk_control

import (
	"Wavelet/core/contracts"
	"Wavelet/pkg/ginutil"
	"Wavelet/pkg/idgen"
	"Wavelet/pkg/response"
	"Wavelet/plugins/domain/risk_control/logstore"
	"encoding/json"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
)

var accessLogEnabled atomic.Bool

// SetAccessLogEnabled enables or disables access log collection.
func SetAccessLogEnabled(enabled bool) {
	accessLogEnabled.Store(enabled)
}

// IsAccessLogEnabled reports whether access log collection is enabled.
func IsAccessLogEnabled() bool {
	return accessLogEnabled.Load()
}

// Middleware is an alias for RiskControlMiddleware.
var Middleware = RiskControlMiddleware

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

		// 1. 限流背压检测（检测本地缓冲队列是否已满）
		if IsBufferFull() {
			response.AbortTooManyRequests(c, errSystemBusy)
			return
		}

		start := time.Now()

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

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

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

		var headersStr string
		if c.Request.Header != nil {
			// 克隆 Header，避免污染原 HTTP 请求的 Header 对象
			clonedHeaders := make(http.Header)
			for k, v := range c.Request.Header {
				clonedHeaders[k] = v
			}
			clonedHeaders.Del("Cookie")

			if headersBytes, err := json.Marshal(clonedHeaders); err == nil {
				headersStr = string(headersBytes)
			}
		}

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

		logItem := &logstore.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)
	}
}