Compare commits
No commits in common. "b5c3c0954ff5e595e9e65db8e746fbb4393eecfb" and "39ff086236184147a0b58ff4b825c02ad4abcf00" have entirely different histories.
b5c3c0954f
...
39ff086236
@ -46,18 +46,6 @@ jobs:
|
||||
echo "=== Unit Tests ==="
|
||||
go test -race -count=1 ./internal/...
|
||||
|
||||
- name: Install golangci-lint
|
||||
run: |
|
||||
curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin latest
|
||||
|
||||
- name: Lint
|
||||
run: |
|
||||
export PATH="$(go env GOPATH)/bin:$PATH"
|
||||
golangci-lint run ./...
|
||||
|
||||
- name: Integration Tests
|
||||
run: go test -race -tags=integration -count=1 ./internal/integration/...
|
||||
|
||||
build:
|
||||
name: Build
|
||||
needs: [test]
|
||||
|
||||
5
Makefile
5
Makefile
@ -128,12 +128,11 @@ build-all: build-linux build-darwin build-windows build-freebsd build-openbsd
|
||||
|
||||
run:
|
||||
@echo "Running $(APP_NAME) in development mode..."
|
||||
go run $(MAIN_PATH) -c lolly.yaml
|
||||
go run $(MAIN_PATH) -c configs/lolly.yaml
|
||||
|
||||
test-config:
|
||||
@echo "Testing configuration..."
|
||||
go run $(MAIN_PATH) --generate-config -o /tmp/lolly-generated.yaml
|
||||
@echo "Config syntax OK (generated template validated)"
|
||||
go run $(MAIN_PATH) -t -c configs/lolly.yaml
|
||||
|
||||
version:
|
||||
@echo "$(APP_NAME) version $(VERSION)"
|
||||
|
||||
@ -27,13 +27,12 @@ services:
|
||||
# reservations:
|
||||
# memory: 256M
|
||||
restart: unless-stopped
|
||||
# Scratch 镜像没有 wget/curl,如需健康检查请使用外部 HTTP 探针或挂载包含工具的 sidecar。
|
||||
# healthcheck:
|
||||
# test: ["CMD", "/lolly", "--healthcheck"]
|
||||
# interval: 30s
|
||||
# timeout: 10s
|
||||
# retries: 3
|
||||
# start_period: 10s
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "wget -q --spider http://localhost:8080/_status || exit 1"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
start_period: 10s
|
||||
|
||||
# 示例后端服务(用于反向代理测试)
|
||||
# backend:
|
||||
|
||||
@ -219,7 +219,7 @@ func BenchmarkStreamTCPForward(b *testing.B) {
|
||||
// 设置 upstream 健康
|
||||
srv.SetHealthy("test", 0, true)
|
||||
|
||||
_ = srv.ListenTCP("127.0.0.1:0", "test")
|
||||
_ = srv.ListenTCP("127.0.0.1:0")
|
||||
_ = srv.Start()
|
||||
defer srv.Stop()
|
||||
|
||||
|
||||
1
go.mod
1
go.mod
@ -19,7 +19,6 @@ require (
|
||||
github.com/yuin/gopher-lua v1.1.2
|
||||
golang.org/x/crypto v0.50.0
|
||||
golang.org/x/net v0.53.0
|
||||
golang.org/x/sync v0.21.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
|
||||
2
go.sum
2
go.sum
@ -154,8 +154,6 @@ golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI=
|
||||
golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q=
|
||||
golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA=
|
||||
golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs=
|
||||
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
|
||||
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
|
||||
@ -17,7 +17,6 @@
|
||||
package adapter
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
@ -51,9 +50,6 @@ type CommonAdapter struct {
|
||||
// CtxPool 用于复用 fasthttp.RequestCtx 对象
|
||||
// 每个协议适配器实例独立维护自己的 ctxPool
|
||||
CtxPool sync.Pool
|
||||
|
||||
// MaxBodySize 限制允许读取的请求体最大字节数(<=0 表示不限制)
|
||||
MaxBodySize int64
|
||||
}
|
||||
|
||||
// NewCommonAdapter 创建新的共享适配器实例。
|
||||
@ -94,52 +90,33 @@ func (a *CommonAdapter) ResetContext(ctx *fasthttp.RequestCtx) {
|
||||
//
|
||||
// 对于小于等于 DefaultBodyThreshold(64KB)的请求体,直接读取到内存;
|
||||
// 对于大于阈值的请求体,使用共享 bufferPool 进行流式处理,避免内存峰值。
|
||||
// 超过 MaxBodySize 的请求体会被拒绝并返回 413 错误。
|
||||
//
|
||||
// 参数:
|
||||
// - r: 标准库的 HTTP 请求
|
||||
// - ctx: fasthttp 请求上下文,用于存储读取的请求体
|
||||
//
|
||||
// 返回值:
|
||||
// - error: 读取或限制失败时返回错误
|
||||
func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.RequestCtx) error {
|
||||
func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.RequestCtx) {
|
||||
if r.Body == nil || r.Body == http.NoBody {
|
||||
return nil
|
||||
return
|
||||
}
|
||||
|
||||
defer func() {
|
||||
_ = r.Body.Close()
|
||||
}()
|
||||
|
||||
limit := a.MaxBodySize
|
||||
if limit <= 0 {
|
||||
limit = 1 << 20 // 1MB default when unset
|
||||
}
|
||||
|
||||
// Reject early via Content-Length when possible.
|
||||
if r.ContentLength > 0 && r.ContentLength > limit {
|
||||
ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge)
|
||||
return fmt.Errorf("request body %d exceeds limit %d", r.ContentLength, limit)
|
||||
}
|
||||
|
||||
// Stream with a hard cap so chunked/unknown-length bodies cannot OOM us.
|
||||
limitedBody := io.LimitReader(r.Body, limit+1)
|
||||
|
||||
// 小请求体:直接读取到内存(<= 64KB)
|
||||
if r.ContentLength > 0 && r.ContentLength <= DefaultBodyThreshold {
|
||||
body, err := io.ReadAll(limitedBody)
|
||||
if err != nil {
|
||||
return err
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err == nil {
|
||||
ctx.Request.SetBody(body)
|
||||
}
|
||||
if int64(len(body)) > limit {
|
||||
ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge)
|
||||
return fmt.Errorf("request body exceeds limit %d", limit)
|
||||
}
|
||||
ctx.Request.SetBody(body)
|
||||
return nil
|
||||
return
|
||||
}
|
||||
|
||||
// 大请求体:使用流式缓冲区(> 64KB 或未知长度)
|
||||
// 从全局 pool 获取缓冲区
|
||||
bufPtr, ok := bufferPoolInstance.Get().(*[]byte)
|
||||
if !ok {
|
||||
// 如果类型断言失败,创建新的缓冲区(不应该发生)
|
||||
buf := make([]byte, 4096)
|
||||
bufPtr = &buf
|
||||
}
|
||||
@ -147,33 +124,29 @@ func (a *CommonAdapter) StreamRequestBody(r *http.Request, ctx *fasthttp.Request
|
||||
|
||||
buf := *bufPtr
|
||||
var body []byte
|
||||
|
||||
// 如果已知 ContentLength,预分配精确大小的缓冲区
|
||||
if r.ContentLength > 0 {
|
||||
body = make([]byte, 0, r.ContentLength)
|
||||
}
|
||||
|
||||
var total int64
|
||||
// 分块读取请求体
|
||||
for {
|
||||
n, err := limitedBody.Read(buf)
|
||||
n, err := r.Body.Read(buf)
|
||||
if n > 0 {
|
||||
body = append(body, buf[:n]...)
|
||||
total += int64(n)
|
||||
}
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if total > limit {
|
||||
ctx.Error("Request Entity Too Large", fasthttp.StatusRequestEntityTooLarge)
|
||||
return fmt.Errorf("request body exceeds limit %d", limit)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(body) > 0 {
|
||||
ctx.Request.SetBody(body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetContext 从 pool 获取一个 fasthttp.RequestCtx。
|
||||
|
||||
@ -35,7 +35,7 @@ func TestNewCommonAdapter(t *testing.T) {
|
||||
a := NewCommonAdapter()
|
||||
|
||||
require.NotNil(t, a, "返回值不应为 nil")
|
||||
require.NotNil(t, a.CtxPool.New, "CtxPool.New 应该被初始化")
|
||||
require.NotNil(t, a.CtxPool, "CtxPool 应该被初始化")
|
||||
|
||||
// 验证 pool 能创建 *fasthttp.RequestCtx
|
||||
obj := a.CtxPool.Get()
|
||||
@ -162,8 +162,7 @@ func TestStreamRequestBody(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
require.NoError(t, err, "StreamRequestBody 不应返回错误")
|
||||
a.StreamRequestBody(r, ctx)
|
||||
|
||||
if tt.wantBodySet {
|
||||
assert.Equal(t, tt.wantBody, string(ctx.Request.Body()), "请求体内容应匹配")
|
||||
@ -186,9 +185,8 @@ func TestStreamRequestBody_ReadError(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
// 读取错误应被返回
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
require.Error(t, err, "读取错误应返回错误")
|
||||
// 不应该 panic
|
||||
a.StreamRequestBody(r, ctx)
|
||||
}
|
||||
|
||||
// errorReader 是一个始终返回错误的 io.Reader
|
||||
@ -212,11 +210,10 @@ func TestStreamRequestBody_PartialReadError(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
require.Error(t, err, "部分读取后出错应返回错误")
|
||||
a.StreamRequestBody(r, ctx)
|
||||
|
||||
// 出错后不再保留已读取的部分数据
|
||||
assert.Equal(t, 0, len(ctx.Request.Body()), "出错后请求体应为空")
|
||||
// 部分读取的数据应该保留
|
||||
assert.Equal(t, "partial", string(ctx.Request.Body()), "已读取的部分数据应保留")
|
||||
}
|
||||
|
||||
// partialErrorReader 先返回数据,再返回错误
|
||||
@ -339,8 +336,7 @@ func TestConcurrentStreamRequestBody(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
require.NoError(t, err)
|
||||
a.StreamRequestBody(r, ctx)
|
||||
assert.Equal(t, "concurrent body data", string(ctx.Request.Body()))
|
||||
}()
|
||||
}
|
||||
@ -359,8 +355,7 @@ func TestStreamRequestBody_ClosesBody(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
require.NoError(t, err)
|
||||
a.StreamRequestBody(r, ctx)
|
||||
|
||||
assert.True(t, closeTracker.closed, "请求体应该被关闭")
|
||||
}
|
||||
@ -379,24 +374,3 @@ func (r *trackableReader) Close() error {
|
||||
r.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestStreamRequestBodyEnforcesMaxBodySize 测试请求体超过最大限制时被拒绝。
|
||||
func TestStreamRequestBodyEnforcesMaxBodySize(t *testing.T) {
|
||||
a := NewCommonAdapter()
|
||||
a.MaxBodySize = 10
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
body := strings.NewReader("this body is way too large")
|
||||
r := &http.Request{
|
||||
Body: io.NopCloser(body),
|
||||
ContentLength: int64(body.Len()),
|
||||
}
|
||||
|
||||
err := a.StreamRequestBody(r, ctx)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for oversized body")
|
||||
}
|
||||
if ctx.Response.StatusCode() != fasthttp.StatusRequestEntityTooLarge {
|
||||
t.Fatalf("expected 413, got %d", ctx.Response.StatusCode())
|
||||
}
|
||||
}
|
||||
|
||||
@ -16,7 +16,6 @@ import (
|
||||
"rua.plus/lolly/internal/http2"
|
||||
"rua.plus/lolly/internal/http3"
|
||||
"rua.plus/lolly/internal/logging"
|
||||
"rua.plus/lolly/internal/middleware/bodylimit"
|
||||
"rua.plus/lolly/internal/resolver"
|
||||
"rua.plus/lolly/internal/server"
|
||||
"rua.plus/lolly/internal/stream"
|
||||
@ -161,7 +160,7 @@ func (a *App) initStreamServers() {
|
||||
a.logger.Error().Err(err).Str("listen", sc.Listen).Msg("Failed to listen on UDP")
|
||||
}
|
||||
} else {
|
||||
if err := a.streamSrv.ListenTCP(sc.Listen, sc.Listen); err != nil {
|
||||
if err := a.streamSrv.ListenTCP(sc.Listen); err != nil {
|
||||
a.logger.Error().Err(err).Str("listen", sc.Listen).Msg("Failed to listen on TCP")
|
||||
}
|
||||
}
|
||||
@ -175,15 +174,6 @@ func (a *App) initStreamServers() {
|
||||
}()
|
||||
}
|
||||
|
||||
// clientMaxBodySize parses the first server's client_max_body_size into bytes.
|
||||
func (a *App) clientMaxBodySize() int64 {
|
||||
size, err := bodylimit.ParseSize(a.cfg.Servers[0].ClientMaxBodySize)
|
||||
if err != nil {
|
||||
return bodylimit.DefaultMaxBodySize
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
// initHTTP3 starts the HTTP/3 server if enabled.
|
||||
func (a *App) initHTTP3() {
|
||||
if len(a.cfg.Servers) == 0 || !a.cfg.HTTP3.Enabled || a.cfg.Servers[0].SSL.Cert == "" {
|
||||
@ -196,8 +186,6 @@ func (a *App) initHTTP3() {
|
||||
return
|
||||
}
|
||||
|
||||
a.cfg.HTTP3.MaxBodySize = a.clientMaxBodySize()
|
||||
|
||||
a.http3Srv, err = http3.NewServer(&a.cfg.HTTP3, a.srv.GetHandler(), tlsConfig)
|
||||
if err != nil {
|
||||
a.logger.Error().Err(err).Msg("Failed to create HTTP/3 server")
|
||||
@ -224,8 +212,6 @@ func (a *App) initHTTP2() {
|
||||
return
|
||||
}
|
||||
|
||||
a.cfg.Servers[0].SSL.HTTP2.MaxBodySize = a.clientMaxBodySize()
|
||||
|
||||
a.http2Srv, err = http2.NewServer(&a.cfg.Servers[0].SSL.HTTP2, a.srv.GetHandler(), tlsConfig)
|
||||
if err != nil {
|
||||
a.logger.Error().Err(err).Msg("Failed to create HTTP/2 server")
|
||||
|
||||
@ -23,8 +23,6 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"rua.plus/lolly/internal/config"
|
||||
"rua.plus/lolly/internal/logging"
|
||||
"rua.plus/lolly/internal/version"
|
||||
)
|
||||
|
||||
@ -896,36 +894,6 @@ logging:
|
||||
app.srv.StopWithTimeout(1 * time.Second)
|
||||
}
|
||||
|
||||
// TestInitHTTP3_EmptyServers 验证空服务器配置时 HTTP/3 初始化直接跳过。
|
||||
func TestInitHTTP3_EmptyServers(t *testing.T) {
|
||||
app := &App{
|
||||
cfg: &config.Config{Servers: []config.ServerConfig{}},
|
||||
logger: logging.NewAppLogger(nil),
|
||||
}
|
||||
|
||||
// 不应 panic
|
||||
app.initHTTP3()
|
||||
|
||||
if app.http3Srv != nil {
|
||||
t.Error("http3Srv should remain nil when no servers are configured")
|
||||
}
|
||||
}
|
||||
|
||||
// TestInitHTTP2_EmptyServers 验证空服务器配置时 HTTP/2 初始化直接跳过。
|
||||
func TestInitHTTP2_EmptyServers(t *testing.T) {
|
||||
app := &App{
|
||||
cfg: &config.Config{Servers: []config.ServerConfig{}},
|
||||
logger: logging.NewAppLogger(nil),
|
||||
}
|
||||
|
||||
// 不应 panic
|
||||
app.initHTTP2()
|
||||
|
||||
if app.http2Srv != nil {
|
||||
t.Error("http2Srv should remain nil when no servers are configured")
|
||||
}
|
||||
}
|
||||
|
||||
func waitForAppServerRunning(app *App, timeout time.Duration) bool {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
|
||||
11
internal/cache/file_cache.go
vendored
11
internal/cache/file_cache.go
vendored
@ -148,17 +148,6 @@ func (c *FileCache) Get(path string) (*FileEntry, bool) {
|
||||
return entry, true
|
||||
}
|
||||
|
||||
// 更新近似 LRU 访问时间(每 5 秒或超过 10% inactive 间隔才写,减少锁竞争)
|
||||
const accessUpdateInterval = 5 * time.Second
|
||||
if time.Since(entry.LastAccess) > accessUpdateInterval {
|
||||
c.mu.Lock()
|
||||
if e, ok := c.entries[path]; ok && e == entry {
|
||||
e.LastAccess = time.Now()
|
||||
c.lruList.MoveToFront(e.element)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// 近似 LRU: 不更新 LastAccess/LRU 位置,直接返回
|
||||
// 精确 LRU 由 Set/Delete/evict 操作保证
|
||||
return entry, true
|
||||
|
||||
24
internal/cache/file_cache_test.go
vendored
24
internal/cache/file_cache_test.go
vendored
@ -1,24 +0,0 @@
|
||||
// Package cache 提供文件缓存功能测试。
|
||||
package cache
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestFileCacheLastAccessUpdatedOnHit(t *testing.T) {
|
||||
c := NewFileCache(1, 1<<20, 10*time.Second)
|
||||
c.Set("/x", []byte("data"), 4, time.Now(), "text/plain")
|
||||
|
||||
entry := c.entries["/x"]
|
||||
before := entry.LastAccess
|
||||
// 将访问时间设为过去,确保超过 accessUpdateInterval 阈值但不超过 inactive
|
||||
entry.LastAccess = before.Add(-7 * time.Second)
|
||||
|
||||
c.Get("/x")
|
||||
after := c.entries["/x"].LastAccess
|
||||
|
||||
if !after.After(before) {
|
||||
t.Fatal("LastAccess was not updated on cache hit")
|
||||
}
|
||||
}
|
||||
@ -36,7 +36,6 @@ type ProxyCacheConfig struct {
|
||||
CacheLockTimeout time.Duration `yaml:"cache_lock_timeout"` // 缓存锁超时时间
|
||||
BackgroundUpdateDisable bool `yaml:"background_update_disable"` // 禁用后台更新(默认 false = 启用后台更新)
|
||||
CacheIgnoreHeaders []string `yaml:"cache_ignore_headers"` // 缓存时忽略的响应头
|
||||
Vary []string `yaml:"vary"` // 参与缓存键的 Vary 请求头
|
||||
Revalidate bool `yaml:"revalidate"` // 启用条件请求(If-Modified-Since/If-None-Match)
|
||||
}
|
||||
|
||||
|
||||
@ -40,7 +40,6 @@ type HTTP2Config struct {
|
||||
PushEnabled bool `yaml:"push_enabled"`
|
||||
H2CEnabled bool `yaml:"h2c_enabled"`
|
||||
GracefulShutdownTimeout time.Duration `yaml:"graceful_shutdown_timeout"`
|
||||
MaxBodySize int64 `yaml:"max_body_size,omitempty"`
|
||||
}
|
||||
|
||||
// HTTP3Config HTTP/3 (QUIC) 配置。
|
||||
@ -68,7 +67,6 @@ type HTTP3Config struct {
|
||||
IdleTimeout time.Duration `yaml:"idle_timeout"`
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Enable0RTT bool `yaml:"enable_0rtt"`
|
||||
MaxBodySize int64 `yaml:"max_body_size,omitempty"`
|
||||
}
|
||||
|
||||
// PerformanceConfig 性能配置。
|
||||
|
||||
@ -945,66 +945,3 @@ func TestSendFile_NegativeLength(t *testing.T) {
|
||||
t.Errorf("Expected body %s, got %s", content, ctx.Response.Body())
|
||||
}
|
||||
}
|
||||
|
||||
// TestLinuxSendfile_DataIntegrity 验证 sendfile 真正传输了完整文件内容,即 socket FD 在使用前未被关闭。
|
||||
func TestLinuxSendfile_DataIntegrity(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
tmpFile := filepath.Join(tmpDir, "integrity.bin")
|
||||
|
||||
content := make([]byte, 16*1024)
|
||||
for i := range content {
|
||||
content[i] = byte(i % 256)
|
||||
}
|
||||
if err := os.WriteFile(tmpFile, content, 0o644); err != nil {
|
||||
t.Fatalf("Failed to create temp file: %v", err)
|
||||
}
|
||||
|
||||
file, err := os.Open(tmpFile)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to open file: %v", err)
|
||||
}
|
||||
defer func() { _ = file.Close() }()
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to listen: %v", err)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
received := make(chan []byte, 1)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
c, err := ln.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
_ = c.SetReadDeadline(time.Now().Add(5 * time.Second))
|
||||
buf, _ := io.ReadAll(c)
|
||||
received <- buf
|
||||
}()
|
||||
|
||||
clientConn, err := net.Dial("tcp", ln.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to dial: %v", err)
|
||||
}
|
||||
defer clientConn.Close()
|
||||
|
||||
_ = clientConn.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
||||
|
||||
err = linuxSendfile(clientConn, file.Fd(), 0, int64(len(content)))
|
||||
if err != nil && err != syscall.EPIPE && err != syscall.ECONNRESET {
|
||||
t.Fatalf("linuxSendfile failed: %v", err)
|
||||
}
|
||||
|
||||
_ = clientConn.Close()
|
||||
wg.Wait()
|
||||
close(received)
|
||||
|
||||
buf := <-received
|
||||
if !bytes.Equal(buf, content) {
|
||||
t.Errorf("sendfile data integrity mismatch: sent %d bytes, received %d bytes", len(content), len(buf))
|
||||
}
|
||||
}
|
||||
|
||||
@ -359,10 +359,8 @@ func (h *StaticHandler) Handle(ctx *fasthttp.RequestCtx) {
|
||||
return
|
||||
}
|
||||
|
||||
// 安全检查:解析后的绝对路径必须位于 root/alias 目录下
|
||||
requestedPath := string(ctx.Path())
|
||||
rawPath := string(ctx.Request.URI().PathOriginal())
|
||||
if !h.isPathSafe(requestedPath, rawPath) {
|
||||
// 安全检查:防止目录遍历
|
||||
if strings.Contains(reqPath, "..") {
|
||||
utils.SendError(ctx, utils.ErrForbidden)
|
||||
return
|
||||
}
|
||||
@ -377,54 +375,6 @@ func (h *StaticHandler) Handle(ctx *fasthttp.RequestCtx) {
|
||||
h.handleStandard(ctx, reqPath)
|
||||
}
|
||||
|
||||
// isPathSafe 检查请求路径解析后是否始终位于 root/alias 目录范围内。
|
||||
// 它同时检查规范化后的路径和原始编码路径,防止通过 URL 编码的 ..
|
||||
// 或经过 fasthttp 规范化后超出目录边界。
|
||||
func (h *StaticHandler) isPathSafe(reqPath, rawPath string) bool {
|
||||
// 快速拒绝原始路径中的遍历片段(在 URL 解码/规范化之前)。
|
||||
// 特别拦截编码的 traversal,如 /..%2fsecret.txt。
|
||||
if strings.Contains(rawPath, "..") && strings.Contains(rawPath, "%") {
|
||||
return false
|
||||
}
|
||||
|
||||
// 拒绝规范化路径中仍包含 .. 的字符串(如 file..txt)。
|
||||
if strings.Contains(reqPath, "..") {
|
||||
return false
|
||||
}
|
||||
|
||||
base := h.root
|
||||
if h.alias != "" {
|
||||
base = h.alias
|
||||
}
|
||||
if base == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
absRoot, err := filepath.Abs(base)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
// 使用与后续文件服务一致的相对路径计算方式,并单独清理相对部分,
|
||||
// 避免 filepath.Clean 删除对前缀剥离有意义的尾部斜杠。
|
||||
relPath := h.stripPathPrefix(reqPath)
|
||||
if relPath == "" {
|
||||
relPath = "."
|
||||
}
|
||||
target := filepath.Join(absRoot, filepath.Clean(relPath))
|
||||
absTarget, err := filepath.Abs(target)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
sep := string(filepath.Separator)
|
||||
if !strings.HasPrefix(absTarget, absRoot+sep) && absTarget != absRoot {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// tryServeFromFileCache 尝试从文件缓存直接响应。
|
||||
// 命中缓存且文件未修改时直接写入响应并返回 true。
|
||||
func (h *StaticHandler) tryServeFromFileCache(ctx *fasthttp.RequestCtx, filePath string, info os.FileInfo) bool {
|
||||
|
||||
@ -20,7 +20,6 @@ package handler
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -2053,49 +2052,6 @@ func TestSetCacheHeaders(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestStaticHandler_ConcurrentSetterAndHandle 验证 setter 与 Handle 并发访问无数据竞争。
|
||||
func TestStaticHandler_ConcurrentSetterAndHandle(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(tmpDir, "test.txt"), []byte("hello"), 0o644); err != nil {
|
||||
t.Fatalf("创建测试文件失败: %v", err)
|
||||
}
|
||||
|
||||
handler := newTestHandler(t, tmpDir)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
// 并发执行 setter
|
||||
for i := range 10 {
|
||||
wg.Add(1)
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
switch id % 5 {
|
||||
case 0:
|
||||
handler.SetExpires("1h")
|
||||
case 1:
|
||||
handler.SetCacheTTL(time.Second)
|
||||
case 2:
|
||||
handler.SetFileCache(nil)
|
||||
case 3:
|
||||
handler.SetGzipStatic(false, nil, nil)
|
||||
case 4:
|
||||
handler.SetTryFiles([]string{"$uri"}, false, nil)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
// 并发执行 Handle
|
||||
for range 20 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ctx := newTestContext(t, "/test.txt")
|
||||
handler.Handle(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// TestParseExpires 测试 parseExpires 函数。
|
||||
func TestParseExpires(t *testing.T) {
|
||||
tests := []struct {
|
||||
@ -2125,22 +2081,3 @@ func TestParseExpires(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestStaticHandlerPathTraversalEncoded 测试 URL 编码的路径遍历被拦截。
|
||||
func TestStaticHandlerPathTraversalEncoded(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(root, "secret.txt"), []byte("secret"), 0o600); err != nil {
|
||||
t.Fatalf("创建测试文件失败: %v", err)
|
||||
}
|
||||
|
||||
handler := NewStaticHandler(root, "/", []string{"index.html", "index.htm"}, false)
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Request.Header.SetMethod("GET")
|
||||
ctx.Request.SetRequestURI("/..%2fsecret.txt")
|
||||
handler.Handle(ctx)
|
||||
|
||||
if ctx.Response.StatusCode() != fasthttp.StatusForbidden {
|
||||
t.Fatalf("expected 403, got %d", ctx.Response.StatusCode())
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,73 +0,0 @@
|
||||
package hash
|
||||
|
||||
import (
|
||||
"hash/fnv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFNV64a(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
want uint64
|
||||
}{
|
||||
{"", 0xcbf29ce484222325},
|
||||
{"hello", 0xa430d84680aabd0b},
|
||||
{"a", 0xaf63dc4c8601ec8c},
|
||||
{"The quick brown fox jumps over the lazy dog", 0xf3f9b7f5e7e47110},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
got := FNV64a(tc.input)
|
||||
if got != tc.want {
|
||||
t.Errorf("FNV64a(%q) = %#x, want %#x", tc.input, got, tc.want)
|
||||
}
|
||||
|
||||
// Cross-check with stdlib FNV-1a implementation.
|
||||
h := fnv.New64a()
|
||||
_, _ = h.Write([]byte(tc.input))
|
||||
if got != h.Sum64() {
|
||||
t.Errorf("FNV64a(%q) = %#x, stdlib FNV-1a = %#x", tc.input, got, h.Sum64())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFNV64aBytes(t *testing.T) {
|
||||
cases := []struct {
|
||||
input []byte
|
||||
want uint64
|
||||
}{
|
||||
{[]byte{}, 0xcbf29ce484222325},
|
||||
{[]byte("hello"), 0xa430d84680aabd0b},
|
||||
{[]byte{0x00, 0x01, 0xff}, 0xd94a37186c0d38bf},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
got := FNV64aBytes(tc.input)
|
||||
if got != tc.want {
|
||||
t.Errorf("FNV64aBytes(%q) = %#x, want %#x", tc.input, got, tc.want)
|
||||
}
|
||||
|
||||
h := fnv.New64a()
|
||||
_, _ = h.Write(tc.input)
|
||||
if got != h.Sum64() {
|
||||
t.Errorf("FNV64aBytes(%q) = %#x, stdlib FNV-1a = %#x", tc.input, got, h.Sum64())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFNV64aAndBytesConsistency(t *testing.T) {
|
||||
inputs := []string{
|
||||
"",
|
||||
"hello",
|
||||
"lolly",
|
||||
"The quick brown fox jumps over the lazy dog",
|
||||
}
|
||||
|
||||
for _, s := range inputs {
|
||||
fromString := FNV64a(s)
|
||||
fromBytes := FNV64aBytes([]byte(s))
|
||||
if fromString != fromBytes {
|
||||
t.Errorf("FNV64a(%q) = %#x, FNV64aBytes(%q) = %#x", s, fromString, s, fromBytes)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -63,11 +63,7 @@ func (a *FastHTTPHandlerAdapter) ServeHTTP(w http.ResponseWriter, r *http.Reques
|
||||
a.convertRequest(r, ctx)
|
||||
|
||||
// 流式处理请求体
|
||||
if err := a.StreamRequestBody(r, ctx); err != nil {
|
||||
// ctx already contains the 413 response; write it back.
|
||||
a.convertResponse(ctx, w)
|
||||
return
|
||||
}
|
||||
a.StreamRequestBody(r, ctx)
|
||||
|
||||
// 调用 fasthttp handler
|
||||
a.handler(ctx)
|
||||
|
||||
@ -42,7 +42,6 @@ type Server struct {
|
||||
stopChan chan struct{}
|
||||
connWg sync.WaitGroup
|
||||
GracefulShutdownTimeout time.Duration
|
||||
maxBodySize int64
|
||||
mu sync.RWMutex
|
||||
running bool
|
||||
}
|
||||
@ -105,7 +104,6 @@ func NewServer(cfg *config.HTTP2Config, handler fasthttp.RequestHandler, tlsConf
|
||||
handler: handler,
|
||||
pool: newConnectionPool(),
|
||||
GracefulShutdownTimeout: gracefulTimeout,
|
||||
maxBodySize: cfg.MaxBodySize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@ -213,7 +211,6 @@ func (s *Server) handleConnection(conn net.Conn) {
|
||||
// serveHTTP2 使用 HTTP/2 协议服务连接。
|
||||
func (s *Server) serveHTTP2(conn net.Conn) {
|
||||
adapter := NewFastHTTPHandlerAdapter(s.handler)
|
||||
adapter.MaxBodySize = s.maxBodySize
|
||||
|
||||
opts := &http2.ServeConnOpts{
|
||||
Context: context.Background(),
|
||||
|
||||
@ -67,11 +67,7 @@ func (a *Adapter) Wrap(handler fasthttp.RequestHandler) http.Handler {
|
||||
a.ResetContext(ctx)
|
||||
|
||||
// 转换请求
|
||||
if err := a.convertRequest(r, ctx); err != nil {
|
||||
// Write a 413 response directly.
|
||||
http.Error(w, "Request Entity Too Large", http.StatusRequestEntityTooLarge)
|
||||
return
|
||||
}
|
||||
a.convertRequest(r, ctx)
|
||||
|
||||
// 设置 ResponseWriter 用于后续写入
|
||||
ctx.SetUserValue("http3_response_writer", w)
|
||||
@ -89,10 +85,7 @@ func (a *Adapter) Wrap(handler fasthttp.RequestHandler) http.Handler {
|
||||
// 参数:
|
||||
// - r: 标准库 HTTP 请求
|
||||
// - ctx: FastHTTP 请求上下文
|
||||
//
|
||||
// 返回值:
|
||||
// - error: 请求体超过限制时返回错误
|
||||
func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) error {
|
||||
func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) {
|
||||
// 设置方法
|
||||
ctx.Request.Header.SetMethod(r.Method)
|
||||
|
||||
@ -114,9 +107,7 @@ func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) erro
|
||||
}
|
||||
|
||||
// 设置请求体(使用流式处理优化)
|
||||
if err := a.StreamRequestBody(r, ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
a.StreamRequestBody(r, ctx)
|
||||
|
||||
// 设置远程地址
|
||||
if r.RemoteAddr != "" {
|
||||
@ -127,7 +118,6 @@ func (a *Adapter) convertRequest(r *http.Request, ctx *fasthttp.RequestCtx) erro
|
||||
|
||||
// 设置协议版本
|
||||
ctx.Request.Header.SetProtocol("HTTP/3")
|
||||
return nil
|
||||
}
|
||||
|
||||
// convertResponse 将 fasthttp.RequestCtx 响应写入 http.ResponseWriter。
|
||||
|
||||
@ -59,9 +59,6 @@ type Server struct {
|
||||
|
||||
// mu 读写锁
|
||||
mu sync.RWMutex
|
||||
|
||||
// maxBodySize 请求体大小限制
|
||||
maxBodySize int64
|
||||
}
|
||||
|
||||
// NewServer 创建 HTTP/3 服务器。
|
||||
@ -88,14 +85,12 @@ func NewServer(cfg *config.HTTP3Config, handler fasthttp.RequestHandler, tlsConf
|
||||
}
|
||||
|
||||
adapter := NewAdapter()
|
||||
adapter.MaxBodySize = cfg.MaxBodySize
|
||||
|
||||
return &Server{
|
||||
config: cfg,
|
||||
handler: handler,
|
||||
adapter: adapter,
|
||||
tlsConfig: tlsConfig,
|
||||
maxBodySize: cfg.MaxBodySize,
|
||||
config: cfg,
|
||||
handler: handler,
|
||||
adapter: adapter,
|
||||
tlsConfig: tlsConfig,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
//go:build integration
|
||||
|
||||
// Package integration 提供端到端集成基准测试。
|
||||
//
|
||||
// 该文件测试完整请求路径的吞吐量,涵盖静态文件、代理转发、
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
//go:build integration
|
||||
|
||||
// Package integration 提供后端故障切换 E2E 基准测试。
|
||||
//
|
||||
// 该文件测试负载均衡器剔除/恢复后端的开销。
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
//go:build integration
|
||||
|
||||
package integration
|
||||
|
||||
import (
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
//go:build integration
|
||||
|
||||
// resolver_integration_test.go - DNS 解析器集成测试
|
||||
//
|
||||
// 测试 DNS 解析器与代理的集成
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
//go:build integration
|
||||
|
||||
// variable_integration_test.go - 变量系统集成测试
|
||||
//
|
||||
// 测试变量系统与日志、代理、重写的端到端集成
|
||||
@ -210,11 +208,9 @@ func TestVariablePerformance(t *testing.T) {
|
||||
avg := elapsed / time.Duration(iterations)
|
||||
t.Logf("Average expansion time: %v (iterations: %d)", avg, iterations)
|
||||
|
||||
// 验证性能在合理范围内。
|
||||
// race 检测会显著放大执行时间,因此阈值放宽到 50μs。
|
||||
threshold := 50 * time.Microsecond
|
||||
if avg > threshold {
|
||||
t.Errorf("average time %v exceeds %v", avg, threshold)
|
||||
// 验证性能在合理范围内(< 1μs 每次)
|
||||
if avg > time.Microsecond {
|
||||
t.Errorf("average time %v exceeds 1μs", avg)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -89,32 +89,19 @@ func (m *SlowStartManager) OnTargetUnhealthy(target *Target) {
|
||||
// Start 启动后台权重更新。
|
||||
//
|
||||
// 定期遍历所有慢启动中的目标,计算并更新 EffectiveWeight。
|
||||
// 支持 Start-Stop-Start 周期:每次启动都会重新创建 stopCh,
|
||||
// 避免 Stop 关闭通道后 updateLoop 立即退出。
|
||||
func (m *SlowStartManager) Start() {
|
||||
if m.running.Swap(true) {
|
||||
return // 已经在运行
|
||||
}
|
||||
|
||||
// 重新创建 stopCh,确保新一轮更新循环使用未关闭的通道
|
||||
m.mu.Lock()
|
||||
m.stopCh = make(chan struct{})
|
||||
stopCh := m.stopCh
|
||||
m.mu.Unlock()
|
||||
go m.updateLoop(stopCh)
|
||||
}
|
||||
// 重建 stopCh 以支持 Start-Stop-Start 周期
|
||||
select {
|
||||
case <-m.stopCh:
|
||||
m.stopCh = make(chan struct{})
|
||||
default:
|
||||
}
|
||||
|
||||
// SetFindTarget 设置目标查找回调。
|
||||
//
|
||||
// 慢启动管理器通过该回调在内部状态与负载均衡目标之间建立映射,
|
||||
// 用于更新 EffectiveWeight。必须在 Start 之前设置。
|
||||
//
|
||||
// 参数:
|
||||
// - findTarget: 根据 target URL 查找对应 Target 的回调函数
|
||||
func (m *SlowStartManager) SetFindTarget(findTarget func(url string) *Target) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.findTarget = findTarget
|
||||
go m.updateLoop()
|
||||
}
|
||||
|
||||
// Stop 停止后台更新。
|
||||
@ -131,7 +118,7 @@ func (m *SlowStartManager) Stop() {
|
||||
}
|
||||
|
||||
// updateLoop 后台更新循环。
|
||||
func (m *SlowStartManager) updateLoop(stopCh <-chan struct{}) {
|
||||
func (m *SlowStartManager) updateLoop() {
|
||||
ticker := time.NewTicker(m.interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
@ -139,7 +126,7 @@ func (m *SlowStartManager) updateLoop(stopCh <-chan struct{}) {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
m.updateEffectiveWeights()
|
||||
case <-stopCh:
|
||||
case <-m.stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,67 +1,6 @@
|
||||
package loadbalance
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSlowStartManager_StartStopStart(t *testing.T) {
|
||||
m := NewSlowStartManager(10 * time.Millisecond)
|
||||
|
||||
// 第一次启动
|
||||
m.Start()
|
||||
if !m.running.Load() {
|
||||
t.Fatal("manager should be running after first Start")
|
||||
}
|
||||
|
||||
// 停止
|
||||
m.Stop()
|
||||
if m.running.Load() {
|
||||
t.Fatal("manager should be stopped after Stop")
|
||||
}
|
||||
|
||||
// 再次启动:验证 stopCh 被正确重建,updateLoop 不会立即退出
|
||||
m.Start()
|
||||
if !m.running.Load() {
|
||||
t.Fatal("manager should be running after second Start")
|
||||
}
|
||||
|
||||
// 等待一段时间,确认 updateLoop 仍在运行
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
if !m.running.Load() {
|
||||
t.Fatal("manager should still be running after sleep")
|
||||
}
|
||||
|
||||
m.Stop()
|
||||
}
|
||||
|
||||
// TestSlowStartManager_SetFindTarget 验证 findTarget 回调被正确使用。
|
||||
func TestSlowStartManager_SetFindTarget(t *testing.T) {
|
||||
m := NewSlowStartManager(10 * time.Millisecond)
|
||||
|
||||
target := &Target{URL: "http://backend:8080", Weight: 10, SlowStart: 100 * time.Millisecond}
|
||||
m.OnTargetHealthy(target)
|
||||
|
||||
called := atomic.Bool{}
|
||||
m.SetFindTarget(func(url string) *Target {
|
||||
called.Store(true)
|
||||
if url == "http://backend:8080" {
|
||||
return target
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
m.Start()
|
||||
defer m.Stop()
|
||||
|
||||
// 等待 updateLoop 执行
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if !called.Load() {
|
||||
t.Error("findTarget callback should be called by updateLoop")
|
||||
}
|
||||
}
|
||||
import "testing"
|
||||
|
||||
func TestTarget_GetEffectiveWeight(t *testing.T) {
|
||||
tests := []struct {
|
||||
|
||||
@ -434,57 +434,3 @@ func TestLoggerClose(t *testing.T) {
|
||||
t.Errorf("Unexpected error on Close: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoggerClose_ClosesFileHandles 验证 Close 真正关闭了底层文件句柄。
|
||||
func TestLoggerClose_ClosesFileHandles(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
accessPath := filepath.Join(tmpDir, "access.log")
|
||||
errorPath := filepath.Join(tmpDir, "error.log")
|
||||
|
||||
cfg := &config.LoggingConfig{
|
||||
Access: config.AccessLogConfig{Path: accessPath},
|
||||
Error: config.ErrorLogConfig{Path: errorPath},
|
||||
}
|
||||
|
||||
logger := New(cfg)
|
||||
|
||||
// 先写入一些日志
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Request.SetRequestURI("/closed")
|
||||
ctx.Request.Header.SetMethod("GET")
|
||||
logger.LogAccess(ctx, 200, 10, 100*time.Millisecond)
|
||||
logger.Error().Msg("error before close")
|
||||
|
||||
// 关闭后应能正常删除文件(Windows 下需要句柄关闭)
|
||||
if err := logger.Close(); err != nil {
|
||||
t.Fatalf("Close failed: %v", err)
|
||||
}
|
||||
|
||||
if err := os.Remove(accessPath); err != nil {
|
||||
t.Errorf("should be able to remove access log after Close: %v", err)
|
||||
}
|
||||
if err := os.Remove(errorPath); err != nil {
|
||||
t.Errorf("should be able to remove error log after Close: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLogAccessDoesNotMutateBuffer 验证访问日志不会修改 fasthttp 内部缓冲区。
|
||||
func TestLogAccessDoesNotMutateBuffer(t *testing.T) {
|
||||
logger := New(&config.LoggingConfig{Access: config.AccessLogConfig{Format: "json"}})
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Request.SetRequestURI("/api/users")
|
||||
ctx.Request.Header.SetMethod("POST")
|
||||
|
||||
methodBefore := string(ctx.Method())
|
||||
pathBefore := string(ctx.Path())
|
||||
|
||||
logger.LogAccess(ctx, 200, 10, 100*time.Millisecond)
|
||||
|
||||
if string(ctx.Method()) != methodBefore {
|
||||
t.Errorf("Method buffer was mutated: got %q, want %q", ctx.Method(), methodBefore)
|
||||
}
|
||||
if string(ctx.Path()) != pathBefore {
|
||||
t.Errorf("Path buffer was mutated: got %q, want %q", ctx.Path(), pathBefore)
|
||||
}
|
||||
}
|
||||
|
||||
@ -3,7 +3,6 @@ package lua
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
@ -39,36 +38,6 @@ func TestNgxLogAPIWithoutLogger(t *testing.T) {
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestNgxLogEmergDoesNotExit 验证 ngx.EMERG 不会导致进程退出。
|
||||
func TestNgxLogEmergDoesNotExit(t *testing.T) {
|
||||
engine, err := NewEngine(DefaultConfig())
|
||||
require.NoError(t, err)
|
||||
defer engine.Close()
|
||||
|
||||
ctx := mockRequestCtxForLog()
|
||||
luaCtx := NewContext(engine, ctx)
|
||||
|
||||
var buf bytes.Buffer
|
||||
logger := zerolog.New(&buf)
|
||||
|
||||
api := newNgxLogAPI(ctx, luaCtx, &logger)
|
||||
L := engine.GetLStateForTest()
|
||||
RegisterNgxLogAPI(L, api)
|
||||
|
||||
// 不应 panic 或退出进程
|
||||
err = L.DoString(`
|
||||
ngx.log(ngx.EMERG, "emergency message")
|
||||
ngx.log(ngx.ALERT, "alert message")
|
||||
ngx.log(ngx.CRIT, "critical message")
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
|
||||
output := buf.String()
|
||||
if !strings.Contains(output, "emergency message") {
|
||||
t.Error("expected emerg message to be logged")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNgxLogAPIIntegration 集成测试
|
||||
func TestNgxLogAPIIntegration(t *testing.T) {
|
||||
engine, err := NewEngine(DefaultConfig())
|
||||
|
||||
@ -110,19 +110,10 @@ func NewTCPSocket(manager *CosocketManager) *TCPSocket {
|
||||
// 返回值:
|
||||
// - error: 状态不正确或地址解析失败时返回错误
|
||||
func (s *TCPSocket) Connect(host string, port int) error {
|
||||
_, err := s.connectWithOp(host, port)
|
||||
return err
|
||||
}
|
||||
|
||||
// connectWithOp 连接到指定地址并返回 SocketOperation。
|
||||
//
|
||||
// 将 SocketOperation 直接返回给调用方,避免 ConnectAsync 在 Connect 返回后
|
||||
// 读取 s.currentOp 时与后台 goroutine 设置 s.currentOp = nil 产生竞态。
|
||||
func (s *TCPSocket) connectWithOp(host string, port int) (*SocketOperation, error) {
|
||||
s.mu.Lock()
|
||||
if s.state != SocketStateIdle {
|
||||
s.mu.Unlock()
|
||||
return nil, fmt.Errorf("socket not idle, current state: %s", s.state)
|
||||
return fmt.Errorf("socket not idle, current state: %s", s.state)
|
||||
}
|
||||
s.state = SocketStateConnecting
|
||||
s.mu.Unlock()
|
||||
@ -131,13 +122,13 @@ func (s *TCPSocket) connectWithOp(host string, port int) (*SocketOperation, erro
|
||||
addr, err := s.manager.TCPAddr(host, port)
|
||||
if err != nil {
|
||||
s.setState(SocketStateError)
|
||||
return nil, fmt.Errorf("resolve address: %w", err)
|
||||
return fmt.Errorf("resolve address: %w", err)
|
||||
}
|
||||
|
||||
// IP 字面量:立即检查受限地址
|
||||
if !s.manager.DisableSSRFGuard && addr.IP != nil && isRestrictedIP(addr.IP) {
|
||||
s.setState(SocketStateError)
|
||||
return nil, fmt.Errorf("connection to restricted address denied: %s", addr.IP)
|
||||
return fmt.Errorf("connection to restricted address denied: %s", addr.IP)
|
||||
}
|
||||
|
||||
s.addr = addr
|
||||
@ -192,12 +183,12 @@ func (s *TCPSocket) connectWithOp(host string, port int) (*SocketOperation, erro
|
||||
s.manager.CompleteOperation(op.ID, conn, nil)
|
||||
}()
|
||||
|
||||
return op, nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// ConnectAsync 异步连接(用于 Lua yield/resume)。
|
||||
//
|
||||
// 调用 connectWithOp 并返回关联的 SocketOperation,供 Lua 协程 yield 等待。
|
||||
// 调用 Connect 并返回关联的 SocketOperation,供 Lua 协程 yield 等待。
|
||||
//
|
||||
// 返回值:
|
||||
// - *SocketOperation: 连接操作实例
|
||||
@ -210,10 +201,13 @@ func (s *TCPSocket) ConnectAsync(_ *glua.LState, host string, port int) (*Socket
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
op, err := s.connectWithOp(host, port)
|
||||
err := s.Connect(host, port)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.mu.RLock()
|
||||
op := s.currentOp
|
||||
s.mu.RUnlock()
|
||||
return op, nil
|
||||
}
|
||||
|
||||
|
||||
@ -1348,33 +1348,3 @@ func TestSocketOperation_Touch(t *testing.T) {
|
||||
op.Touch()
|
||||
assert.True(t, op.LastActivity.After(oldTime))
|
||||
}
|
||||
|
||||
// TestTCPSocket_ConnectAsync_Concurrent 验证并发 ConnectAsync 不会返回 nil op 导致 panic。
|
||||
func TestTCPSocket_ConnectAsync_Concurrent(t *testing.T) {
|
||||
_, cleanup := mockEchoServer(t, "127.0.0.1:18820")
|
||||
defer cleanup()
|
||||
|
||||
cm := NewCosocketManager()
|
||||
defer cm.Close()
|
||||
|
||||
L := glua.NewState()
|
||||
defer L.Close()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 50 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
socket := NewTCPSocket(cm)
|
||||
defer socket.Close()
|
||||
|
||||
op, err := socket.ConnectAsync(L, "127.0.0.1", 18820)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
require.NotNil(t, op)
|
||||
_, _ = op.Wait(context.Background())
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
@ -258,10 +258,11 @@ func (m *TimerManager) executeTimer(entry *TimerEntry) {
|
||||
}
|
||||
select {
|
||||
case m.callbackQueue <- cbEntry:
|
||||
m.queueMu.Unlock()
|
||||
default:
|
||||
m.queueMu.Unlock()
|
||||
logging.Warn().Msg("[lua] timer callback dropped: queue full")
|
||||
}
|
||||
m.queueMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
@ -308,12 +309,23 @@ func (m *TimerManager) Cancel(handle *TimerHandle) bool {
|
||||
return false // 定时器不存在或已执行
|
||||
}
|
||||
|
||||
// 发送取消信号;executeTimer 检查到后会在 defer 中递减 active
|
||||
// 停止定时器;如果成功停止,回调不会执行
|
||||
stopped := true
|
||||
if entry.timer != nil {
|
||||
stopped = entry.timer.Stop()
|
||||
}
|
||||
|
||||
// 发送取消信号
|
||||
close(entry.cancel)
|
||||
|
||||
// 清理
|
||||
delete(m.timers, entry.id)
|
||||
|
||||
// 如果成功停止定时器(回调不会执行),递减 active
|
||||
if stopped {
|
||||
m.active.Add(-1)
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
@ -335,14 +347,13 @@ func (m *TimerManager) WaitAll(timeout time.Duration) bool {
|
||||
start := time.Now()
|
||||
for m.active.Load() > 0 {
|
||||
if time.Since(start) > timeout {
|
||||
// 超时,发送取消信号给所有剩余定时器
|
||||
// 超时,强制取消所有
|
||||
m.mu.Lock()
|
||||
for _, entry := range m.timers {
|
||||
select {
|
||||
case <-entry.cancel:
|
||||
default:
|
||||
close(entry.cancel)
|
||||
if entry.timer != nil {
|
||||
entry.timer.Stop()
|
||||
}
|
||||
close(entry.cancel)
|
||||
}
|
||||
m.timers = make(map[uint64]*TimerEntry)
|
||||
m.mu.Unlock()
|
||||
|
||||
@ -2,7 +2,6 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -151,37 +150,3 @@ func TestTimerRunningCount(t *testing.T) {
|
||||
// 应该回到 0
|
||||
assert.Equal(t, int32(0), manager.ActiveCount())
|
||||
}
|
||||
|
||||
// TestTimerManager_CancelConcurrent 验证并发 Cancel 与定时器触发不会产生负 active 计数或 panic。
|
||||
func TestTimerManager_CancelConcurrent(t *testing.T) {
|
||||
engine, err := NewEngine(DefaultConfig())
|
||||
require.NoError(t, err)
|
||||
defer engine.Close()
|
||||
|
||||
manager := engine.TimerManager()
|
||||
|
||||
L := engine.GetLStateForTest()
|
||||
defer engine.PutLStateForTest(L)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 100 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
callback := L.NewFunction(func(_ *glua.LState) int {
|
||||
return 0
|
||||
})
|
||||
handle, err := manager.At(10*time.Millisecond, callback, nil)
|
||||
require.NoError(t, err)
|
||||
manager.Cancel(handle)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// 等待所有定时器触发/取消完成
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
|
||||
// active 计数不应为负
|
||||
require.GreaterOrEqual(t, manager.ActiveCount(), int32(0))
|
||||
assert.Equal(t, int32(0), manager.ActiveCount())
|
||||
}
|
||||
|
||||
@ -80,8 +80,6 @@ func TestBoundaryCoroutineYieldAllowed(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestBoundaryTimerHandleCancel 测试定时器句柄取消。
|
||||
//
|
||||
// Cancel 仅关闭 cancel 通道,active 计数由 executeTimer 的 defer 在到期后统一递减。
|
||||
func TestBoundaryTimerHandleCancel(t *testing.T) {
|
||||
engine, err := NewEngine(DefaultConfig())
|
||||
require.NoError(t, err)
|
||||
@ -102,7 +100,7 @@ func TestBoundaryTimerHandleCancel(t *testing.T) {
|
||||
function timer_callback()
|
||||
-- 空回调
|
||||
end
|
||||
local handle, err = ngx.timer.at(0.01, timer_callback)
|
||||
local handle, err = ngx.timer.at(10, timer_callback)
|
||||
if not handle then
|
||||
error("Failed to create timer: " .. tostring(err))
|
||||
end
|
||||
@ -115,10 +113,8 @@ func TestBoundaryTimerHandleCancel(t *testing.T) {
|
||||
`)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 等待定时器到期后 active 计数递减到 0
|
||||
require.Eventually(t, func() bool {
|
||||
return timerMgr.ActiveCount() == 0
|
||||
}, 500*time.Millisecond, 10*time.Millisecond)
|
||||
// 验证定时器已取消
|
||||
assert.Equal(t, int32(0), timerMgr.ActiveCount())
|
||||
}
|
||||
|
||||
// TestBoundaryTimerHandleDoubleCancel 测试重复取消定时器。
|
||||
@ -155,9 +151,6 @@ func TestBoundaryTimerHandleDoubleCancel(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestBoundaryTimerRunningCount 测试定时器运行计数。
|
||||
//
|
||||
// Cancel 仅关闭 cancel 通道,active 计数由 executeTimer 的 defer 统一递减,
|
||||
// 因此取消后定时器仍会在到期时触发并减少计数。
|
||||
func TestBoundaryTimerRunningCount(t *testing.T) {
|
||||
engine, err := NewEngine(DefaultConfig())
|
||||
require.NoError(t, err)
|
||||
@ -178,25 +171,28 @@ func TestBoundaryTimerRunningCount(t *testing.T) {
|
||||
error("Initial count should be 0")
|
||||
end
|
||||
|
||||
-- 创建定时器(使用较短延迟便于测试完成)
|
||||
local h1 = ngx.timer.at(0.01, function() end)
|
||||
local h2 = ngx.timer.at(0.01, function() end)
|
||||
local h3 = ngx.timer.at(0.01, function() end)
|
||||
-- 创建定时器
|
||||
local h1 = ngx.timer.at(10, function() end)
|
||||
local h2 = ngx.timer.at(10, function() end)
|
||||
local h3 = ngx.timer.at(10, function() end)
|
||||
|
||||
count = ngx.timer.running_count()
|
||||
if count ~= 3 then
|
||||
error("Count should be 3, got " .. tostring(count))
|
||||
end
|
||||
|
||||
-- 取消一个;计数仍由 executeTimer 在到期时统一递减
|
||||
-- 取消一个
|
||||
h1:cancel()
|
||||
count = ngx.timer.running_count()
|
||||
if count ~= 2 then
|
||||
error("Count should be 2 after cancel, got " .. tostring(count))
|
||||
end
|
||||
|
||||
-- 清理
|
||||
h2:cancel()
|
||||
h3:cancel()
|
||||
`)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// 等待定时器到期并递减 active 计数
|
||||
require.Eventually(t, func() bool {
|
||||
return timerMgr.ActiveCount() == 0
|
||||
}, 500*time.Millisecond, 10*time.Millisecond)
|
||||
}
|
||||
|
||||
// TestBoundaryTimerUpvalueRejected 测试定时器拒绝闭包变量。
|
||||
|
||||
@ -17,41 +17,26 @@ func init() {
|
||||
testingSSRFGuardDisabled = true
|
||||
}
|
||||
|
||||
// mockEchoServer 模拟 echo 服务器。
|
||||
//
|
||||
// 使用短 deadline 的 Accept 循环,确保收到停止信号后能在关闭 listener 前退出,
|
||||
// 避免 listener.Close 与 listener.Accept 并发执行触发 race detector。
|
||||
// mockEchoServer 模拟 echo 服务器
|
||||
func mockEchoServer(t *testing.T, addr string) (net.Listener, func()) {
|
||||
ln, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to listen: %v", err)
|
||||
}
|
||||
|
||||
tcpLn, ok := ln.(*net.TCPListener)
|
||||
if !ok {
|
||||
t.Fatalf("Expected TCPListener, got %T", ln)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
stop := make(chan struct{})
|
||||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
_ = tcpLn.SetDeadline(time.Now().Add(100 * time.Millisecond))
|
||||
conn, err := tcpLn.Accept()
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
@ -75,8 +60,7 @@ func mockEchoServer(t *testing.T, addr string) (net.Listener, func()) {
|
||||
|
||||
cleanup := func() {
|
||||
close(stop)
|
||||
<-done
|
||||
_ = ln.Close()
|
||||
ln.Close()
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
|
||||
@ -14,7 +14,6 @@ import (
|
||||
"bytes"
|
||||
"io"
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/andybalholm/brotli"
|
||||
@ -545,32 +544,3 @@ func TestMiddleware_CompressWhenNoPrecompressed(t *testing.T) {
|
||||
t.Errorf("Expected compressed body smaller than original, got %d >= %d", len(body), len(originalBody))
|
||||
}
|
||||
}
|
||||
|
||||
// TestCompressorPoolConcurrentGet 验证并发从池中获取 writer 不会产生竞态。
|
||||
//
|
||||
// 该测试针对 pool.New 在构造函数中初始化的场景,确保多个 goroutine 同时
|
||||
// Get/Put 时 sync.Pool 的行为是安全的。
|
||||
func TestCompressorPoolConcurrentGet(t *testing.T) {
|
||||
m, err := New(&config.CompressionConfig{
|
||||
Type: "gzip",
|
||||
Level: 6,
|
||||
MinSize: 10,
|
||||
Types: []string{"text/html"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New() error: %v", err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 100 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
w, ok := m.gzipPool.Get()
|
||||
if ok {
|
||||
m.gzipPool.Put(w)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
@ -348,10 +348,49 @@ func (ac *AccessControl) SetDefault(action string) error {
|
||||
// 返回值:
|
||||
// - net.IP: 客户端 IP 地址,无法获取时返回 nil
|
||||
func (ac *AccessControl) getClientIP(ctx *fasthttp.RequestCtx) net.IP {
|
||||
ac.mu.RLock()
|
||||
trusted := ac.trustedProxies
|
||||
ac.mu.RUnlock()
|
||||
return netutil.ExtractClientIPWithTrustedProxies(ctx, trusted)
|
||||
remoteIP := netutil.GetRemoteAddrIP(ctx)
|
||||
|
||||
if len(ac.trustedProxies) == 0 || remoteIP == nil {
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
isTrusted := false
|
||||
for _, network := range ac.trustedProxies {
|
||||
if network.Contains(remoteIP) {
|
||||
isTrusted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !isTrusted {
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
if xff := ctx.Request.Header.Peek("X-Forwarded-For"); len(xff) > 0 {
|
||||
ips := strings.Split(string(xff), ",")
|
||||
for _, v := range slices.Backward(ips) {
|
||||
ipStr := strings.TrimSpace(v)
|
||||
if ip := net.ParseIP(ipStr); ip != nil {
|
||||
trusted := false
|
||||
for _, network := range ac.trustedProxies {
|
||||
if network.Contains(ip) {
|
||||
trusted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !trusted {
|
||||
return ip
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if xri := ctx.Request.Header.Peek("X-Real-IP"); len(xri) > 0 {
|
||||
if ip := net.ParseIP(string(xri)); ip != nil {
|
||||
return ip
|
||||
}
|
||||
}
|
||||
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
// AccessStats 访问控制统计信息结构。
|
||||
|
||||
@ -202,7 +202,7 @@ func TestAccessControlProcess_TrustedProxies_UntrustedSource(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccessControlProcess_TrustedProxies_XRealIP 测试可信代理 X-Real-IP 不被单独信任
|
||||
// TestAccessControlProcess_TrustedProxies_XRealIP 测试可信代理 X-Real-IP
|
||||
func TestAccessControlProcess_TrustedProxies_XRealIP(t *testing.T) {
|
||||
ac, err := NewAccessControl(&config.AccessConfig{
|
||||
TrustedProxies: []string{"10.0.0.0/8"},
|
||||
@ -225,8 +225,8 @@ func TestAccessControlProcess_TrustedProxies_XRealIP(t *testing.T) {
|
||||
ctx.Request.Header.Set("X-Real-IP", "192.168.1.50")
|
||||
handler(ctx)
|
||||
|
||||
if called {
|
||||
t.Error("Process() should NOT use X-Real-IP without X-Forwarded-For")
|
||||
if !called {
|
||||
t.Error("Process() should use X-Real-IP from trusted proxy")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -133,15 +133,6 @@ func (sh *HeadersMiddleware) addHeaders(ctx *fasthttp.RequestCtx) {
|
||||
hstsValue := sh.hsts
|
||||
sh.mu.RUnlock()
|
||||
|
||||
// nil 保护:防止并发更新将配置置为 nil 后 panic
|
||||
if cfg == nil {
|
||||
cfg = &config.SecurityHeaders{
|
||||
XFrameOptions: "DENY",
|
||||
XContentTypeOptions: "nosniff",
|
||||
ReferrerPolicy: "strict-origin-when-cross-origin",
|
||||
}
|
||||
}
|
||||
|
||||
// X-Frame-Options
|
||||
if cfg.XFrameOptions != "" {
|
||||
headers.Set("X-Frame-Options", cfg.XFrameOptions)
|
||||
|
||||
@ -2,37 +2,11 @@
|
||||
//
|
||||
// 该文件测试安全头部模块的各项功能,包括:
|
||||
// - HSTS 值格式化
|
||||
// - nil 配置保护
|
||||
//
|
||||
// 作者:xfy
|
||||
package security
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
// TestAddHeadersNilConfig 验证配置在运行时被置为 nil 也不会 panic。
|
||||
func TestAddHeadersNilConfig(t *testing.T) {
|
||||
sh := NewHeadersWithHSTS(nil, nil)
|
||||
// 故意将内部配置置为 nil,模拟异常状态
|
||||
sh.config = nil
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Request.Header.SetMethod("GET")
|
||||
ctx.Request.Header.SetRequestURI("/")
|
||||
|
||||
// 不应 panic
|
||||
sh.addHeaders(ctx)
|
||||
|
||||
if string(ctx.Response.Header.Peek("X-Frame-Options")) != "DENY" {
|
||||
t.Errorf("expected default X-Frame-Options=DENY, got %q", ctx.Response.Header.Peek("X-Frame-Options"))
|
||||
}
|
||||
if string(ctx.Response.Header.Peek("X-Content-Type-Options")) != "nosniff" {
|
||||
t.Errorf("expected default X-Content-Type-Options=nosniff, got %q", ctx.Response.Header.Peek("X-Content-Type-Options"))
|
||||
}
|
||||
}
|
||||
import "testing"
|
||||
|
||||
func TestFormatHSTSValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
|
||||
@ -32,7 +32,6 @@ package security
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@ -66,15 +65,14 @@ type shardedBucket struct {
|
||||
// - 所有方法均为并发安全
|
||||
// - 启动后会自动后台清理过期的桶
|
||||
type RateLimiter struct {
|
||||
shards [16]shardedBucket
|
||||
keyFunc KeyFunc
|
||||
cleanupTicker *time.Ticker
|
||||
stopCleanupCh chan struct{}
|
||||
cleanupDone chan struct{}
|
||||
stopOnce sync.Once
|
||||
rate float64
|
||||
burst float64
|
||||
trustedProxies []net.IPNet
|
||||
shards [16]shardedBucket
|
||||
keyFunc KeyFunc
|
||||
cleanupTicker *time.Ticker
|
||||
stopCleanupCh chan struct{}
|
||||
cleanupDone chan struct{}
|
||||
stopOnce sync.Once
|
||||
rate float64
|
||||
burst float64
|
||||
}
|
||||
|
||||
// tokenBucket 表示单个限流键的令牌桶。
|
||||
@ -97,12 +95,11 @@ type KeyFunc func(ctx *fasthttp.RequestCtx) string
|
||||
//
|
||||
// 参数:
|
||||
// - cfg: 限流配置,包含速率、突发量和键类型
|
||||
// - trustedProxies: 可信代理 CIDR 列表,用于安全提取客户端真实 IP
|
||||
//
|
||||
// 返回值:
|
||||
// - *RateLimiter: 配置好的限流器实例
|
||||
// - error: 配置无效时返回错误(如速率小于 0)
|
||||
func NewRateLimiter(cfg *config.RateLimitConfig, trustedProxies ...string) (middleware.Middleware, error) {
|
||||
func NewRateLimiter(cfg *config.RateLimitConfig) (middleware.Middleware, error) {
|
||||
if cfg == nil {
|
||||
return nil, errors.New("rate limit config is nil")
|
||||
}
|
||||
@ -119,7 +116,7 @@ func NewRateLimiter(cfg *config.RateLimitConfig, trustedProxies ...string) (midd
|
||||
|
||||
switch algorithm {
|
||||
case "token_bucket", "":
|
||||
return newTokenBucketLimiter(cfg, trustedProxies)
|
||||
return newTokenBucketLimiter(cfg)
|
||||
case "sliding_window":
|
||||
window := time.Duration(cfg.SlidingWindow) * time.Second
|
||||
if window <= 0 {
|
||||
@ -133,7 +130,7 @@ func NewRateLimiter(cfg *config.RateLimitConfig, trustedProxies ...string) (midd
|
||||
}
|
||||
|
||||
// newTokenBucketLimiter 创建令牌桶限流器。
|
||||
func newTokenBucketLimiter(cfg *config.RateLimitConfig, trustedProxies []string) (*RateLimiter, error) {
|
||||
func newTokenBucketLimiter(cfg *config.RateLimitConfig) (*RateLimiter, error) {
|
||||
if cfg.Burst < cfg.RequestRate {
|
||||
return nil, errors.New("burst must be at least equal to request rate")
|
||||
}
|
||||
@ -152,24 +149,12 @@ func newTokenBucketLimiter(cfg *config.RateLimitConfig, trustedProxies []string)
|
||||
}
|
||||
}
|
||||
|
||||
// 解析可信代理 CIDR
|
||||
for _, cidr := range trustedProxies {
|
||||
_, ipNet, err := net.ParseCIDR(cidr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid trusted proxy CIDR %s: %w", cidr, err)
|
||||
}
|
||||
rl.trustedProxies = append(rl.trustedProxies, *ipNet)
|
||||
}
|
||||
|
||||
// 根据配置设置键提取函数
|
||||
switch cfg.Key {
|
||||
case "ip", "":
|
||||
rl.keyFunc = rl.keyByIP
|
||||
case rateLimitHeader:
|
||||
rl.keyFunc = rl.keyByHeader
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown key type: %s", cfg.Key)
|
||||
keyFunc, err := parseKeyFunc(cfg.Key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rl.keyFunc = keyFunc
|
||||
|
||||
// 启动后台清理 goroutine
|
||||
rl.startCleanup(10 * time.Minute)
|
||||
@ -411,29 +396,6 @@ func keyByHeader(ctx *fasthttp.RequestCtx) string {
|
||||
return string(key)
|
||||
}
|
||||
|
||||
// keyByIP 提取客户端 IP 作为限流键(令牌桶限流器方法)。
|
||||
//
|
||||
// 仅当请求来自可信代理时,才解析 X-Forwarded-For 头部获取真实客户端 IP。
|
||||
func (rl *RateLimiter) keyByIP(ctx *fasthttp.RequestCtx) string {
|
||||
ip := netutil.ExtractClientIPWithTrustedProxies(ctx, rl.trustedProxies)
|
||||
if ip == nil {
|
||||
return accessUnknown
|
||||
}
|
||||
return ip.String()
|
||||
}
|
||||
|
||||
// keyByHeader 提取头部值作为限流键(令牌桶限流器方法)。
|
||||
//
|
||||
// 默认使用 X-RateLimit-Key 头部,如果不存在则回退到 IP。
|
||||
func (rl *RateLimiter) keyByHeader(ctx *fasthttp.RequestCtx) string {
|
||||
key := ctx.Request.Header.Peek("X-RateLimit-Key")
|
||||
if len(key) == 0 {
|
||||
// 头部不存在时回退到 IP
|
||||
return rl.keyByIP(ctx)
|
||||
}
|
||||
return string(key)
|
||||
}
|
||||
|
||||
// parseKeyFunc 根据配置字符串解析键提取函数。
|
||||
//
|
||||
// 支持的键类型:
|
||||
|
||||
@ -12,8 +12,6 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"net"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -869,97 +867,3 @@ func TestConnLimiter_MiddlewareIdentity(t *testing.T) {
|
||||
t.Errorf("Name() = %s, want 'conn_limiter'", middleware.Name())
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_MultipleCreateDestroy 验证多次创建/销毁限流器不会泄漏 goroutine 也不会 panic。
|
||||
func TestRateLimiter_MultipleCreateDestroy(t *testing.T) {
|
||||
cfg := &config.RateLimitConfig{
|
||||
RequestRate: 100,
|
||||
Burst: 200,
|
||||
}
|
||||
|
||||
// 预热:先创建并销毁一次,让 runtime 稳定
|
||||
mw, err := NewRateLimiter(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||||
}
|
||||
rl := mw.(*RateLimiter)
|
||||
rl.StopCleanup()
|
||||
|
||||
before := runtime.NumGoroutine()
|
||||
|
||||
for range 10 {
|
||||
mw, err = NewRateLimiter(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||||
}
|
||||
mw.(*RateLimiter).StopCleanup()
|
||||
// 重复调用不应 panic
|
||||
mw.(*RateLimiter).StopCleanup()
|
||||
}
|
||||
|
||||
after := runtime.NumGoroutine()
|
||||
// 允许少量波动,但不应有明显泄漏
|
||||
if after-before > 5 {
|
||||
t.Fatalf("goroutine leak detected: before=%d, after=%d", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_TrustedProxies 验证令牌桶限流器使用可信代理安全提取客户端 IP。
|
||||
func TestRateLimiter_TrustedProxies(t *testing.T) {
|
||||
mw, err := NewRateLimiter(&config.RateLimitConfig{
|
||||
RequestRate: 1,
|
||||
Burst: 1,
|
||||
}, "10.0.0.0/8")
|
||||
if err != nil {
|
||||
t.Fatalf("NewRateLimiter() error: %v", err)
|
||||
}
|
||||
|
||||
rl, ok := mw.(*RateLimiter)
|
||||
if !ok {
|
||||
t.Fatalf("Expected *RateLimiter, got %T", mw)
|
||||
}
|
||||
defer rl.StopCleanup()
|
||||
|
||||
callCount := 0
|
||||
nextHandler := func(ctx *fasthttp.RequestCtx) {
|
||||
callCount++
|
||||
ctx.SetStatusCode(200)
|
||||
}
|
||||
handler := rl.Process(nextHandler)
|
||||
|
||||
// 来自可信代理的请求,使用 X-Forwarded-For 中的真实客户端 IP
|
||||
ctx1 := testutil.NewRequestCtx("GET", "/test")
|
||||
ctx1.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")})
|
||||
ctx1.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
handler(ctx1)
|
||||
if ctx1.Response.StatusCode() != 200 {
|
||||
t.Errorf("first trusted request should be allowed, got %d", ctx1.Response.StatusCode())
|
||||
}
|
||||
|
||||
// 同一个真实客户端 IP 的第二个请求应被拒绝
|
||||
ctx2 := testutil.NewRequestCtx("GET", "/test")
|
||||
ctx2.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.2")})
|
||||
ctx2.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
handler(ctx2)
|
||||
if ctx2.Response.StatusCode() != 429 {
|
||||
t.Errorf("second request from same client IP should be rate limited, got %d", ctx2.Response.StatusCode())
|
||||
}
|
||||
|
||||
// 不同真实客户端 IP 的请求应被允许
|
||||
ctx3 := testutil.NewRequestCtx("GET", "/test")
|
||||
ctx3.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")})
|
||||
ctx3.Request.Header.Set("X-Forwarded-For", "5.6.7.8")
|
||||
handler(ctx3)
|
||||
if ctx3.Response.StatusCode() != 200 {
|
||||
t.Errorf("request from different client IP should be allowed, got %d", ctx3.Response.StatusCode())
|
||||
}
|
||||
|
||||
// 不可信代理:忽略 XFF,使用 remoteAddr 作为限流键
|
||||
ctx4 := testutil.NewRequestCtx("GET", "/test")
|
||||
ctx4.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("192.168.1.1")})
|
||||
ctx4.Request.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
handler(ctx4)
|
||||
if ctx4.Response.StatusCode() != 200 {
|
||||
t.Errorf("request from untrusted proxy should use remote IP and be allowed, got %d", ctx4.Response.StatusCode())
|
||||
}
|
||||
}
|
||||
|
||||
@ -36,27 +36,6 @@ func TestNewSlidingWindowLimiter(t *testing.T) {
|
||||
t.Error("precise should be true")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("window 为零时使用默认值避免除零", func(t *testing.T) {
|
||||
limiter := NewSlidingWindowLimiter(0, 10, false)
|
||||
if limiter.window <= 0 {
|
||||
t.Errorf("window = %v, want > 0", limiter.window)
|
||||
}
|
||||
// 验证不会触发除零 panic
|
||||
if !limiter.Allow("test-key") {
|
||||
t.Error("第一个请求应该被允许")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("window 为负时使用默认值避免除零", func(t *testing.T) {
|
||||
limiter := NewSlidingWindowLimiter(-time.Second, 10, false)
|
||||
if limiter.window <= 0 {
|
||||
t.Errorf("window = %v, want > 0", limiter.window)
|
||||
}
|
||||
if !limiter.Allow("test-key") {
|
||||
t.Error("第一个请求应该被允许")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSlidingWindowLimiter_Allow(t *testing.T) {
|
||||
|
||||
@ -107,16 +107,3 @@ func TestAddTypes(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDetectContentTypeUnknownCached 测试未知扩展名多次查询均回退到默认 MIME 类型。
|
||||
// 防止缓存中写入空字符串导致后续查询返回错误值。
|
||||
func TestDetectContentTypeUnknownCached(t *testing.T) {
|
||||
const unknownFile = "test.unknownxyz"
|
||||
|
||||
for range 3 {
|
||||
got := DetectContentType(unknownFile)
|
||||
if got != defaultMIME {
|
||||
t.Errorf("DetectContentType(%q) = %q, want %q", unknownFile, got, defaultMIME)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -112,48 +112,6 @@ func GetRemoteAddrIP(ctx *fasthttp.RequestCtx) net.IP {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ExtractClientIPWithTrustedProxies 从请求上下文安全提取客户端 IP。
|
||||
// 仅当 remoteAddr 属于 trustedProxies 时,才解析 X-Forwarded-For 头部,
|
||||
// 并取右侧(最接近服务器)第一个非可信 IP 作为真实客户端 IP。
|
||||
func ExtractClientIPWithTrustedProxies(ctx *fasthttp.RequestCtx, trustedProxies []net.IPNet) net.IP {
|
||||
remoteIP := GetRemoteAddrIP(ctx)
|
||||
if remoteIP == nil || len(trustedProxies) == 0 {
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
trusted := false
|
||||
for i := range trustedProxies {
|
||||
if trustedProxies[i].Contains(remoteIP) {
|
||||
trusted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !trusted {
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
if xff := ctx.Request.Header.Peek("X-Forwarded-For"); len(xff) > 0 {
|
||||
ips := strings.Split(string(xff), ",")
|
||||
for i := len(ips) - 1; i >= 0; i-- {
|
||||
ipStr := strings.TrimSpace(ips[i])
|
||||
if ip := net.ParseIP(ipStr); ip != nil {
|
||||
isTrusted := false
|
||||
for j := range trustedProxies {
|
||||
if trustedProxies[j].Contains(ip) {
|
||||
isTrusted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !isTrusted {
|
||||
return ip
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return remoteIP
|
||||
}
|
||||
|
||||
// remoteAddrCache 缓存 RemoteAddr 字符串化结果,避免重复的 net.TCPAddr.String() 分配。
|
||||
type remoteAddrCache struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
@ -130,32 +130,3 @@ func TestGetRemoteAddrIP(_ *testing.T) {
|
||||
// Just verify it doesn't panic
|
||||
_ = got
|
||||
}
|
||||
|
||||
func TestExtractClientIPWithTrustedProxies(t *testing.T) {
|
||||
trusted := []net.IPNet{mustParseCIDR(t, "10.0.0.0/8")}
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.0.0.1")})
|
||||
ctx.Request.Header.Set("X-Forwarded-For", "1.2.3.4, 10.0.0.2")
|
||||
|
||||
ip := ExtractClientIPWithTrustedProxies(ctx, trusted)
|
||||
if ip == nil || ip.String() != "1.2.3.4" {
|
||||
t.Fatalf("expected 1.2.3.4, got %v", ip)
|
||||
}
|
||||
|
||||
// Untrusted proxy: ignore XFF.
|
||||
ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("192.168.1.1")})
|
||||
ip = ExtractClientIPWithTrustedProxies(ctx, trusted)
|
||||
if ip == nil || ip.String() != "192.168.1.1" {
|
||||
t.Fatalf("expected 192.168.1.1, got %v", ip)
|
||||
}
|
||||
}
|
||||
|
||||
func mustParseCIDR(t *testing.T, s string) net.IPNet {
|
||||
t.Helper()
|
||||
_, ipNet, err := net.ParseCIDR(s)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse CIDR %s: %v", s, err)
|
||||
}
|
||||
return *ipNet
|
||||
}
|
||||
|
||||
@ -15,46 +15,41 @@ import (
|
||||
|
||||
// buildCacheKeyHash 使用 FNV-64a 计算缓存键的 uint64 哈希值。
|
||||
// 使用零分配方式构建哈希,避免 []byte(origKey) 转换。
|
||||
// 缓存键包含请求方法、Host、URI 以及配置的 Vary 请求头,防止缓存中毒。
|
||||
func (p *Proxy) buildCacheKeyHash(ctx *fasthttp.RequestCtx, varyHeaders []string) (uint64, string) {
|
||||
return p.buildCacheKeyHashWithHost(ctx, ctx.Request.Header.Host(), varyHeaders)
|
||||
}
|
||||
|
||||
// buildCacheKeyHashWithHost 与 buildCacheKeyHash 相同,但使用显式传入的 Host,
|
||||
// 避免请求头在缓存键计算前被修改(例如 proxy 将 Host 改写为上游目标)。
|
||||
func (p *Proxy) buildCacheKeyHashWithHost(ctx *fasthttp.RequestCtx, host []byte, varyHeaders []string) (uint64, string) {
|
||||
func (p *Proxy) buildCacheKeyHash(ctx *fasthttp.RequestCtx) (uint64, string) {
|
||||
method := ctx.Request.Header.Method()
|
||||
uri := ctx.Request.URI().RequestURI()
|
||||
|
||||
var h uint64 = 14695981039346656037
|
||||
for _, b := range [][]byte{method, host, uri} {
|
||||
h ^= uint64(':')
|
||||
for i := 0; i < len(method); i++ {
|
||||
h ^= uint64(method[i])
|
||||
h *= 1099511628211
|
||||
}
|
||||
h ^= uint64(':')
|
||||
h *= 1099511628211
|
||||
for i := 0; i < len(uri); i++ {
|
||||
h ^= uint64(uri[i])
|
||||
h *= 1099511628211
|
||||
for i := 0; i < len(b); i++ {
|
||||
h ^= uint64(b[i])
|
||||
h *= 1099511628211
|
||||
}
|
||||
}
|
||||
|
||||
for _, name := range varyHeaders {
|
||||
value := ctx.Request.Header.Peek(name)
|
||||
h ^= uint64(':')
|
||||
h *= 1099511628211
|
||||
for i := 0; i < len(value); i++ {
|
||||
h ^= uint64(value[i])
|
||||
h *= 1099511628211
|
||||
}
|
||||
}
|
||||
|
||||
origKey := b2s(method) + ":" + b2s(host) + b2s(uri)
|
||||
for _, name := range varyHeaders {
|
||||
origKey += ":" + name + "=" + b2s(ctx.Request.Header.Peek(name))
|
||||
}
|
||||
origKey := b2s(method) + ":" + b2s(uri)
|
||||
return h, origKey
|
||||
}
|
||||
|
||||
func (p *Proxy) buildCacheKeyHashValue(ctx *fasthttp.RequestCtx, varyHeaders []string) uint64 {
|
||||
h, _ := p.buildCacheKeyHash(ctx, varyHeaders)
|
||||
func (p *Proxy) buildCacheKeyHashValue(ctx *fasthttp.RequestCtx) uint64 {
|
||||
method := ctx.Request.Header.Method()
|
||||
uri := ctx.Request.URI().RequestURI()
|
||||
|
||||
var h uint64 = 14695981039346656037
|
||||
for i := 0; i < len(method); i++ {
|
||||
h ^= uint64(method[i])
|
||||
h *= 1099511628211
|
||||
}
|
||||
h ^= uint64(':')
|
||||
h *= 1099511628211
|
||||
for i := 0; i < len(uri); i++ {
|
||||
h ^= uint64(uri[i])
|
||||
h *= 1099511628211
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
@ -126,8 +121,14 @@ func (p *Proxy) backgroundRefresh(req *fasthttp.Request, target *loadbalance.Tar
|
||||
return
|
||||
}
|
||||
|
||||
// 提取响应头
|
||||
headers := make(map[string]string, 20)
|
||||
// 提取响应头(使用 pool 复用 map)
|
||||
headers, ok := headersPool.Get().(map[string]string)
|
||||
if !ok {
|
||||
headers = make(map[string]string, 20)
|
||||
}
|
||||
for k := range headers {
|
||||
delete(headers, k)
|
||||
}
|
||||
for key, value := range resp.Header.All() {
|
||||
headers[string(key)] = string(value)
|
||||
}
|
||||
|
||||
@ -1,28 +0,0 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
func TestBuildCacheKeyIncludesHostAndVary(t *testing.T) {
|
||||
p := &Proxy{}
|
||||
ctx1 := &fasthttp.RequestCtx{}
|
||||
ctx1.Request.Header.SetMethod("GET")
|
||||
ctx1.Request.Header.SetHost("a.example.com")
|
||||
ctx1.Request.SetRequestURI("/api/data")
|
||||
|
||||
ctx2 := &fasthttp.RequestCtx{}
|
||||
ctx2.Request.Header.SetMethod("GET")
|
||||
ctx2.Request.Header.SetHost("b.example.com")
|
||||
ctx2.Request.SetRequestURI("/api/data")
|
||||
|
||||
vary := []string{"Accept-Encoding"}
|
||||
h1, key1 := p.buildCacheKeyHash(ctx1, vary)
|
||||
h2, key2 := p.buildCacheKeyHash(ctx2, vary)
|
||||
|
||||
if h1 == h2 || key1 == key2 {
|
||||
t.Fatal("cache keys must differ by Host")
|
||||
}
|
||||
}
|
||||
@ -42,7 +42,7 @@ func BenchmarkCacheKeyHashValue_ZeroAlloc(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
|
||||
for b.Loop() {
|
||||
hash := p.buildCacheKeyHashValue(ctx, nil)
|
||||
hash := p.buildCacheKeyHashValue(ctx)
|
||||
_ = hash
|
||||
}
|
||||
}
|
||||
@ -72,7 +72,7 @@ func BenchmarkCacheKeyHash_WithAlloc(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
|
||||
for b.Loop() {
|
||||
hash, key := p.buildCacheKeyHash(ctx, nil)
|
||||
hash, key := p.buildCacheKeyHash(ctx)
|
||||
_ = hash
|
||||
_ = key
|
||||
}
|
||||
@ -96,7 +96,7 @@ func BenchmarkCacheKeyHash_Compare(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
_ = p.buildCacheKeyHashValue(ctx, nil)
|
||||
_ = p.buildCacheKeyHashValue(ctx)
|
||||
}
|
||||
})
|
||||
|
||||
@ -104,7 +104,7 @@ func BenchmarkCacheKeyHash_Compare(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
_, _ = p.buildCacheKeyHash(ctx, nil)
|
||||
_, _ = p.buildCacheKeyHash(ctx)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@ -115,7 +115,7 @@ func NewHealthChecker(targets []*loadbalance.Target, cfg *config.HealthCheckConf
|
||||
slowStartManager = loadbalance.NewSlowStartManager(cfg.SlowStart)
|
||||
}
|
||||
|
||||
h := &HealthChecker{
|
||||
return &HealthChecker{
|
||||
targets: targets,
|
||||
interval: interval,
|
||||
timeout: timeout,
|
||||
@ -128,13 +128,6 @@ func NewHealthChecker(targets []*loadbalance.Target, cfg *config.HealthCheckConf
|
||||
WriteTimeout: timeout,
|
||||
},
|
||||
}
|
||||
|
||||
// 初始化慢启动管理器的目标查找回调
|
||||
if h.slowStartManager != nil {
|
||||
h.slowStartManager.SetFindTarget(h.findTargetByURL)
|
||||
}
|
||||
|
||||
return h
|
||||
}
|
||||
|
||||
// Start 启动后台健康检查进程。
|
||||
@ -288,22 +281,3 @@ func (h *HealthChecker) MarkHealthy(target *loadbalance.Target) {
|
||||
h.slowStartManager.OnTargetHealthy(target)
|
||||
}
|
||||
}
|
||||
|
||||
// findTargetByURL 根据 URL 查找对应的后端目标。
|
||||
//
|
||||
// 作为 SlowStartManager 的目标查找回调,用于将慢启动状态
|
||||
// 同步到负载均衡目标的 EffectiveWeight。
|
||||
//
|
||||
// 参数:
|
||||
// - url: 后端目标 URL
|
||||
//
|
||||
// 返回值:
|
||||
// - *loadbalance.Target: 找到的目标,不存在时返回 nil
|
||||
func (h *HealthChecker) findTargetByURL(url string) *loadbalance.Target {
|
||||
for _, t := range h.targets {
|
||||
if t.URL == url {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -207,27 +207,6 @@ func TestCheckTarget(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewHealthChecker_NilConfig 验证 nil 健康检查配置不会 panic。
|
||||
func TestNewHealthChecker_NilConfig(t *testing.T) {
|
||||
target := &loadbalance.Target{URL: "http://127.0.0.1:8080"}
|
||||
|
||||
// 不应 panic
|
||||
checker := NewHealthChecker([]*loadbalance.Target{target}, nil)
|
||||
if checker == nil {
|
||||
t.Fatal("NewHealthChecker(nil) should return non-nil checker")
|
||||
}
|
||||
|
||||
if checker.interval != 10*time.Second {
|
||||
t.Errorf("expected default interval 10s, got %v", checker.interval)
|
||||
}
|
||||
if checker.timeout != 5*time.Second {
|
||||
t.Errorf("expected default timeout 5s, got %v", checker.timeout)
|
||||
}
|
||||
if checker.path != healthPath {
|
||||
t.Errorf("expected default path %q, got %q", healthPath, checker.path)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMarkUnhealthy 测试 MarkUnhealthy 方法。
|
||||
func TestMarkUnhealthy(t *testing.T) {
|
||||
t.Run("标记不健康", func(t *testing.T) {
|
||||
|
||||
@ -42,7 +42,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
"golang.org/x/sync/singleflight"
|
||||
"rua.plus/lolly/internal/cache"
|
||||
"rua.plus/lolly/internal/config"
|
||||
"rua.plus/lolly/internal/loadbalance"
|
||||
@ -93,6 +92,16 @@ const (
|
||||
lbSticky = "sticky" // 会话粘性
|
||||
)
|
||||
|
||||
// headersPool 复用缓存 headers map,减少分配。
|
||||
// 预容量 20 覆盖大多数 HTTP 响应头数量。
|
||||
// 注意:从 pool 获取的 map 使用后不能 Put 回 pool,
|
||||
// 因为 cache.Set 存储了 map 引用。
|
||||
var headersPool = sync.Pool{
|
||||
New: func() any {
|
||||
return make(map[string]string, 20)
|
||||
},
|
||||
}
|
||||
|
||||
var upstreamTimingPool = sync.Pool{
|
||||
New: func() any {
|
||||
return NewUpstreamTiming()
|
||||
@ -122,7 +131,6 @@ type Proxy struct {
|
||||
mu sync.RWMutex // 保护并发访问的读写锁
|
||||
started atomic.Bool // 代理启动标志
|
||||
cacheIgnoreSet map[string]bool // 缓存时忽略的响应头集合
|
||||
refreshGroup singleflight.Group // 合并并发后台缓存刷新
|
||||
}
|
||||
|
||||
// NewProxy 使用给定的配置和后台目标创建一个新的反向代理实例。
|
||||
@ -221,39 +229,6 @@ func NewProxy(cfg *config.ProxyConfig, targets []*loadbalance.Target, transportC
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// cacheVaryHeaders 返回用于构建缓存键的 Vary 请求头列表。
|
||||
//
|
||||
// 当 varyResponse 非空时,将响应中的 Vary 头与配置的值合并去重。
|
||||
// 请求头名称统一按大小写不敏感方式比较,避免重复。
|
||||
func (p *Proxy) cacheVaryHeaders(varyResponse []byte) []string {
|
||||
configured := p.config.Cache.Vary
|
||||
if len(varyResponse) == 0 {
|
||||
return configured
|
||||
}
|
||||
|
||||
result := make([]string, 0, len(configured)+bytes.Count(varyResponse, []byte(","))+1)
|
||||
result = append(result, configured...)
|
||||
seen := make(map[string]struct{}, len(result))
|
||||
for _, h := range result {
|
||||
seen[strings.ToLower(h)] = struct{}{}
|
||||
}
|
||||
|
||||
for _, part := range strings.Split(b2s(varyResponse), ",") {
|
||||
name := strings.TrimSpace(part)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(name)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
result = append(result, name)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// stickyBalancer wraps StickySession to implement loadbalance.Balancer.
|
||||
// It delegates Select/SelectExcluding to the fallback balancer while
|
||||
// allowing the proxy to access the StickySession for cookie-based routing.
|
||||
@ -608,11 +583,15 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
}
|
||||
timing.reset()
|
||||
|
||||
// 在请求头被修改前捕获原始 Host,用于构建缓存键。
|
||||
// 防止 modifyRequestHeaders 将 Host 改为上游目标后导致缓存键丢失客户端 Host 区分。
|
||||
cacheHost := ctx.Request.Header.Host()
|
||||
computeCacheKey := func(varyHeaders []string) (uint64, string) {
|
||||
return p.buildCacheKeyHashWithHost(ctx, cacheHost, varyHeaders)
|
||||
var cacheHashKey uint64
|
||||
var cacheOrigKey string
|
||||
cacheKeyComputed := false
|
||||
computeCacheKey := func() (uint64, string) {
|
||||
if !cacheKeyComputed {
|
||||
cacheHashKey, cacheOrigKey = p.buildCacheKeyHash(ctx)
|
||||
cacheKeyComputed = true
|
||||
}
|
||||
return cacheHashKey, cacheOrigKey
|
||||
}
|
||||
|
||||
// 创建变量上下文用于设置上游变量
|
||||
@ -702,18 +681,18 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
|
||||
// 检查是否为 WebSocket 升级请求
|
||||
if isWebSocketRequest(ctx) {
|
||||
defer func() {
|
||||
loadbalance.DecrementConnections(target)
|
||||
}()
|
||||
defer timing.MarkConnectEnd()
|
||||
// WebSocket 使用 defer 确保连接计数释放
|
||||
defer loadbalance.DecrementConnections(target)
|
||||
timing.MarkConnectStart()
|
||||
err := WebSocket(ctx, target, p.config.Timeout.Connect, &p.config.Headers, p.config.ProxySSL)
|
||||
err := WebSocket(ctx, target, p.config.Timeout.Connect, &p.config.Headers)
|
||||
timing.MarkConnectEnd()
|
||||
if err != nil {
|
||||
upstreamStatus = fasthttp.StatusBadGateway
|
||||
logging.Error().Err(err).Msg("WebSocket proxy error")
|
||||
upstreamStatus = 502
|
||||
logging.Error().Msgf("WebSocket proxy error: %v", err)
|
||||
return
|
||||
}
|
||||
upstreamStatus = fasthttp.StatusSwitchingProtocols
|
||||
// WebSocket 成功
|
||||
upstreamStatus = 101
|
||||
return
|
||||
}
|
||||
|
||||
@ -770,7 +749,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
path := string(ctx.Request.URI().Path())
|
||||
rule := p.cache.MatchRule(path, method, 0)
|
||||
if rule != nil {
|
||||
hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(nil))
|
||||
hashKey, origKey := computeCacheKey()
|
||||
if entry, ok, stale := p.cache.Get(hashKey, origKey); ok {
|
||||
// 缓存命中
|
||||
loadbalance.DecrementConnections(target)
|
||||
@ -789,17 +768,10 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
entry.Updating.Store(true)
|
||||
reqCopy := fasthttp.AcquireRequest()
|
||||
ctx.Request.CopyTo(reqCopy)
|
||||
key := origKey
|
||||
go func() {
|
||||
defer entry.Updating.Store(false)
|
||||
defer fasthttp.ReleaseRequest(reqCopy)
|
||||
_, err, _ := p.refreshGroup.Do(key, func() (interface{}, error) {
|
||||
p.backgroundRefresh(reqCopy, target, hashKey, origKey)
|
||||
return nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
logging.Error().Err(err).Msg("background cache refresh failed")
|
||||
}
|
||||
p.backgroundRefresh(reqCopy, target, hashKey, origKey)
|
||||
}()
|
||||
}
|
||||
upstreamAddr = upstreamCache
|
||||
@ -868,7 +840,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
|
||||
// 尝试使用 stale 缓存
|
||||
if p.cache != nil {
|
||||
hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(nil))
|
||||
hashKey, origKey := computeCacheKey()
|
||||
isTimeout := errors.Is(err, fasthttp.ErrTimeout)
|
||||
if staleEntry, ok := p.cache.GetStale(hashKey, origKey, isTimeout); ok {
|
||||
logging.Info().Msgf("[PROXY] 使用 stale 缓存: key=%s, isTimeout=%v", origKey, isTimeout)
|
||||
@ -881,7 +853,8 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
|
||||
// 释放缓存锁
|
||||
if p.cache != nil && attempt == 0 {
|
||||
hashKey, _ := computeCacheKey(p.cacheVaryHeaders(nil))
|
||||
computeCacheKey()
|
||||
hashKey := cacheHashKey
|
||||
p.cache.ReleaseLock(hashKey, err)
|
||||
}
|
||||
|
||||
@ -931,7 +904,8 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
if shouldRetry {
|
||||
// 释放缓存锁
|
||||
if p.cache != nil && attempt == 0 {
|
||||
hashKey, _ := computeCacheKey(p.cacheVaryHeaders(nil))
|
||||
computeCacheKey()
|
||||
hashKey := cacheHashKey
|
||||
p.cache.ReleaseLock(hashKey, fmt.Errorf("HTTP %d", statusCode))
|
||||
}
|
||||
|
||||
@ -960,7 +934,7 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
return
|
||||
}
|
||||
|
||||
hashKey, origKey := computeCacheKey(p.cacheVaryHeaders(ctx.Response.Header.Peek("Vary")))
|
||||
hashKey, origKey := computeCacheKey()
|
||||
if statusCode >= 200 && statusCode < 300 {
|
||||
// 检查 MinUses 阈值
|
||||
if entry, ok, _ := p.cache.Get(hashKey, origKey); ok {
|
||||
@ -971,8 +945,14 @@ func (p *Proxy) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
}
|
||||
}
|
||||
|
||||
// 提取响应头
|
||||
headers := make(map[string]string, 20)
|
||||
// 提取响应头(使用 pool 复用 map)
|
||||
headers, ok := headersPool.Get().(map[string]string)
|
||||
if !ok {
|
||||
headers = make(map[string]string, 20)
|
||||
}
|
||||
for k := range headers {
|
||||
delete(headers, k)
|
||||
}
|
||||
|
||||
var lastModified, etag string
|
||||
for key, value := range ctx.Response.Header.All() {
|
||||
|
||||
@ -451,7 +451,7 @@ func BenchmarkBuildCacheKeyHash(b *testing.B) {
|
||||
b.Run("buildCacheKeyHash_with_string", func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
hashKey, _ := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, _ := p.buildCacheKeyHash(ctx)
|
||||
_ = hashKey
|
||||
}
|
||||
})
|
||||
@ -459,7 +459,7 @@ func BenchmarkBuildCacheKeyHash(b *testing.B) {
|
||||
b.Run("buildCacheKeyHashValue_direct", func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
hashKey := p.buildCacheKeyHashValue(ctx, nil)
|
||||
hashKey := p.buildCacheKeyHashValue(ctx)
|
||||
_ = hashKey
|
||||
}
|
||||
})
|
||||
@ -585,7 +585,7 @@ func BenchmarkProxyZeroAllocPath(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
hash := p.buildCacheKeyHashValue(ctx, nil)
|
||||
hash := p.buildCacheKeyHashValue(ctx)
|
||||
_ = hash
|
||||
}
|
||||
})
|
||||
@ -594,7 +594,7 @@ func BenchmarkProxyZeroAllocPath(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
hash, key := p.buildCacheKeyHash(ctx, nil)
|
||||
hash, key := p.buildCacheKeyHash(ctx)
|
||||
_ = hash
|
||||
_ = key
|
||||
}
|
||||
|
||||
@ -848,14 +848,14 @@ func TestSelectTarget_EmptyTargets(t *testing.T) {
|
||||
func TestDialTarget(t *testing.T) {
|
||||
t.Run("连接超时", func(t *testing.T) {
|
||||
// 使用不可达地址测试超时
|
||||
_, err := dialTarget("http://10.255.255.1:9999", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("http://10.255.255.1:9999", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("dialTarget() should return error for unreachable address")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("HTTPS 连接失败", func(t *testing.T) {
|
||||
_, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("dialTarget() should return error for unreachable HTTPS address")
|
||||
}
|
||||
@ -876,7 +876,7 @@ func TestDialTarget(t *testing.T) {
|
||||
}()
|
||||
|
||||
addr := ln.Addr().String()
|
||||
conn, err := dialTarget("http://"+addr, 1*time.Second, nil)
|
||||
conn, err := dialTarget("http://"+addr, 1*time.Second)
|
||||
if err != nil {
|
||||
t.Errorf("dialTarget() error: %v", err)
|
||||
}
|
||||
@ -972,7 +972,7 @@ func TestWebSocket_ErrorCases(t *testing.T) {
|
||||
target.Healthy.Store(true)
|
||||
|
||||
// 使用很短的超时
|
||||
err := WebSocket(ctx, target, 10*time.Millisecond, nil, nil)
|
||||
err := WebSocket(ctx, target, 10*time.Millisecond, nil)
|
||||
if err == nil {
|
||||
t.Error("WebSocket() should return error for invalid backend")
|
||||
}
|
||||
@ -983,7 +983,7 @@ func TestWebSocket_ErrorCases(t *testing.T) {
|
||||
func TestDialTarget_TLS_Extra(t *testing.T) {
|
||||
t.Run("TLS 握手失败", func(t *testing.T) {
|
||||
// 使用不可达的 HTTPS 地址
|
||||
_, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("https://10.255.255.1:9999", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("dialTarget() should return error for unreachable HTTPS address")
|
||||
}
|
||||
@ -1469,7 +1469,7 @@ func TestBackgroundRefresh_304(t *testing.T) {
|
||||
|
||||
// 预先设置缓存条目
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/test")
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
p.cache.Set(hashKey, origKey, []byte("cached"), map[string]string{
|
||||
"Last-Modified": "Tue, 20 Oct 2015 07:28:00 GMT",
|
||||
"ETag": "\"old\"",
|
||||
|
||||
@ -235,7 +235,7 @@ func TestBuildCacheKeyHash(t *testing.T) {
|
||||
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/test")
|
||||
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
if hashKey == 0 {
|
||||
t.Error("buildCacheKeyHash() should return non-zero hash")
|
||||
}
|
||||
@ -245,14 +245,14 @@ func TestBuildCacheKeyHash(t *testing.T) {
|
||||
|
||||
// 相同请求应产生相同哈希
|
||||
ctx2 := testutil.NewRequestCtx("GET", "/api/test")
|
||||
hashKey2, _ := p.buildCacheKeyHash(ctx2, nil)
|
||||
hashKey2, _ := p.buildCacheKeyHash(ctx2)
|
||||
if hashKey != hashKey2 {
|
||||
t.Error("Same request should produce same hash")
|
||||
}
|
||||
|
||||
// 不同请求应产生不同哈希
|
||||
ctx3 := testutil.NewRequestCtx("POST", "/api/other")
|
||||
hashKey3, _ := p.buildCacheKeyHash(ctx3, nil)
|
||||
hashKey3, _ := p.buildCacheKeyHash(ctx3)
|
||||
if hashKey == hashKey3 {
|
||||
t.Error("Different request should produce different hash")
|
||||
}
|
||||
@ -273,13 +273,13 @@ func TestBuildCacheKeyHashValue(t *testing.T) {
|
||||
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/test")
|
||||
|
||||
hashValue := p.buildCacheKeyHashValue(ctx, nil)
|
||||
hashValue := p.buildCacheKeyHashValue(ctx)
|
||||
if hashValue == 0 {
|
||||
t.Error("buildCacheKeyHashValue() should return non-zero hash")
|
||||
}
|
||||
|
||||
// 应该与 buildCacheKeyHash 结果一致
|
||||
hashKey, _ := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, _ := p.buildCacheKeyHash(ctx)
|
||||
if hashValue != hashKey {
|
||||
t.Error("buildCacheKeyHashValue() should match buildCacheKeyHash()")
|
||||
}
|
||||
@ -624,8 +624,7 @@ func TestServeHTTP_CacheHit(t *testing.T) {
|
||||
|
||||
// 预填充缓存
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/cached")
|
||||
ctx.Request.Header.SetHost("localhost:8080")
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
p.cache.Set(hashKey, origKey, []byte("cached!"), map[string]string{
|
||||
"Content-Type": "text/plain",
|
||||
}, 200, 10*time.Second)
|
||||
@ -706,8 +705,7 @@ func TestServeHTTP_WithRedirectRewrite_CacheHit(t *testing.T) {
|
||||
}
|
||||
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/test")
|
||||
ctx.Request.Header.SetHost("localhost:8080")
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
p.cache.Set(hashKey, origKey, []byte("ok"), map[string]string{
|
||||
"Content-Type": "text/plain",
|
||||
}, 200, 10*time.Second)
|
||||
|
||||
@ -23,7 +23,6 @@ import (
|
||||
"net/url"
|
||||
"time"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
"rua.plus/lolly/internal/loadbalance"
|
||||
"rua.plus/lolly/internal/logging"
|
||||
"rua.plus/lolly/internal/resolver"
|
||||
@ -142,10 +141,8 @@ func (p *Proxy) refreshDNS() {
|
||||
|
||||
// updateHostClientAddr 更新 HostClient 的连接地址。
|
||||
//
|
||||
// 从目标 URL 中解析出端口,与新的 IP 地址组合后创建新的 HostClient
|
||||
// 并替换连接池中的旧实例。旧连接不受影响,新连接将使用新地址。
|
||||
// 通过重建 client 而不是直接修改 Addr 字段,避免与正在使用旧 client
|
||||
// 的 goroutine 产生数据竞争。
|
||||
// 从目标 URL 中解析出端口,与新的 IP 地址组合后更新对应
|
||||
// HostClient 的 Addr 字段。旧连接不受影响,新连接将使用新地址。
|
||||
//
|
||||
// 参数:
|
||||
// - target: 负载均衡目标,包含原始 URL
|
||||
@ -172,49 +169,12 @@ func (p *Proxy) updateHostClientAddr(target *loadbalance.Target, ip string) {
|
||||
|
||||
newAddr := net.JoinHostPort(ip, port)
|
||||
|
||||
// 与 getClient 保持一致的 key 计算逻辑
|
||||
key := target.URL
|
||||
if p.config.ProxyBind != "" {
|
||||
key = target.URL + "|" + p.config.ProxyBind
|
||||
// 更新 HostClient 的 Addr
|
||||
// 注意:新连接将使用新 IP,旧连接继续使用旧 IP 直到超时
|
||||
if client, ok := p.clients[target.URL]; ok {
|
||||
client.Addr = newAddr
|
||||
logging.Debug().Msgf("Updated HostClient addr for %s to %s", target.URL, newAddr)
|
||||
}
|
||||
|
||||
oldClient, ok := p.clients[key]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// 创建新的 HostClient,复制旧 client 的配置,仅修改 Addr
|
||||
// 新连接将使用新 IP,旧连接继续使用旧 IP 直到超时
|
||||
newClient := &fasthttp.HostClient{
|
||||
Addr: newAddr,
|
||||
IsTLS: oldClient.IsTLS,
|
||||
ReadTimeout: oldClient.ReadTimeout,
|
||||
WriteTimeout: oldClient.WriteTimeout,
|
||||
MaxIdleConnDuration: oldClient.MaxIdleConnDuration,
|
||||
MaxConns: oldClient.MaxConns,
|
||||
MaxConnWaitTimeout: oldClient.MaxConnWaitTimeout,
|
||||
DisablePathNormalizing: oldClient.DisablePathNormalizing,
|
||||
SecureErrorLogMessage: oldClient.SecureErrorLogMessage,
|
||||
Dial: oldClient.Dial,
|
||||
StreamResponseBody: oldClient.StreamResponseBody,
|
||||
ReadBufferSize: oldClient.ReadBufferSize,
|
||||
WriteBufferSize: oldClient.WriteBufferSize,
|
||||
TLSConfig: oldClient.TLSConfig,
|
||||
}
|
||||
|
||||
// 兼容旧版 RetryIf 配置,转换为 RetryIfErr。
|
||||
// RetryIf 已被 fasthttp 标记为 deprecated,新连接使用 RetryIfErr 替代。
|
||||
//nolint:staticcheck // 需要读取旧 client 的 deprecated RetryIf 以兼容已有配置。
|
||||
if oldClient.RetryIf != nil {
|
||||
//nolint:staticcheck // 同上,读取旧 client 的 deprecated RetryIf。
|
||||
oldRetryIf := oldClient.RetryIf
|
||||
newClient.RetryIfErr = func(req *fasthttp.Request, _ int, _ error) (bool, bool) {
|
||||
return false, oldRetryIf(req)
|
||||
}
|
||||
}
|
||||
|
||||
p.clients[key] = newClient
|
||||
logging.Debug().Msgf("Updated HostClient addr for %s to %s", target.URL, newAddr)
|
||||
}
|
||||
|
||||
// getResolverTTL 获取 DNS 解析记录的过期时间。
|
||||
|
||||
@ -161,116 +161,4 @@ func TestSetResolver(t *testing.T) {
|
||||
|
||||
// TestMultipleTargets_Refresh 测试多目标刷新。
|
||||
|
||||
// TestUpdateHostClientAddr_ReplacesClient 验证 DNS 更新时重建 HostClient 而不是修改 Addr。
|
||||
func TestUpdateHostClientAddr_ReplacesClient(t *testing.T) {
|
||||
cfg := &config.ProxyConfig{
|
||||
Path: "/api",
|
||||
LoadBalance: "round_robin",
|
||||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||||
}
|
||||
targets := []*loadbalance.Target{
|
||||
{URL: "http://backend.example.com:8080"},
|
||||
}
|
||||
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewProxy() error: %v", err)
|
||||
}
|
||||
|
||||
oldClient := p.clients["http://backend.example.com:8080"]
|
||||
if oldClient == nil {
|
||||
t.Fatal("old client should not be nil")
|
||||
}
|
||||
oldAddr := oldClient.Addr
|
||||
|
||||
p.updateHostClientAddr(targets[0], "192.168.1.100")
|
||||
|
||||
newClient := p.clients["http://backend.example.com:8080"]
|
||||
if newClient == nil {
|
||||
t.Fatal("new client should not be nil")
|
||||
}
|
||||
if newClient == oldClient {
|
||||
t.Error("updateHostClientAddr should replace the HostClient instead of mutating it")
|
||||
}
|
||||
if newClient.Addr != "192.168.1.100:8080" {
|
||||
t.Errorf("new client addr = %q, want %q", newClient.Addr, "192.168.1.100:8080")
|
||||
}
|
||||
// 旧 client 的 Addr 不应被修改,旧连接继续使用
|
||||
if oldClient.Addr != oldAddr {
|
||||
t.Errorf("old client addr was mutated: got %q, want %q", oldClient.Addr, oldAddr)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateHostClientAddr_WithProxyBind 验证配置了 proxy_bind 时使用正确的 client key。
|
||||
func TestUpdateHostClientAddr_WithProxyBind(t *testing.T) {
|
||||
cfg := &config.ProxyConfig{
|
||||
Path: "/api",
|
||||
LoadBalance: "round_robin",
|
||||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||||
ProxyBind: "127.0.0.1",
|
||||
}
|
||||
targets := []*loadbalance.Target{
|
||||
{URL: "http://backend.example.com:8080"},
|
||||
}
|
||||
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewProxy() error: %v", err)
|
||||
}
|
||||
|
||||
key := "http://backend.example.com:8080|127.0.0.1"
|
||||
if p.clients[key] == nil {
|
||||
t.Fatalf("client with proxy_bind key should exist")
|
||||
}
|
||||
|
||||
p.updateHostClientAddr(targets[0], "192.168.1.200")
|
||||
|
||||
newClient := p.clients[key]
|
||||
if newClient == nil {
|
||||
t.Fatal("new client should not be nil")
|
||||
}
|
||||
if newClient.Addr != "192.168.1.200:8080" {
|
||||
t.Errorf("new client addr = %q, want %q", newClient.Addr, "192.168.1.200:8080")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateHostClientAddr_ConcurrentRead 验证更新 Addr 时不会与 getClient 产生数据竞争。
|
||||
func TestUpdateHostClientAddr_ConcurrentRead(t *testing.T) {
|
||||
cfg := &config.ProxyConfig{
|
||||
Path: "/api",
|
||||
LoadBalance: "round_robin",
|
||||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||||
}
|
||||
targets := []*loadbalance.Target{
|
||||
{URL: "http://backend.example.com:8080"},
|
||||
}
|
||||
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewProxy() error: %v", err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
// 并发读取 client
|
||||
for range 10 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
client := p.getClient(targets[0].URL)
|
||||
_ = client.Addr
|
||||
}()
|
||||
}
|
||||
|
||||
// 并发更新 client
|
||||
for i := range 10 {
|
||||
wg.Add(1)
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
p.updateHostClientAddr(targets[0], "192.168.1."+string(rune('0'+id)))
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// TestStopResolverFails 测试停止解析器失败时返回错误。
|
||||
|
||||
@ -654,7 +654,7 @@ func TestBackgroundRefresh_WithCacheEntry(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/test")
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
|
||||
p.cache.Set(hashKey, origKey, []byte("old content"), map[string]string{
|
||||
"Content-Type": "text/plain",
|
||||
@ -712,7 +712,7 @@ func TestDialTarget_Success(t *testing.T) {
|
||||
}
|
||||
}()
|
||||
|
||||
conn, err := dialTarget("http://"+ln.Addr().String(), 1*time.Second, nil)
|
||||
conn, err := dialTarget("http://"+ln.Addr().String(), 1*time.Second)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, conn)
|
||||
_ = conn.Close()
|
||||
@ -720,7 +720,7 @@ func TestDialTarget_Success(t *testing.T) {
|
||||
|
||||
// TestDialTarget_Timeout 测试连接超时。
|
||||
func TestDialTarget_Timeout(t *testing.T) {
|
||||
_, err := dialTarget("http://10.255.255.1:9999", 50*time.Millisecond, nil)
|
||||
_, err := dialTarget("http://10.255.255.1:9999", 50*time.Millisecond)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "failed to connect")
|
||||
}
|
||||
@ -770,7 +770,7 @@ func TestWebSocket_BackendSuccess(t *testing.T) {
|
||||
target := &loadbalance.Target{URL: "http://127.0.0.1:1"}
|
||||
target.Healthy.Store(true)
|
||||
|
||||
err := WebSocket(ctx, target, 2*time.Second, nil, nil)
|
||||
err := WebSocket(ctx, target, 2*time.Second, nil)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "hijack")
|
||||
}
|
||||
@ -1068,7 +1068,7 @@ func TestServeHTTP_CacheStaleWhileRevalidate(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx := testutil.NewRequestCtx("GET", "/api/cached")
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx, nil)
|
||||
hashKey, origKey := p.buildCacheKeyHash(ctx)
|
||||
|
||||
headers := map[string]string{"Content-Type": "text/plain"}
|
||||
p.cache.Set(hashKey, origKey, []byte("stale content"), headers, 200, -1*time.Second)
|
||||
@ -1399,7 +1399,7 @@ func TestWebSocket_ReadResponseError(t *testing.T) {
|
||||
target := &loadbalance.Target{URL: "http://" + ln.Addr().String()}
|
||||
target.Healthy.Store(true)
|
||||
|
||||
err = WebSocket(ctx, target, 1*time.Second, nil, nil)
|
||||
err = WebSocket(ctx, target, 1*time.Second, nil)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
@ -1434,7 +1434,7 @@ func TestWebSocket_HijackFails(t *testing.T) {
|
||||
target := &loadbalance.Target{URL: "http://127.0.0.1:1"}
|
||||
target.Healthy.Store(true)
|
||||
|
||||
err := WebSocket(ctx, target, 100*time.Millisecond, nil, nil)
|
||||
err := WebSocket(ctx, target, 100*time.Millisecond, nil)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
|
||||
@ -20,7 +20,6 @@ package proxy
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -774,31 +773,6 @@ func TestHandleWebSocket(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestWebSocketConnectionCountReleasedOnFailure 测试 WebSocket 升级失败后连接计数被释放。
|
||||
// testutil 创建的请求上下文不支持 Hijack,会导致 WebSocket 代理失败;
|
||||
// 失败后应确保 IncrementConnections 对应的 DecrementConnections 被调用。
|
||||
func TestWebSocketConnectionCountReleasedOnFailure(t *testing.T) {
|
||||
cfg := testutil.NewTestProxyConfig("/ws")
|
||||
targets := testutil.NewTestTargets("http://127.0.0.1:1")
|
||||
targets[0].MaxFails = 1
|
||||
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewProxy() error: %v", err)
|
||||
}
|
||||
|
||||
ctx := testutil.NewRequestCtxWithHeader("GET", "/ws", map[string]string{
|
||||
"Upgrade": "websocket",
|
||||
"Connection": "Upgrade",
|
||||
})
|
||||
|
||||
p.ServeHTTP(ctx)
|
||||
|
||||
if got := atomic.LoadInt64(&targets[0].Connections); got != 0 {
|
||||
t.Fatalf("expected connection count 0 after failed WebSocket, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSetHealthChecker 测试健康检查器设置
|
||||
// 注意:SetHealthChecker 是公开方法,但 healthChecker 是私有字段
|
||||
// 此测试验证方法可以正常调用
|
||||
@ -1151,87 +1125,3 @@ func TestUpstreamTiming_PartialMarks(t *testing.T) {
|
||||
t.Errorf("GetResponseTime() without connectEnd = %v, want 0", timing.GetResponseTime())
|
||||
}
|
||||
}
|
||||
|
||||
// TestProxyConnectionCount_InternalRedirect 验证 X-Accel-Redirect 后连接计数归零。
|
||||
func TestProxyConnectionCount_InternalRedirect(t *testing.T) {
|
||||
backend := &fasthttp.Server{
|
||||
Handler: func(ctx *fasthttp.RequestCtx) {
|
||||
ctx.Response.Header.Set("X-Accel-Redirect", "/redirect/new-path")
|
||||
ctx.SetStatusCode(200)
|
||||
},
|
||||
}
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create listener: %v", err)
|
||||
}
|
||||
defer func() { _ = ln.Close() }()
|
||||
|
||||
go func() { _ = backend.Serve(ln) }()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
targets := testutil.NewTestHealthyTargets("http://" + ln.Addr().String())
|
||||
cfg := testutil.NewTestProxyConfig("/")
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create proxy: %v", err)
|
||||
}
|
||||
|
||||
ctx := testutil.NewRequestCtxWithHeader("GET", "/api", map[string]string{
|
||||
"Host": "example.com",
|
||||
})
|
||||
|
||||
p.ServeHTTP(ctx)
|
||||
|
||||
if atomic.LoadInt64(&targets[0].Connections) != 0 {
|
||||
t.Errorf("expected connections to return to 0 after internal redirect, got %d", targets[0].Connections)
|
||||
}
|
||||
}
|
||||
|
||||
// TestProxyConnectionCount_Retry 验证重试路径后连接计数归零。
|
||||
func TestProxyConnectionCount_Retry(t *testing.T) {
|
||||
backend := &fasthttp.Server{
|
||||
Handler: func(ctx *fasthttp.RequestCtx) {
|
||||
ctx.SetStatusCode(503)
|
||||
},
|
||||
}
|
||||
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create listener: %v", err)
|
||||
}
|
||||
defer func() { _ = ln.Close() }()
|
||||
|
||||
go func() { _ = backend.Serve(ln) }()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
targets := testutil.NewTestHealthyTargets("http://" + ln.Addr().String())
|
||||
cfg := &config.ProxyConfig{
|
||||
Path: "/",
|
||||
LoadBalance: "round_robin",
|
||||
Timeout: config.ProxyTimeout{
|
||||
Connect: 5 * time.Second,
|
||||
Read: 5 * time.Second,
|
||||
Write: 5 * time.Second,
|
||||
},
|
||||
NextUpstream: config.NextUpstreamConfig{
|
||||
Tries: 2,
|
||||
HTTPCodes: []int{503},
|
||||
},
|
||||
}
|
||||
|
||||
p, err := NewProxy(cfg, targets, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create proxy: %v", err)
|
||||
}
|
||||
|
||||
ctx := testutil.NewRequestCtxWithHeader("GET", "/api", map[string]string{
|
||||
"Host": "example.com",
|
||||
})
|
||||
|
||||
p.ServeHTTP(ctx)
|
||||
|
||||
if atomic.LoadInt64(&targets[0].Connections) != 0 {
|
||||
t.Errorf("expected connections to return to 0 after retry, got %d", targets[0].Connections)
|
||||
}
|
||||
}
|
||||
|
||||
@ -219,7 +219,7 @@ func isConnectionClosedError(err error) bool {
|
||||
// 返回值:
|
||||
// - net.Conn: 建立的连接(TLS 连接或普通 TCP 连接)
|
||||
// - error: 连接失败时返回错误
|
||||
func dialTarget(targetURL string, timeout time.Duration, proxySSL *config.ProxySSLConfig) (net.Conn, error) {
|
||||
func dialTarget(targetURL string, timeout time.Duration) (net.Conn, error) {
|
||||
// 解析目标 URL
|
||||
addr, isTLS := netutil.ParseTargetURL(targetURL, true)
|
||||
|
||||
@ -233,34 +233,24 @@ func dialTarget(targetURL string, timeout time.Duration, proxySSL *config.ProxyS
|
||||
return nil, fmt.Errorf("failed to connect to target: %w", err)
|
||||
}
|
||||
|
||||
if !isTLS {
|
||||
return conn, nil
|
||||
// 如果是 HTTPS,建立 TLS 连接
|
||||
if isTLS {
|
||||
tlsConn := tls.Client(conn, &tls.Config{
|
||||
InsecureSkipVerify: false,
|
||||
ServerName: strings.Split(addr, ":")[0],
|
||||
})
|
||||
if err := tlsConn.SetDeadline(time.Now().Add(timeout)); err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, fmt.Errorf("failed to set TLS deadline: %w", err)
|
||||
}
|
||||
if err := tlsConn.Handshake(); err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, fmt.Errorf("TLS handshake failed: %w", err)
|
||||
}
|
||||
return tlsConn, nil
|
||||
}
|
||||
|
||||
tlsCfg, err := CreateTLSConfig(proxySSL, addr)
|
||||
if err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
host = addr
|
||||
}
|
||||
if tlsCfg.ServerName == "" {
|
||||
tlsCfg.ServerName = host
|
||||
}
|
||||
|
||||
tlsConn := tls.Client(conn, tlsCfg)
|
||||
if err := tlsConn.SetDeadline(time.Now().Add(timeout)); err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, err
|
||||
}
|
||||
if err := tlsConn.Handshake(); err != nil {
|
||||
_ = conn.Close()
|
||||
return nil, fmt.Errorf("TLS handshake failed: %w", err)
|
||||
}
|
||||
return tlsConn, nil
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// buildWebSocketUpgradeRequest 构建 WebSocket 升级 HTTP 请求。
|
||||
@ -375,7 +365,7 @@ func readWebSocketUpgradeResponse(conn net.Conn, timeout time.Duration) (*http.R
|
||||
//
|
||||
// 返回值:
|
||||
// - error: 代理过程中的错误
|
||||
func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout time.Duration, headersConfig *config.ProxyHeaders, proxySSL *config.ProxySSLConfig) error {
|
||||
func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout time.Duration, headersConfig *config.ProxyHeaders) error {
|
||||
// 使用 Hijack 获取客户端 TCP 连接
|
||||
var clientConn net.Conn
|
||||
|
||||
@ -388,7 +378,7 @@ func WebSocket(ctx *fasthttp.RequestCtx, target *loadbalance.Target, timeout tim
|
||||
}
|
||||
|
||||
// 步骤1: 建立到后端目标的连接
|
||||
targetConn, err := dialTarget(target.URL, timeout, proxySSL)
|
||||
targetConn, err := dialTarget(target.URL, timeout)
|
||||
if err != nil {
|
||||
_ = clientConn.Close()
|
||||
return fmt.Errorf("failed to connect to backend: %w", err)
|
||||
|
||||
@ -21,7 +21,6 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@ -136,7 +135,7 @@ func TestIsConnectionClosedError(t *testing.T) {
|
||||
// TestDialTarget_InvalidAddress 测试无效地址的拨号
|
||||
func TestDialTarget_InvalidAddress(t *testing.T) {
|
||||
// 测试连接到无效端口
|
||||
_, err := dialTarget("http://127.0.0.1:1", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("http://127.0.0.1:1", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid address")
|
||||
}
|
||||
@ -145,7 +144,7 @@ func TestDialTarget_InvalidAddress(t *testing.T) {
|
||||
// TestDialTarget_HTTPS 测试 HTTPS 连接(会失败,但验证错误处理)
|
||||
func TestDialTarget_HTTPS(t *testing.T) {
|
||||
// 测试 HTTPS 连接到无效端口
|
||||
_, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid HTTPS address")
|
||||
}
|
||||
@ -258,7 +257,7 @@ func TestDialTarget_URLParsing(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := dialTarget(tt.url, 10*time.Millisecond, nil)
|
||||
_, err := dialTarget(tt.url, 10*time.Millisecond)
|
||||
if tt.expectError {
|
||||
if err == nil {
|
||||
t.Error("Expected error, got nil")
|
||||
@ -589,7 +588,7 @@ func TestReadWebSocketUpgradeResponse_Timeout(t *testing.T) {
|
||||
// TestDialTarget_TLS 测试 TLS 连接(连接无效端口应失败)
|
||||
func TestDialTarget_TLS(t *testing.T) {
|
||||
// 测试 HTTPS 连接到无效端口
|
||||
_, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond, nil)
|
||||
_, err := dialTarget("https://127.0.0.1:1", 100*time.Millisecond)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid HTTPS address")
|
||||
}
|
||||
@ -783,45 +782,3 @@ func TestCopyData_WriteError(t *testing.T) {
|
||||
|
||||
_ = src1.Close()
|
||||
}
|
||||
|
||||
// TestReadWebSocketUpgradeResponse_BuffersTrailingData 验证升级响应后紧跟的帧数据不会被丢弃。
|
||||
func TestReadWebSocketUpgradeResponse_BuffersTrailingData(t *testing.T) {
|
||||
conn1, conn2 := net.Pipe()
|
||||
defer func() { _ = conn1.Close() }()
|
||||
defer func() { _ = conn2.Close() }()
|
||||
|
||||
trailing := []byte("first websocket frame data")
|
||||
go func() {
|
||||
response := "HTTP/1.1 101 Switching Protocols\r\n" +
|
||||
"Upgrade: websocket\r\n" +
|
||||
"Connection: Upgrade\r\n" +
|
||||
"\r\n"
|
||||
// 一次性写入响应头和后续帧数据,模拟真实后端在同一 TCP 段发送的场景
|
||||
_, _ = conn2.Write(append([]byte(response), trailing...))
|
||||
}()
|
||||
|
||||
resp, reader, err := readWebSocketUpgradeResponse(conn1, 1*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("readWebSocketUpgradeResponse failed: %v", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusSwitchingProtocols {
|
||||
t.Errorf("expected 101, got %d", resp.StatusCode)
|
||||
}
|
||||
if reader == nil {
|
||||
t.Fatal("expected buffered reader")
|
||||
}
|
||||
if reader.Buffered() != len(trailing) {
|
||||
t.Errorf("expected %d buffered trailing bytes, got %d", len(trailing), reader.Buffered())
|
||||
}
|
||||
|
||||
// 使用 bufferedConn 消费缓冲区数据
|
||||
bc := &bufferedConn{Conn: conn1, reader: reader}
|
||||
buf := make([]byte, len(trailing))
|
||||
n, err := io.ReadFull(bc, buf)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to read buffered trailing data: %v", err)
|
||||
}
|
||||
if !bytes.Equal(buf[:n], trailing) {
|
||||
t.Errorf("trailing data mismatch: expected %q, got %q", trailing, buf[:n])
|
||||
}
|
||||
}
|
||||
|
||||
@ -357,20 +357,21 @@ func (r *DNSResolver) Start() error {
|
||||
|
||||
r.started.Store(true)
|
||||
|
||||
// 重新创建 stopCh,确保新一轮刷新循环使用未关闭的通道
|
||||
r.mu.Lock()
|
||||
r.stopCh = make(chan struct{})
|
||||
stopCh := r.stopCh
|
||||
r.mu.Unlock()
|
||||
// 重建 stopCh 以支持 Start-Stop-Start 周期
|
||||
select {
|
||||
case <-r.stopCh:
|
||||
r.stopCh = make(chan struct{})
|
||||
default:
|
||||
}
|
||||
|
||||
// 启动后台刷新协程
|
||||
go r.refreshLoop(stopCh)
|
||||
go r.refreshLoop()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// refreshLoop 后台刷新循环。
|
||||
func (r *DNSResolver) refreshLoop(stopCh <-chan struct{}) {
|
||||
func (r *DNSResolver) refreshLoop() {
|
||||
// 刷新间隔为 TTL / 2
|
||||
interval := max(r.config.TTL()/2, time.Second)
|
||||
|
||||
@ -381,7 +382,7 @@ func (r *DNSResolver) refreshLoop(stopCh <-chan struct{}) {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
r.doRefresh()
|
||||
case <-stopCh:
|
||||
case <-r.stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@ -10,17 +10,6 @@ import (
|
||||
"rua.plus/lolly/internal/config"
|
||||
)
|
||||
|
||||
// TestNewResolver_NilConfig 验证 nil 解析器配置不会 panic。
|
||||
func TestNewResolver_NilConfig(t *testing.T) {
|
||||
r := New(nil)
|
||||
if r == nil {
|
||||
t.Fatal("New(nil) should return non-nil resolver")
|
||||
}
|
||||
if _, ok := r.(*noopResolver); !ok {
|
||||
t.Error("New(nil) should return *noopResolver")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewResolver 测试解析器创建。
|
||||
func TestNewResolver(t *testing.T) {
|
||||
// 测试启用状态
|
||||
@ -330,49 +319,6 @@ func TestStartStop(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestStartStopStart 验证 resolver 在 Stop 后可以重新 Start。
|
||||
func TestStartStopStart(t *testing.T) {
|
||||
cfg := &config.ResolverConfig{
|
||||
Enabled: true,
|
||||
Addresses: []string{"8.8.8.8:53"},
|
||||
Valid: 30 * time.Second,
|
||||
}
|
||||
|
||||
r := New(cfg).(*DNSResolver)
|
||||
|
||||
// 第一次启动
|
||||
if err := r.Start(); err != nil {
|
||||
t.Fatalf("Start failed: %v", err)
|
||||
}
|
||||
if !r.started.Load() {
|
||||
t.Error("resolver should be started")
|
||||
}
|
||||
|
||||
// 停止
|
||||
if err := r.Stop(); err != nil {
|
||||
t.Fatalf("Stop failed: %v", err)
|
||||
}
|
||||
if r.started.Load() {
|
||||
t.Error("resolver should be stopped")
|
||||
}
|
||||
|
||||
// 再次启动:验证 stopCh 被正确重建,refreshLoop 不会立即退出
|
||||
if err := r.Start(); err != nil {
|
||||
t.Fatalf("second Start failed: %v", err)
|
||||
}
|
||||
if !r.started.Load() {
|
||||
t.Error("resolver should be started again")
|
||||
}
|
||||
|
||||
// 等待一段时间,确认 refreshLoop 仍在运行
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
if !r.started.Load() {
|
||||
t.Error("resolver should still be running after sleep")
|
||||
}
|
||||
|
||||
_ = r.Stop()
|
||||
}
|
||||
|
||||
// TestDeleteCacheEntry 测试删除缓存条目。
|
||||
func TestDeleteCacheEntry(t *testing.T) {
|
||||
cfg := &config.ResolverConfig{
|
||||
|
||||
@ -97,6 +97,8 @@ func TestWrapRoutedHandler_有AccessLog(t *testing.T) {
|
||||
}
|
||||
s := New(cfg)
|
||||
|
||||
s.accessLogMiddleware = s.accessLogMiddleware
|
||||
|
||||
called := false
|
||||
original := func(ctx *fasthttp.RequestCtx) {
|
||||
called = true
|
||||
|
||||
@ -43,14 +43,6 @@ func (s *Server) cleanupResources() {
|
||||
s.tlsManager.Close()
|
||||
}
|
||||
|
||||
// 关闭多服务器模式下创建的子 TLS 管理器
|
||||
s.tlsManagersMu.Lock()
|
||||
for _, m := range s.tlsManagers {
|
||||
m.Close()
|
||||
}
|
||||
s.tlsManagers = nil
|
||||
s.tlsManagersMu.Unlock()
|
||||
|
||||
// 关闭 AccessControl (释放 GeoIP 资源)
|
||||
s.accessControlsMu.Lock()
|
||||
for _, ac := range s.accessControls {
|
||||
@ -60,9 +52,6 @@ func (s *Server) cleanupResources() {
|
||||
}
|
||||
s.accessControlsMu.Unlock()
|
||||
|
||||
// 停止 RateLimiter 的后台清理 goroutine
|
||||
s.stopRateLimiters()
|
||||
|
||||
// 关闭 Lua 引擎
|
||||
if s.luaEngine != nil {
|
||||
s.luaEngine.Close()
|
||||
|
||||
@ -1,66 +0,0 @@
|
||||
// Package server 提供服务器生命周期相关功能的测试。
|
||||
//
|
||||
// 该文件测试生命周期管理中的并发安全场景,包括:
|
||||
// - 代理缓存统计与代理创建/清理的并发访问
|
||||
//
|
||||
// 作者:xfy
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"rua.plus/lolly/internal/config"
|
||||
"rua.plus/lolly/internal/proxy"
|
||||
"rua.plus/lolly/internal/testutil"
|
||||
)
|
||||
|
||||
// TestGetProxyCacheStats_Concurrent 验证并发统计代理缓存与修改 proxies 切片不产生竞态。
|
||||
func TestGetProxyCacheStats_Concurrent(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Servers: []config.ServerConfig{{Listen: ":0"}},
|
||||
}
|
||||
s := New(cfg)
|
||||
|
||||
proxyCfg := &config.ProxyConfig{
|
||||
Path: "/api",
|
||||
LoadBalance: "round_robin",
|
||||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||||
Cache: config.ProxyCacheConfig{
|
||||
Enabled: true,
|
||||
MaxAge: 10 * time.Second,
|
||||
},
|
||||
}
|
||||
targets := testutil.NewTestTargets("http://localhost:8080")
|
||||
p, err := proxy.NewProxy(proxyCfg, targets, nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
s.proxies = []*proxy.Proxy{p}
|
||||
|
||||
purgeHandler := &PurgeHandler{server: s}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 100 {
|
||||
wg.Add(3)
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = s.getProxyCacheStats()
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = purgeHandler.purgeByPath("/api/test", "GET")
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
s.proxiesMu.Lock()
|
||||
s.proxies = append(s.proxies, p)
|
||||
s.proxiesMu.Unlock()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
@ -61,15 +61,11 @@ func (s *Server) buildMiddlewareChain(serverCfg *config.ServerConfig) (*middlewa
|
||||
|
||||
// 3. Security: RateLimiter (速率限制)
|
||||
if serverCfg.Security.RateLimit.RequestRate > 0 {
|
||||
rl, err := security.NewRateLimiter(&serverCfg.Security.RateLimit, serverCfg.Security.Access.TrustedProxies...)
|
||||
rl, err := security.NewRateLimiter(&serverCfg.Security.RateLimit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create rate limiter middleware: %w", err)
|
||||
}
|
||||
middlewares = append(middlewares, rl)
|
||||
// 跟踪 token_bucket 限流器的清理 goroutine,在服务器停止/重载时释放
|
||||
if tokenRl, ok := rl.(*security.RateLimiter); ok {
|
||||
s.trackRateLimiter(tokenRl)
|
||||
}
|
||||
}
|
||||
|
||||
// 3.5 Security: ConnLimiter (连接数限制)
|
||||
|
||||
@ -450,44 +450,3 @@ func TestPoolSubmit_StartWorkerWhenNoIdle(t *testing.T) {
|
||||
|
||||
close(blockCh)
|
||||
}
|
||||
|
||||
// TestPoolSubmit_ContinuousFullLoad 验证队列满载时持续 Submit 不会阻塞。
|
||||
func TestPoolSubmit_ContinuousFullLoad(t *testing.T) {
|
||||
p := NewGoroutinePool(PoolConfig{
|
||||
MaxWorkers: 2,
|
||||
MinWorkers: 1,
|
||||
QueueSize: 1,
|
||||
IdleTimeout: 5 * time.Second,
|
||||
})
|
||||
|
||||
p.Start()
|
||||
defer p.Stop()
|
||||
|
||||
// 阻塞唯一 worker,使队列保持满载
|
||||
blockCh := make(chan struct{})
|
||||
started := make(chan struct{})
|
||||
_ = p.Submit(nil, func(*fasthttp.RequestCtx) {
|
||||
close(started)
|
||||
<-blockCh
|
||||
})
|
||||
<-started
|
||||
|
||||
// 填满队列
|
||||
_ = p.Submit(nil, func(*fasthttp.RequestCtx) {})
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
for range 100 {
|
||||
_ = p.Submit(nil, func(*fasthttp.RequestCtx) {})
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("continuous Submit under full load blocked")
|
||||
}
|
||||
|
||||
close(blockCh)
|
||||
}
|
||||
|
||||
@ -80,7 +80,7 @@ func (h *PurgeHandler) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
}
|
||||
|
||||
// 检查 IP 访问权限
|
||||
if !utils.CheckIPAccess(ctx, h.allowed, nil) {
|
||||
if !utils.CheckIPAccess(ctx, h.allowed) {
|
||||
utils.SendJSONError(ctx, fasthttp.StatusForbidden, "forbidden")
|
||||
return
|
||||
}
|
||||
|
||||
@ -236,7 +236,7 @@ func TestPurgeHandler_checkAccess(t *testing.T) {
|
||||
|
||||
if len(h.allowed) == 0 {
|
||||
// 无白名单时应允许所有访问
|
||||
if !utils.CheckIPAccess(nil, h.allowed, nil) {
|
||||
if !utils.CheckIPAccess(nil, h.allowed) {
|
||||
t.Error("expected access to be true when no allow list configured")
|
||||
}
|
||||
return
|
||||
@ -654,7 +654,7 @@ func TestPurgeHandler_checkAccess_NilContext(t *testing.T) {
|
||||
}
|
||||
|
||||
// Empty allow list should allow access (returns true even with nil context)
|
||||
if !utils.CheckIPAccess(nil, h.allowed, nil) {
|
||||
if !utils.CheckIPAccess(nil, h.allowed) {
|
||||
t.Error("expected checkAccess to return true with empty allow list")
|
||||
}
|
||||
})
|
||||
@ -705,7 +705,7 @@ func TestPurgeHandler_checkAccess_WithAllowedIP(t *testing.T) {
|
||||
ctx.Init(&fasthttp.Request{}, nil, nil)
|
||||
|
||||
// context with nil remote address - should return false (no client IP)
|
||||
if utils.CheckIPAccess(ctx, h.allowed, nil) {
|
||||
if utils.CheckIPAccess(ctx, h.allowed) {
|
||||
t.Error("expected checkAccess to return false with no client IP")
|
||||
}
|
||||
})
|
||||
|
||||
@ -63,15 +63,11 @@ type Server struct {
|
||||
handler fasthttp.RequestHandler
|
||||
resolver resolver.Resolver
|
||||
tlsManager *ssl.TLSManager
|
||||
tlsManagers []*ssl.TLSManager
|
||||
tlsManagersMu sync.Mutex
|
||||
accessLogMiddleware *accesslog.AccessLog
|
||||
luaEngine *lua.LuaEngine
|
||||
accessControl *security.AccessControl
|
||||
accessControls []*security.AccessControl
|
||||
accessControlsMu sync.Mutex
|
||||
rateLimiters []*security.RateLimiter
|
||||
rateLimitersMu sync.Mutex
|
||||
errorPageManager *handler.ErrorPageManager
|
||||
fileCache *cache.FileCache
|
||||
pool *GoroutinePool
|
||||
@ -115,23 +111,6 @@ func (s *Server) Running() bool {
|
||||
return s.running.Load()
|
||||
}
|
||||
|
||||
// trackRateLimiter 记录创建的令牌桶限流器,便于后续统一释放后台清理 goroutine。
|
||||
func (s *Server) trackRateLimiter(rl *security.RateLimiter) {
|
||||
s.rateLimitersMu.Lock()
|
||||
defer s.rateLimitersMu.Unlock()
|
||||
s.rateLimiters = append(s.rateLimiters, rl)
|
||||
}
|
||||
|
||||
// stopRateLimiters 停止所有已记录的令牌桶限流器。
|
||||
func (s *Server) stopRateLimiters() {
|
||||
s.rateLimitersMu.Lock()
|
||||
defer s.rateLimitersMu.Unlock()
|
||||
for _, rl := range s.rateLimiters {
|
||||
rl.StopCleanup()
|
||||
}
|
||||
s.rateLimiters = s.rateLimiters[:0]
|
||||
}
|
||||
|
||||
func (s *Server) handleRegistrationError(source, path string, err error) error {
|
||||
var ce *matcher.ConflictError
|
||||
if errors.As(err, &ce) {
|
||||
@ -292,9 +271,6 @@ func (s *Server) Start() error {
|
||||
if s.config == nil {
|
||||
return fmt.Errorf("server config is nil")
|
||||
}
|
||||
if len(s.config.Servers) == 0 {
|
||||
return fmt.Errorf("no servers configured")
|
||||
}
|
||||
logging.Init(s.config.Logging.Error.Level, s.config.Logging.Format)
|
||||
|
||||
// 记录启动时间
|
||||
@ -493,10 +469,6 @@ func DupListener(ln net.Listener) (net.Listener, error) {
|
||||
// - 静态文件服务作为 fallback 处理非代理路径的请求
|
||||
// - 使用零拷贝传输优化大文件传输
|
||||
func (s *Server) startSingleMode() error {
|
||||
if len(s.config.Servers) == 0 {
|
||||
return fmt.Errorf("no servers configured")
|
||||
}
|
||||
|
||||
// 使用 Servers[0] 配置(迁移后 Server 字段为空)
|
||||
serverCfg := &s.config.Servers[0]
|
||||
|
||||
@ -770,10 +742,6 @@ func (s *Server) startMultiServerMode() error {
|
||||
return
|
||||
}
|
||||
fastSrv.TLSConfig = tlsManager.GetTLSConfig()
|
||||
|
||||
s.tlsManagersMu.Lock()
|
||||
s.tlsManagers = append(s.tlsManagers, tlsManager)
|
||||
s.tlsManagersMu.Unlock()
|
||||
}
|
||||
|
||||
s.fastServers[idx] = fastSrv
|
||||
|
||||
@ -656,26 +656,6 @@ func TestServer_TrackStats_EmptyBody(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_NilConfig 测试 nil 配置启动返回错误而非 panic。
|
||||
func TestStart_NilConfig(t *testing.T) {
|
||||
s := New(nil)
|
||||
|
||||
err := s.Start()
|
||||
if err == nil {
|
||||
t.Error("Start() with nil config should return error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_EmptyServers 测试空服务器列表启动返回错误而非 panic。
|
||||
func TestStart_EmptyServers(t *testing.T) {
|
||||
s := New(&config.Config{})
|
||||
|
||||
err := s.Start()
|
||||
if err == nil {
|
||||
t.Error("Start() with empty servers should return error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_Success 测试服务器配置初始化
|
||||
func TestStart_Success(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
|
||||
@ -166,7 +166,7 @@ func (h *StatusHandler) Path() string {
|
||||
// - ctx: FastHTTP 请求上下文
|
||||
func (h *StatusHandler) ServeHTTP(ctx *fasthttp.RequestCtx) {
|
||||
// 步骤1: 检查 IP 访问权限
|
||||
if !utils.CheckIPAccess(ctx, h.allowed, nil) {
|
||||
if !utils.CheckIPAccess(ctx, h.allowed) {
|
||||
utils.SendErrorWithDetail(ctx, utils.ErrForbidden, "Access denied")
|
||||
return
|
||||
}
|
||||
|
||||
@ -918,7 +918,7 @@ func TestStatusHandler_checkAccess_AllowList(t *testing.T) {
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.SetRemoteAddr(&net.TCPAddr{IP: tt.remoteIP, Port: 12345})
|
||||
|
||||
got := utils.CheckIPAccess(ctx, h.allowed, nil)
|
||||
got := utils.CheckIPAccess(ctx, h.allowed)
|
||||
if got != tt.wantAccess {
|
||||
t.Errorf("expected access %v, got %v", tt.wantAccess, got)
|
||||
}
|
||||
@ -1214,7 +1214,7 @@ func TestStatusHandler_ServeHTTP_AccessDenied_XForwardedFor(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusHandler_ServeHTTP_AccessAllowed_RemoteAddr(t *testing.T) {
|
||||
func TestStatusHandler_ServeHTTP_AccessAllowed_XForwardedFor(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := &config.StatusConfig{
|
||||
Path: "/_status",
|
||||
@ -1232,12 +1232,12 @@ func TestStatusHandler_ServeHTTP_AccessAllowed_RemoteAddr(t *testing.T) {
|
||||
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
ctx.Request.SetRequestURI("/_status")
|
||||
ctx.SetRemoteAddr(&net.TCPAddr{IP: net.ParseIP("10.5.5.5"), Port: 12345})
|
||||
ctx.Request.Header.Set("X-Forwarded-For", "10.5.5.5")
|
||||
|
||||
h.ServeHTTP(ctx)
|
||||
|
||||
if ctx.Response.StatusCode() != 200 {
|
||||
t.Errorf("expected status 200 for allowed access via RemoteAddr, got %d", ctx.Response.StatusCode())
|
||||
t.Errorf("expected status 200 for allowed access via X-Forwarded-For, got %d", ctx.Response.StatusCode())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -23,7 +23,6 @@ import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -550,31 +549,3 @@ func TestOCSPManager_StopTwice(t *testing.T) {
|
||||
t.Error("Expected manager to be stopped")
|
||||
}
|
||||
}
|
||||
|
||||
// TestOCSPGetOCSPResponse_Concurrent 验证并发读取 OCSP 响应不会产生竞态。
|
||||
func TestOCSPGetOCSPResponse_Concurrent(t *testing.T) {
|
||||
mgr := NewOCSPManager(nil)
|
||||
|
||||
serial := "12345"
|
||||
testResp := []byte("concurrent-test-response")
|
||||
|
||||
mgr.mu.Lock()
|
||||
mgr.responses[serial] = &ocspResponse{
|
||||
response: testResp,
|
||||
thisUpdate: time.Now(),
|
||||
nextUpdate: time.Now().Add(1 * time.Hour),
|
||||
status: statusValid,
|
||||
fetchedAt: time.Now(),
|
||||
}
|
||||
mgr.mu.Unlock()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for range 100 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = mgr.GetOCSPResponse(serial)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
@ -28,7 +28,7 @@ func TestListenTCP_Success(t *testing.T) {
|
||||
|
||||
s := NewServer()
|
||||
|
||||
err := s.ListenTCP("127.0.0.1:0", "")
|
||||
err := s.ListenTCP("127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
|
||||
s.mu.RLock()
|
||||
@ -47,7 +47,7 @@ func TestListenTCP_InvalidAddress(t *testing.T) {
|
||||
|
||||
s := NewServer()
|
||||
|
||||
err := s.ListenTCP("256.256.256.256:99999", "")
|
||||
err := s.ListenTCP("256.256.256.256:99999")
|
||||
assert.Error(t, err)
|
||||
|
||||
s.mu.RLock()
|
||||
@ -75,7 +75,7 @@ func TestStart_WithTCPListeners(t *testing.T) {
|
||||
}
|
||||
_ = s.AddUpstream("test", targets, "round_robin", HealthCheckSpec{})
|
||||
|
||||
err := s.ListenTCP("127.0.0.1:0", "")
|
||||
err := s.ListenTCP("127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
|
||||
err = s.Start()
|
||||
@ -763,142 +763,3 @@ func TestStart_DoubleStart(t *testing.T) {
|
||||
assert.Error(t, err)
|
||||
s.Stop()
|
||||
}
|
||||
|
||||
// TestListenTCP_AssociatesUpstream 验证 ListenTCP 会建立 listener->upstream 映射。
|
||||
func TestListenTCP_AssociatesUpstream(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
s := NewServer()
|
||||
|
||||
_ = s.AddUpstream("test", []TargetSpec{{Addr: "127.0.0.1:8001", Weight: 1}}, "round_robin", HealthCheckSpec{})
|
||||
|
||||
err := s.ListenTCP("127.0.0.1:0", "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
require.Len(t, s.listeners, 1)
|
||||
require.Len(t, s.listenerUpstreams, 1)
|
||||
for addr, up := range s.listenerUpstreams {
|
||||
assert.NotNil(t, up, "addr=%s", addr)
|
||||
assert.Equal(t, "test", up.name)
|
||||
}
|
||||
}
|
||||
|
||||
// TestListenTCP_UnknownUpstream 验证为不存在的上游建立映射时不会 panic。
|
||||
func TestListenTCP_UnknownUpstream(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
s := NewServer()
|
||||
|
||||
err := s.ListenTCP("127.0.0.1:0", "missing")
|
||||
require.NoError(t, err)
|
||||
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
require.Len(t, s.listenerUpstreams, 1)
|
||||
for _, up := range s.listenerUpstreams {
|
||||
assert.Nil(t, up)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStart_TCPForwardsViaListenerUpstreamMapping 验证 TCP stream proxy 通过 listener->upstream 映射转发到后端。
|
||||
func TestStart_TCPForwardsViaListenerUpstreamMapping(t *testing.T) {
|
||||
s := NewServer()
|
||||
|
||||
backendLn, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
defer backendLn.Close()
|
||||
|
||||
go func() {
|
||||
for {
|
||||
conn, acceptErr := backendLn.Accept()
|
||||
if acceptErr != nil {
|
||||
return
|
||||
}
|
||||
go func(c net.Conn) {
|
||||
defer c.Close()
|
||||
buf := make([]byte, 1024)
|
||||
n, _ := c.Read(buf)
|
||||
_, _ = c.Write(append([]byte("backend:"), buf[:n]...))
|
||||
}(conn)
|
||||
}
|
||||
}()
|
||||
|
||||
_ = s.AddUpstream("test", []TargetSpec{{Addr: backendLn.Addr().String(), Weight: 1}}, "round_robin", HealthCheckSpec{})
|
||||
s.upstreams["test"].targets[0].healthy.Store(true)
|
||||
|
||||
err = s.ListenTCP("127.0.0.1:0", "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
s.mu.RLock()
|
||||
var proxyAddr string
|
||||
for _, ln := range s.listeners {
|
||||
proxyAddr = ln.Addr().String()
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
|
||||
err = s.Start()
|
||||
require.NoError(t, err)
|
||||
defer s.Stop()
|
||||
|
||||
clientConn, err := net.DialTimeout("tcp", proxyAddr, 2*time.Second)
|
||||
require.NoError(t, err)
|
||||
defer clientConn.Close()
|
||||
|
||||
testData := []byte("hello mapping")
|
||||
_, err = clientConn.Write(testData)
|
||||
require.NoError(t, err)
|
||||
|
||||
_ = clientConn.SetReadDeadline(time.Now().Add(3 * time.Second))
|
||||
buf := make([]byte, 1024)
|
||||
n, err := clientConn.Read(buf)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, append([]byte("backend:"), testData...), buf[:n])
|
||||
}
|
||||
|
||||
// TestStop_UDP 验证 UDP stream server 的 Stop 能正常完成,不会死锁。
|
||||
func TestStop_UDP(t *testing.T) {
|
||||
s := NewServer()
|
||||
|
||||
backendAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
backendConn, err := net.ListenUDP("udp", backendAddr)
|
||||
require.NoError(t, err)
|
||||
defer backendConn.Close()
|
||||
|
||||
go func() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, addr, readErr := backendConn.ReadFromUDP(buf)
|
||||
if readErr != nil {
|
||||
return
|
||||
}
|
||||
_, _ = backendConn.WriteToUDP(buf[:n], addr)
|
||||
}
|
||||
}()
|
||||
|
||||
_ = s.AddUpstream("udp-test", []TargetSpec{{Addr: backendConn.LocalAddr().String(), Weight: 1}}, "round_robin", HealthCheckSpec{})
|
||||
|
||||
err = s.ListenUDP("127.0.0.1:0", "udp-test", 100*time.Millisecond)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = s.Start()
|
||||
require.NoError(t, err)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.Stop()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("UDP server Stop did not complete in time")
|
||||
}
|
||||
|
||||
assert.False(t, s.running.Load())
|
||||
}
|
||||
|
||||
@ -21,7 +21,7 @@
|
||||
// log.Fatal(err)
|
||||
// }
|
||||
//
|
||||
// server.ListenTCP(":3306", "mysql")
|
||||
// server.ListenTCP(":3306")
|
||||
// server.Start()
|
||||
// defer server.Stop()
|
||||
//
|
||||
@ -287,15 +287,14 @@ func (i *ipHash) SelectByIP(targets []*Target, clientIP string) *Target {
|
||||
|
||||
// Server TCP/UDP Stream 代理服务器。
|
||||
type Server struct {
|
||||
listeners map[string]net.Listener
|
||||
udpServers map[string]*udpServer
|
||||
upstreams map[string]*Upstream
|
||||
listenerUpstreams map[string]*Upstream // 监听地址到上游的映射
|
||||
connCount atomic.Int64
|
||||
mu sync.RWMutex
|
||||
running atomic.Bool
|
||||
wg sync.WaitGroup
|
||||
stopCh chan struct{}
|
||||
listeners map[string]net.Listener
|
||||
udpServers map[string]*udpServer
|
||||
upstreams map[string]*Upstream
|
||||
connCount atomic.Int64
|
||||
mu sync.RWMutex
|
||||
running atomic.Bool
|
||||
wg sync.WaitGroup
|
||||
stopCh chan struct{}
|
||||
}
|
||||
|
||||
// Upstream Stream 上游配置。
|
||||
@ -388,11 +387,10 @@ type HealthCheckSpec struct {
|
||||
// - *Server: 初始化的 Stream 代理服务器实例
|
||||
func NewServer() *Server {
|
||||
return &Server{
|
||||
listeners: make(map[string]net.Listener),
|
||||
udpServers: make(map[string]*udpServer),
|
||||
upstreams: make(map[string]*Upstream),
|
||||
listenerUpstreams: make(map[string]*Upstream),
|
||||
stopCh: make(chan struct{}),
|
||||
listeners: make(map[string]net.Listener),
|
||||
udpServers: make(map[string]*udpServer),
|
||||
upstreams: make(map[string]*Upstream),
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
@ -448,10 +446,7 @@ func (s *Server) AddUpstream(name string, targets []TargetSpec, lbType string, h
|
||||
}
|
||||
|
||||
// ListenTCP 开始监听 TCP 端口。
|
||||
//
|
||||
// upstreamName 指定该监听地址对应的上游配置名称,建立 listener->upstream 映射,
|
||||
// 使 handleConnection 可以根据监听地址找到正确的上游。
|
||||
func (s *Server) ListenTCP(addr string, upstreamName string) error {
|
||||
func (s *Server) ListenTCP(addr string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
@ -461,7 +456,6 @@ func (s *Server) ListenTCP(addr string, upstreamName string) error {
|
||||
}
|
||||
|
||||
s.listeners[addr] = listener
|
||||
s.listenerUpstreams[addr] = s.upstreams[upstreamName]
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -539,12 +533,10 @@ func (s *Server) Stop() {
|
||||
_ = ln.Close()
|
||||
}
|
||||
for _, udpSrv := range s.udpServers {
|
||||
udpSrv.running.Store(false)
|
||||
close(udpSrv.stopCh)
|
||||
if udpSrv.conn != nil {
|
||||
_ = udpSrv.conn.Close()
|
||||
}
|
||||
udpSrv.closeSessions()
|
||||
}
|
||||
for _, upstream := range s.upstreams {
|
||||
if upstream.healthChk != nil && upstream.healthChk.stopCh != nil {
|
||||
@ -608,11 +600,7 @@ func (s *Server) handleConnection(clientConn net.Conn, addr string) {
|
||||
}()
|
||||
|
||||
s.mu.RLock()
|
||||
upstream := s.listenerUpstreams[addr]
|
||||
if upstream == nil {
|
||||
// 兼容未建立映射的历史用法:回退到按监听地址查找上游
|
||||
upstream = s.upstreams[addr]
|
||||
}
|
||||
upstream := s.upstreams[addr]
|
||||
s.mu.RUnlock()
|
||||
|
||||
if upstream == nil {
|
||||
@ -901,19 +889,6 @@ func (s *udpServer) removeSession(clientAddr *net.UDPAddr) {
|
||||
}
|
||||
}
|
||||
|
||||
// closeSessions 关闭所有 UDP 会话。
|
||||
//
|
||||
// 在服务器停止时调用,释放后端连接并退出响应监听 goroutine。
|
||||
func (s *udpServer) closeSessions() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
for _, session := range s.sessions {
|
||||
session.close()
|
||||
}
|
||||
clear(s.sessions)
|
||||
}
|
||||
|
||||
// close 关闭会话。
|
||||
//
|
||||
// 使用 sync.Once 确保会话只关闭一次。
|
||||
|
||||
@ -257,7 +257,7 @@ func TestConcurrentConnections(t *testing.T) {
|
||||
wg.Wait()
|
||||
|
||||
if s.connCount.Load() != 100 {
|
||||
t.Errorf("Expected 100 connections, got %d", s.connCount.Load())
|
||||
t.Errorf("Expected 100 connections, got %d", s.connCount)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -66,13 +66,12 @@ func SendJSONError(ctx *fasthttp.RequestCtx, status int, errMsg string) {
|
||||
|
||||
// CheckIPAccess checks whether the client IP is in the allowed list.
|
||||
// If allowed is empty, all access is permitted.
|
||||
// trustedProxies is used to safely extract the client IP from X-Forwarded-For.
|
||||
func CheckIPAccess(ctx *fasthttp.RequestCtx, allowed []net.IPNet, trustedProxies []net.IPNet) bool {
|
||||
func CheckIPAccess(ctx *fasthttp.RequestCtx, allowed []net.IPNet) bool {
|
||||
if len(allowed) == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
clientIP := netutil.ExtractClientIPWithTrustedProxies(ctx, trustedProxies)
|
||||
clientIP := netutil.ExtractClientIPNet(ctx)
|
||||
if clientIP == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
@ -440,42 +440,6 @@ func TestInit_ArgsGetterBytes(t *testing.T) {
|
||||
assert.Equal(t, []byte("key=val"), result)
|
||||
}
|
||||
|
||||
// TestRequestID_GeneratedOnce 验证 request_id 在一次请求中只生成一次并存储在 UserValue。
|
||||
func TestRequestID_GeneratedOnce(t *testing.T) {
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
vc := NewContext(ctx)
|
||||
defer ReleaseContext(vc)
|
||||
|
||||
id1, ok1 := vc.Get(VarRequestID)
|
||||
id2, ok2 := vc.Get(VarRequestID)
|
||||
|
||||
require.True(t, ok1, "第一次获取 request_id 应该成功")
|
||||
require.True(t, ok2, "第二次获取 request_id 应该成功")
|
||||
assert.Equal(t, id1, id2, "同一请求的 request_id 应该相同")
|
||||
assert.NotEmpty(t, id1, "request_id 不应为空")
|
||||
|
||||
// 验证已存储在 UserValue 中
|
||||
uv := ctx.UserValue(VarRequestID)
|
||||
require.NotNil(t, uv)
|
||||
assert.Equal(t, id1, uv)
|
||||
}
|
||||
|
||||
// TestTimeLocal_Format 验证 time_local 使用 -0700 时区格式。
|
||||
func TestTimeLocal_Format(t *testing.T) {
|
||||
ctx := &fasthttp.RequestCtx{}
|
||||
vc := NewContext(ctx)
|
||||
defer ReleaseContext(vc)
|
||||
|
||||
value, ok := vc.Get(VarTimeLocal)
|
||||
require.True(t, ok, "time_local 应该存在")
|
||||
assert.NotEmpty(t, value, "time_local 不应为空")
|
||||
|
||||
// 格式示例:16/Jun/2026:10:23:42 +0000
|
||||
// 必须包含时区偏移(+/-HHMM)
|
||||
assert.Regexp(t, `\d{2}/[A-Za-z]{3}/\d{4}:\d{2}:\d{2}:\d{2} [+-]\d{4}`, value,
|
||||
"time_local 格式应为 02/Jan/2006:15:04:05 -0700")
|
||||
}
|
||||
|
||||
// TestInit_MethodGetterBytes 测试 request_method 变量有 GetterBytes
|
||||
func TestInit_MethodGetterBytes(t *testing.T) {
|
||||
builtin := GetBuiltin(VarRequestMethod)
|
||||
|
||||
@ -134,16 +134,6 @@ func NewContext(ctx *fasthttp.RequestCtx) *Context {
|
||||
vc.upstreamResponseTime = 0
|
||||
vc.upstreamConnectTime = 0
|
||||
vc.upstreamHeaderTime = 0
|
||||
// 防御性初始化:防止从池中取出的 Context 因 map 为 nil 而 panic
|
||||
if vc.store == nil {
|
||||
vc.store = make(map[string]string)
|
||||
}
|
||||
if vc.cache == nil {
|
||||
vc.cache = make(map[string]string)
|
||||
}
|
||||
if vc.bytesCache == nil {
|
||||
vc.bytesCache = make(map[string][]byte)
|
||||
}
|
||||
// 清空内置变量缓存
|
||||
for k := range vc.cache {
|
||||
delete(vc.cache, k)
|
||||
|
||||
@ -457,33 +457,6 @@ func TestExpandOnlyDollar(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewContext_NilMaps 验证 NewContext 能防御性地初始化 nil map。
|
||||
func TestNewContext_NilMaps(t *testing.T) {
|
||||
ctx := mockRequestCtx(t)
|
||||
|
||||
// 向池中放入多个 map 为 nil 的 Context,模拟异常状态
|
||||
for range 10 {
|
||||
pool.Put(&Context{})
|
||||
}
|
||||
|
||||
// 连续获取,确保所有返回的 Context 都能安全使用,不会 panic
|
||||
for range 10 {
|
||||
vc := NewContext(ctx)
|
||||
|
||||
vc.Set("key", "value")
|
||||
if v, ok := vc.Get("key"); !ok || v != "value" {
|
||||
t.Errorf("expected key=value, got %q", v)
|
||||
}
|
||||
|
||||
// 验证缓存也能正常工作
|
||||
if v, ok := vc.Get("host"); !ok || v != "example.com" {
|
||||
t.Errorf("expected host=example.com, got %q", v)
|
||||
}
|
||||
|
||||
ReleaseContext(vc)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPoolFunctions 测试 Pool 相关函数
|
||||
func TestPoolFunctions(t *testing.T) {
|
||||
ctx := mockRequestCtx(t)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user