Files
OpenFlare/internal/apps/risk_control/middleware.go
T

93 lines
2.6 KiB
Go

/*
Copyright 2026 Arctel.net
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
// Package risk_control 提供风险控制中间件
package risk_control
import (
"encoding/json"
"net/http"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"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/util"
"github.com/gin-gonic/gin"
)
// RiskControlMiddleware 全局日志采集中间件
func RiskControlMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
// 如果未启用 ClickHouse,直接放行
if !config.Config.ClickHouse.Enabled {
c.Next()
return
}
// 1. 限流背压检测(检测本地缓冲队列是否已满)
if IsBufferFull() {
c.AbortWithStatusJSON(http.StatusTooManyRequests, util.Err("系统繁忙,请稍后再试"))
return
}
start := time.Now()
// 2. 执行后续请求(穿过业务处理和认证中间件)
c.Next()
// 3. 后置身份检查:仅记录通过认证的请求
userObj, exists := util.GetFromContext[*model.User](c, oauth.UserObjKey)
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)
}
}
logItem := &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(c.Writer.Status()),
Latency: latency,
CreatedAt: time.Now(),
}
// 非阻塞地推入缓存队列
QueueAccessLog(logItem)
}
}