mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-06 23:56:37 +08:00
refactor(repo): consolidate openflare-server to root and move subprojects to internal/apps
- Merge all files inside openflare-server to the repository root directory. - Relocate agent, relay, and flared subprojects from internal/ to internal/apps/. - Combine docker-compose files and update build context paths to root. - Update GitHub workflows and Dockerfiles to refer to new directories and package names. - Rewrite Go package imports across all files. - Resolve database renew test race condition and clean up docs.
This commit is contained in:
ryan
@@ -0,0 +1,92 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package risk_control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/Rain-kl/Wavelet/internal/config"
|
||||
"github.com/Rain-kl/Wavelet/internal/model/analytics"
|
||||
analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics"
|
||||
"github.com/Rain-kl/Wavelet/pkg/logger"
|
||||
)
|
||||
|
||||
var logChan chan *analytics.UserAccessLog
|
||||
|
||||
const (
|
||||
defaultQueueSize = 10000
|
||||
maxBatchSize = 1000
|
||||
flushInterval = 1 * time.Second
|
||||
)
|
||||
|
||||
// InitLogWriter 初始化日志写入通道和后台写入协程
|
||||
func InitLogWriter(ctx context.Context) {
|
||||
if !config.Config.ClickHouse.Enabled {
|
||||
return
|
||||
}
|
||||
|
||||
logChan = make(chan *analytics.UserAccessLog, defaultQueueSize)
|
||||
go startBatchWorker(context.WithoutCancel(ctx))
|
||||
}
|
||||
|
||||
// IsBufferFull 检查当前本地缓冲队列是否已满
|
||||
// 如果没有启用 ClickHouse，默认返回 false，不触发限流
|
||||
func IsBufferFull() bool {
|
||||
if !config.Config.ClickHouse.Enabled || logChan == nil {
|
||||
return false
|
||||
}
|
||||
return len(logChan) >= cap(logChan)
|
||||
}
|
||||
|
||||
// QueueAccessLog 异步非阻塞地将日志推入缓冲队列
|
||||
func QueueAccessLog(logItem *analytics.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)
|
||||
}
|
||||
}
|
||||
|
||||
func startBatchWorker(ctx context.Context) {
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
var batch []*analytics.UserAccessLog
|
||||
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
items := make([]analytics.UserAccessLog, len(batch))
|
||||
for i, item := range batch {
|
||||
items[i] = *item
|
||||
}
|
||||
if err := analyticsrepo.BatchInsert(ctx, items); 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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// 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/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"
|
||||
)
|
||||
|
||||
// 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()
|
||||
|
||||
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 := &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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
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/model/analytics"
|
||||
"github.com/Rain-kl/Wavelet/internal/testhelper"
|
||||
"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 := testhelper.NewTestGinEngine(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 *analytics.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}
|
||||
oauth.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 *analytics.UserAccessLog, defaultQueueSize)
|
||||
defer func() {
|
||||
config.Config.ClickHouse.Enabled = false
|
||||
logChan = nil
|
||||
}()
|
||||
|
||||
r := testhelper.NewTestGinEngine(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 *analytics.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 <- &analytics.UserAccessLog{}
|
||||
}
|
||||
|
||||
r := testhelper.NewTestGinEngine(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"], "系统繁忙")
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user