go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

docs(limiter/traffic): add doc comments, unit tests

Browse Source
This commit is contained in:
ginuerzh
2026-05-24 20:54:34 +08:00
parent e9372ea44a
commit dc99f4731f
12 changed files with 1511 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
limiter/traffic/traffic_test.go
+578
View File
@@ -0,0 +1,578 @@
package traffic
import (
"context"
"io"
"strings"
"sync"
"testing"
"time"
"github.com/go-gost/core/limiter"
traffic "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/x/internal/loader"
xlogger "github.com/go-gost/x/logger"
"github.com/patrickmn/go-cache"
"github.com/yl2chen/cidranger"
)
// --- test helpers ---
func newTestTrafficLimiter(limits ...string) *trafficLimiter {
ctx, cancel := context.WithCancel(context.Background())
lim := &trafficLimiter{
cidrGenerators: cidranger.NewPCTrieRanger(),
connInLimits: cache.New(defaultExpiration, cleanupInterval),
connOutLimits: cache.New(defaultExpiration, cleanupInterval),
inLimits: cache.New(defaultExpiration, cleanupInterval),
outLimits: cache.New(defaultExpiration, cleanupInterval),
options: options{limits: limits},
cancelFunc: cancel,
logger: xlogger.Nop(),
}
_ = lim.reload(ctx)
return lim
}
type fakeLoader struct {
data string
listErr error
loadErr error
closed bool
}
func (l *fakeLoader) List(ctx context.Context) ([]string, error) {
if l.listErr != nil {
return nil, l.listErr
}
if l.data == "" {
return nil, nil
}
return strings.Split(l.data, "\n"), nil
}
func (l *fakeLoader) Load(ctx context.Context) (io.Reader, error) {
if l.loadErr != nil {
return nil, l.loadErr
}
return strings.NewReader(l.data), nil
}
func (l *fakeLoader) Close() error {
l.closed = true
return nil
}
type errorLoader struct{}
func (e *errorLoader) Load(ctx context.Context) (io.Reader, error) { return nil, io.ErrUnexpectedEOF }
func (e *errorLoader) Close() error { return nil }
// --- parsing tests ---
func TestParseLine(t *testing.T) {
l := newTestTrafficLimiter()
tests := []struct {
input, expected string
}{
{"192.168.1.1 100B 200B", "192.168.1.1 100B 200B"},
{" 192.168.1.1 100B 200B ", "192.168.1.1 100B 200B"},
{"192.168.1.1 100B 200B # comment", "192.168.1.1 100B 200B"},
{"# only comment", ""},
{" # indented comment", ""},
{"", ""},
}
for _, tt := range tests {
got := l.parseLine(tt.input)
if got != tt.expected {
t.Errorf("parseLine(%q) = %q, want %q", tt.input, got, tt.expected)
}
}
}
func TestParseLimit(t *testing.T) {
l := newTestTrafficLimiter()
tests := []struct {
input string
key string
in, out int
}{
{"192.168.1.1 100B 200B", "192.168.1.1", 100, 200},
{"192.168.1.1\t100B\t200B", "192.168.1.1", 100, 200},
{"$ 1KB", "$", 1024, 0},
{"$$ 500B", "$$", 500, 0},
{"10.0.0.0/8 1KB 2KB", "10.0.0.0/8", 1024, 2048},
{"key 100B 200B", "key", 100, 200},
{"key 1MB", "key", 1048576, 0},
{"key 512KB 1MB", "key", 524288, 1048576},
{"key_only", "", 0, 0},
{"", "", 0, 0},
{"key invalid_bytes", "key", 0, 0},
}
for _, tt := range tests {
key, in, out := l.parseLimit(tt.input)
if key != tt.key || in != tt.in || out != tt.out {
t.Errorf("parseLimit(%q) = (%q, %d, %d), want (%q, %d, %d)",
tt.input, key, in, out, tt.key, tt.in, tt.out)
}
}
}
func TestParsePatterns(t *testing.T) {
l := newTestTrafficLimiter()
r := strings.NewReader("192.168.1.1 100B\n# comment\n\n10.0.0.1 200B\n")
patterns, err := l.parsePatterns(r)
if err != nil {
t.Fatal(err)
}
if len(patterns) != 2 {
t.Fatalf("expected 2 patterns, got %d: %v", len(patterns), patterns)
}
if patterns[0] != "192.168.1.1 100B" {
t.Errorf("patterns[0] = %q", patterns[0])
}
if patterns[1] != "10.0.0.1 200B" {
t.Errorf("patterns[1] = %q", patterns[1])
}
}
func TestParsePatterns_NilReader(t *testing.T) {
l := newTestTrafficLimiter()
patterns, err := l.parsePatterns(nil)
if err != nil {
t.Fatal(err)
}
if len(patterns) != 0 {
t.Fatalf("expected 0 patterns, got %d", len(patterns))
}
}
// --- load tests ---
func TestLoad_NoLoaders(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
v, ok := values["$"]
if !ok {
t.Fatal("expected $ key in loaded values")
}
if v.in != 100 || v.out != 0 {
t.Fatalf("expected in=100 out=0, got in=%d out=%d", v.in, v.out)
}
}
func TestLoad_FileLister(t *testing.T) {
l := newTestTrafficLimiter()
l.options.fileLoader = &fakeLoader{data: "$ 50B\n192.168.1.1 200B 300B"}
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
if v := values["$"]; v.in != 50 {
t.Fatalf("expected $ in=50, got %d", v.in)
}
v := values["192.168.1.1"]
if v.in != 200 || v.out != 300 {
t.Fatalf("expected in=200 out=300, got in=%d out=%d", v.in, v.out)
}
}
func TestLoad_FileLoader(t *testing.T) {
l := newTestTrafficLimiter()
l.options.fileLoader = &fakeLoader{data: "$ 30B"}
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
if v := values["$"]; v.in != 30 {
t.Fatalf("expected $ in=30, got %d", v.in)
}
}
func TestLoad_RedisLoader(t *testing.T) {
l := newTestTrafficLimiter()
l.options.redisLoader = &fakeLoader{data: "$ 40B"}
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
if v := values["$"]; v.in != 40 {
t.Fatalf("expected $ in=40, got %d", v.in)
}
}
func TestLoad_HTTPLoader(t *testing.T) {
l := newTestTrafficLimiter()
l.options.httpLoader = &fakeLoader{data: "$ 35B"}
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
if v := values["$"]; v.in != 35 {
t.Fatalf("expected $ in=35, got %d", v.in)
}
}
func TestLoad_ErrorLoader(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
l.options.fileLoader = &errorLoader{}
values, err := l.load(context.Background())
if err != nil {
t.Fatal("error loader should not propagate error")
}
if v := values["$"]; v.in != 100 {
t.Fatalf("expected $ in=100 from static limits, got %d", v.in)
}
}
func TestLoad_MergesLoaders(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
l.options.fileLoader = &fakeLoader{data: "192.168.1.1 50B"}
l.options.httpLoader = &fakeLoader{data: "$ 200B"}
values, err := l.load(context.Background())
if err != nil {
t.Fatal(err)
}
if v := values["$"]; v.in != 200 {
t.Fatalf("expected $ in=200 (HTTP overrides static), got %d", v.in)
}
if v := values["192.168.1.1"]; v.in != 50 {
t.Fatalf("expected 192.168.1.1 in=50, got %d", v.in)
}
}
// --- In/Out tests ---
func TestIn_ServiceScope(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
v := l.In(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeService))
if v == nil {
t.Fatal("expected service-level In limiter")
}
if v.Limit() != 100 {
t.Fatalf("expected limit 100, got %d", v.Limit())
}
}
func TestIn_ClientScope(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
v := l.In(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeClient))
if v != nil {
t.Fatal("ScopeClient should return nil")
}
}
func TestIn_ConnScope_GeneratesLimiter(t *testing.T) {
l := newTestTrafficLimiter("$$ 100B")
v := l.In(context.Background(), "10.0.0.1:12345")
if v == nil {
t.Fatal("expected conn-level In limiter")
}
if v.Limit() != 100 {
t.Fatalf("expected limit 100, got %d", v.Limit())
}
}
func TestIn_ConnScope_CachedResult(t *testing.T) {
l := newTestTrafficLimiter("$$ 100B")
v1 := l.In(context.Background(), "10.0.0.1:12345")
v2 := l.In(context.Background(), "10.0.0.1:12345")
// Both should have the correct limit (cached limiter reused in group).
if v1 == nil || v2 == nil {
t.Fatal("limiters should not be nil")
}
if v1.Limit() != 100 || v2.Limit() != 100 {
t.Fatalf("expected limit 100, got v1=%d v2=%d", v1.Limit(), v2.Limit())
}
// The cache entry should exist for the key.
if _, ok := l.connInLimits.Get("10.0.0.1:12345"); !ok {
t.Fatal("expected cached conn-level In limiter")
}
}
func TestIn_IPMatch(t *testing.T) {
l := newTestTrafficLimiter("192.168.1.1 100B")
v := l.In(context.Background(), "192.168.1.1:12345")
if v == nil {
t.Fatal("expected IP-level In limiter")
}
if v.Limit() != 100 {
t.Fatalf("expected limit 100, got %d", v.Limit())
}
}
func TestIn_CIDRMatch(t *testing.T) {
l := newTestTrafficLimiter("10.0.0.0/8 100B")
v := l.In(context.Background(), "10.0.0.1:12345")
if v == nil {
t.Fatal("expected CIDR-level In limiter")
}
if v.Limit() != 100 {
t.Fatalf("expected limit 100, got %d", v.Limit())
}
}
func TestIn_CIDRNoMatch(t *testing.T) {
l := newTestTrafficLimiter("10.0.0.0/8 100B")
v := l.In(context.Background(), "192.168.1.1:12345")
if v != nil {
t.Fatalf("CIDR should not match different IP range, got limit %d", v.Limit())
}
}
func TestIn_IPTakesPrecedence(t *testing.T) {
l := newTestTrafficLimiter("10.0.0.0/8 5B", "10.0.0.1 2B")
v := l.In(context.Background(), "10.0.0.1:12345")
if v == nil {
t.Fatal("expected limiter for exact IP + CIDR")
}
if v.Limit() != 2 {
t.Fatalf("exact IP should take precedence, expected limit 2, got %d", v.Limit())
}
}
func TestOut_ServiceScope(t *testing.T) {
l := newTestTrafficLimiter("$ 100B 100B")
v := l.Out(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeService))
if v == nil {
t.Fatal("expected service-level Out limiter")
}
if v.Limit() != 100 {
t.Fatalf("expected limit 100, got %d", v.Limit())
}
}
func TestOut_ClientScope(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
v := l.Out(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeClient))
if v != nil {
t.Fatal("ScopeClient Out should return nil")
}
}
func TestOut_ConnScope(t *testing.T) {
l := newTestTrafficLimiter("$$ 200B 200B")
v := l.Out(context.Background(), "10.0.0.1:12345")
if v == nil {
t.Fatal("expected conn-level Out limiter")
}
if v.Limit() != 200 {
t.Fatalf("expected limit 200, got %d", v.Limit())
}
}
// --- reload tests ---
func TestReload_ServiceLimit(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("$ 100B")
v1 := l.In(ctx, "x", limiter.ScopeOption(limiter.ScopeService))
if v1.Limit() != 100 {
t.Fatalf("expected 100, got %d", v1.Limit())
}
l.options.limits = []string{"$ 200B"}
_ = l.reload(ctx)
v2 := l.In(ctx, "x", limiter.ScopeOption(limiter.ScopeService))
if v2.Limit() != 200 {
t.Fatalf("expected 200 after reload, got %d", v2.Limit())
}
}
func TestReload_ConnLimit(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("$$ 100B")
v1 := l.In(ctx, "10.0.0.1:12345")
if v1.Limit() != 100 {
t.Fatalf("expected 100, got %d", v1.Limit())
}
l.options.limits = []string{"$$ 200B"}
_ = l.reload(ctx)
v2 := l.In(ctx, "10.0.0.2:12345")
if v2.Limit() != 200 {
t.Fatalf("expected 200 after reload, got %d", v2.Limit())
}
}
func TestReload_ConnLimit_UpdatesExisting(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("$$ 100B")
v1 := l.In(ctx, "10.0.0.1:12345")
if v1.Limit() != 100 {
t.Fatalf("expected 100, got %d", v1.Limit())
}
l.options.limits = []string{"$$ 200B"}
_ = l.reload(ctx)
// Same cached limiter should have been updated in-place.
if v1.Limit() != 200 {
t.Fatalf("existing cached limiter should be updated to 200, got %d", v1.Limit())
}
}
func TestReload_IPLimit(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("192.168.1.1 100B")
v := l.In(ctx, "192.168.1.1:12345")
if v.Limit() != 100 {
t.Fatalf("expected 100, got %d", v.Limit())
}
l.options.limits = []string{"192.168.1.1 200B"}
_ = l.reload(ctx)
if v.Limit() != 200 {
t.Fatalf("expected 200 after reload, got %d", v.Limit())
}
}
func TestReload_CIDRLimit(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("10.0.0.0/8 100B")
v := l.In(ctx, "10.0.0.1:12345")
if v.Limit() != 100 {
t.Fatalf("expected 100, got %d", v.Limit())
}
l.options.limits = []string{"10.0.0.0/8 200B"}
_ = l.reload(ctx)
if v.Limit() != 200 {
t.Fatalf("expected 200 after CIDR reload, got %d", v.Limit())
}
}
func TestReload_RemovesObsoleteLimit(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("192.168.1.1 100B")
_ = l.In(ctx, "192.168.1.1:12345")
l.options.limits = nil
_ = l.reload(ctx)
v := l.In(ctx, "192.168.1.1:12345")
if v != nil {
t.Fatalf("obsolete IP limit should be removed, got limit %d", v.Limit())
}
}
func TestReload_ClearsConnCacheWhenRemoved(t *testing.T) {
ctx := context.Background()
l := newTestTrafficLimiter("$$ 100B")
_ = l.In(ctx, "10.0.0.1:12345")
l.options.limits = []string{"$$ 0B"}
_ = l.reload(ctx)
v := l.In(ctx, "10.0.0.2:12345")
if v != nil {
t.Fatalf("conn cache should be flushed, got limiter with limit %d", v.Limit())
}
}
// --- close tests ---
func TestClose(t *testing.T) {
fl := &fakeLoader{}
rl := &fakeLoader{}
hl := &fakeLoader{}
l := newTestTrafficLimiter()
l.options.fileLoader = fl
l.options.redisLoader = rl
l.options.httpLoader = hl
l.Close()
if !fl.closed {
t.Error("fileLoader should be closed")
}
if !rl.closed {
t.Error("redisLoader should be closed")
}
if !hl.closed {
t.Error("httpLoader should be closed")
}
}
func TestClose_Idempotent(t *testing.T) {
l := newTestTrafficLimiter()
l.Close()
l.Close()
}
func TestClose_NilLoaders(t *testing.T) {
l := newTestTrafficLimiter()
if err := l.Close(); err != nil {
t.Fatal("Close with nil loaders should return nil error")
}
}
// --- lifecycle tests ---
func TestPeriodReload_ZeroPeriod(t *testing.T) {
l := newTestTrafficLimiter()
l.options.period = 0
ctx := context.Background()
err := l.periodReload(ctx)
if err != nil {
t.Fatalf("zero period should return nil, got %v", err)
}
}
func TestPeriodReload_ContextCancel(t *testing.T) {
l := newTestTrafficLimiter("$ 100B")
l.options.period = time.Second
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := l.periodReload(ctx)
if err != context.Canceled {
t.Fatalf("expected context.Canceled, got %v", err)
}
}
func TestNewTrafficLimiter_Defaults(t *testing.T) {
lim := NewTrafficLimiter()
if lim == nil {
t.Fatal("NewTrafficLimiter should not return nil")
}
if closer, ok := lim.(io.Closer); ok {
closer.Close()
}
}
// --- concurrency ---
func TestIn_ConcurrentAccess(t *testing.T) {
l := newTestTrafficLimiter("$$ 1KB", "192.168.1.1 500B", "10.0.0.0/8 200B", "$ 100B")
ctx := context.Background()
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
_ = l.In(ctx, "10.0.0.1:12345")
_ = l.In(ctx, "192.168.1.1:8080")
_ = l.Out(ctx, "10.0.0.2:9090")
}(i)
}
wg.Wait()
}
// --- interface compliance ---
var (
_ traffic.TrafficLimiter = (*trafficLimiter)(nil)
_ io.Closer = (*trafficLimiter)(nil)
_ loader.Loader = (*fakeLoader)(nil)
_ loader.Lister = (*fakeLoader)(nil)
)