mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-07 08:06:37 +08:00
clickhouse 日志采集
This commit is contained in:
ryan
@@ -0,0 +1,91 @@
|
||||
/*
|
||||
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
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
/*
|
||||
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
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
|
||||
"github.com/Rain-kl/Wavelet/internal/config"
|
||||
"github.com/Rain-kl/Wavelet/internal/model"
|
||||
"github.com/Rain-kl/Wavelet/internal/util"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestRiskControlMiddleware(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("ClickHouse disabled", func(t *testing.T) {
|
||||
config.Config.ClickHouse.Enabled = false
|
||||
defer func() { config.Config.ClickHouse.Enabled = false }()
|
||||
|
||||
r := gin.New()
|
||||
r.Use(RiskControlMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
assert.Equal(t, "ok", w.Body.String())
|
||||
})
|
||||
|
||||
t.Run("ClickHouse enabled - Normal Authenticated Request", func(t *testing.T) {
|
||||
config.Config.ClickHouse.Enabled = true
|
||||
logChan = make(chan *UserAccessLog, defaultQueueSize)
|
||||
defer func() {
|
||||
config.Config.ClickHouse.Enabled = false
|
||||
logChan = nil
|
||||
}()
|
||||
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
// Mock authentication middleware placing user in context
|
||||
user := &model.User{ID: 12345}
|
||||
util.SetToContext(c, oauth.UserObjKey, user)
|
||||
c.Next()
|
||||
})
|
||||
r.Use(RiskControlMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.Header.Set("X-Test-Header", "hello")
|
||||
req.Header.Set("Cookie", "session_id=abcdef123456")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
assert.Equal(t, "ok", w.Body.String())
|
||||
|
||||
// Verify log is enqueued
|
||||
select {
|
||||
case logItem := <-logChan:
|
||||
assert.Equal(t, uint64(12345), logItem.UserID)
|
||||
assert.Equal(t, "/test", logItem.Path)
|
||||
assert.Equal(t, http.MethodGet, logItem.Method)
|
||||
assert.Equal(t, int32(http.StatusOK), logItem.Status)
|
||||
assert.NotEmpty(t, logItem.Headers)
|
||||
assert.Contains(t, logItem.Headers, "X-Test-Header")
|
||||
assert.NotContains(t, logItem.Headers, "Cookie")
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("expected log item in logChan, but got none")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ClickHouse enabled - Unauthenticated Request", func(t *testing.T) {
|
||||
config.Config.ClickHouse.Enabled = true
|
||||
logChan = make(chan *UserAccessLog, defaultQueueSize)
|
||||
defer func() {
|
||||
config.Config.ClickHouse.Enabled = false
|
||||
logChan = nil
|
||||
}()
|
||||
|
||||
r := gin.New()
|
||||
r.Use(RiskControlMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
assert.Equal(t, "ok", w.Body.String())
|
||||
|
||||
// Verify no log is enqueued
|
||||
select {
|
||||
case <-logChan:
|
||||
t.Fatal("expected no log item for unauthenticated request")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
// Success
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ClickHouse enabled - Buffer Full Rate Limiting", func(t *testing.T) {
|
||||
config.Config.ClickHouse.Enabled = true
|
||||
logChan = make(chan *UserAccessLog, 2) // small capacity for quick fill
|
||||
defer func() {
|
||||
config.Config.ClickHouse.Enabled = false
|
||||
logChan = nil
|
||||
}()
|
||||
|
||||
// fill logChan up to cap to simulate buffer full
|
||||
for len(logChan) < cap(logChan) {
|
||||
logChan <- &UserAccessLog{}
|
||||
}
|
||||
|
||||
r := gin.New()
|
||||
r.Use(RiskControlMiddleware())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, "ok")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req, _ := http.NewRequest(http.MethodGet, "/test", nil)
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusTooManyRequests, w.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
err := json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, resp["error_msg"], "系统繁忙")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
/*
|
||||
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
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/config"
|
||||
"github.com/Rain-kl/Wavelet/internal/db"
|
||||
"github.com/Rain-kl/Wavelet/internal/logger"
|
||||
)
|
||||
|
||||
// UserAccessLog 用户访问记录
|
||||
type UserAccessLog struct {
|
||||
ID uint64 `json:"id,string"`
|
||||
UserID uint64 `json:"user_id,string"`
|
||||
Path string `json:"path"`
|
||||
Method string `json:"method"`
|
||||
IP string `json:"ip"`
|
||||
UserAgent string `json:"user_agent"`
|
||||
Headers string `json:"headers"`
|
||||
Status int32 `json:"status"`
|
||||
Latency int64 `json:"latency"` // 耗时毫秒
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
var (
|
||||
logChan chan *UserAccessLog
|
||||
)
|
||||
|
||||
const (
|
||||
defaultQueueSize = 10000
|
||||
maxBatchSize = 1000
|
||||
flushInterval = 1 * time.Second
|
||||
)
|
||||
|
||||
// InitLogWriter 初始化日志写入通道和后台写入协程
|
||||
func InitLogWriter() {
|
||||
if !config.Config.ClickHouse.Enabled {
|
||||
return
|
||||
}
|
||||
|
||||
logChan = make(chan *UserAccessLog, defaultQueueSize)
|
||||
go startBatchWorker()
|
||||
}
|
||||
|
||||
// IsBufferFull 检查当前本地缓冲队列是否已满
|
||||
// 如果没有启用 ClickHouse，默认返回 false，不触发限流
|
||||
func IsBufferFull() bool {
|
||||
if !config.Config.ClickHouse.Enabled || logChan == nil {
|
||||
return false
|
||||
}
|
||||
return len(logChan) >= cap(logChan)
|
||||
}
|
||||
|
||||
// QueueAccessLog 异步非阻塞地将日志推入缓冲队列
|
||||
func QueueAccessLog(logItem *UserAccessLog) {
|
||||
if !config.Config.ClickHouse.Enabled || logChan == nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case logChan <- logItem:
|
||||
default:
|
||||
// 如果在极端并发下仍然写满了，这里做非阻塞丢弃，防止卡死
|
||||
logger.WarnF(context.Background(), "[RiskControl] Log queue full, dropping log item for path: %s", logItem.Path)
|
||||
}
|
||||
}
|
||||
|
||||
// startBatchWorker 后台批量写入 ClickHouse 的工作协程
|
||||
func startBatchWorker() {
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
var batch []*UserAccessLog
|
||||
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
if db.ChConn == nil {
|
||||
batch = nil
|
||||
return
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
b, err := db.ChConn.PrepareBatch(ctx, "INSERT INTO user_access_logs (id, user_id, path, method, ip, user_agent, headers, status, latency, created_at)")
|
||||
if err != nil {
|
||||
logger.ErrorF(ctx, "[RiskControl] Prepare ClickHouse batch failed: %v", err)
|
||||
batch = nil
|
||||
return
|
||||
}
|
||||
|
||||
for _, item := range batch {
|
||||
err = b.Append(
|
||||
item.ID,
|
||||
item.UserID,
|
||||
item.Path,
|
||||
item.Method,
|
||||
item.IP,
|
||||
item.UserAgent,
|
||||
item.Headers,
|
||||
item.Status,
|
||||
item.Latency,
|
||||
item.CreatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
logger.ErrorF(ctx, "[RiskControl] Append item to ClickHouse batch failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := b.Send(); err != nil {
|
||||
logger.ErrorF(ctx, "[RiskControl] Send ClickHouse batch failed: %v", err)
|
||||
}
|
||||
batch = nil
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case item, ok := <-logChan:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
batch = append(batch, item)
|
||||
if len(batch) >= maxBatchSize {
|
||||
flush()
|
||||
}
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user