From dc99f4731f9e1d4ff1f817ec4bb344cd9a780699 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sun, 24 May 2026 20:54:34 +0800 Subject: [PATCH] docs(limiter/traffic): add doc comments, unit tests --- limiter/traffic/cache/limiter.go | 13 + limiter/traffic/cache/limiter_test.go | 175 +++++++ limiter/traffic/generator_test.go | 63 +++ limiter/traffic/limiter.go | 2 + limiter/traffic/limiter_test.go | 135 ++++++ limiter/traffic/traffic.go | 20 +- limiter/traffic/traffic_test.go | 578 +++++++++++++++++++++++ limiter/traffic/wrapper/conn.go | 9 +- limiter/traffic/wrapper/conn_test.go | 448 ++++++++++++++++++ limiter/traffic/wrapper/io.go | 4 +- limiter/traffic/wrapper/listener.go | 2 + limiter/traffic/wrapper/listener_test.go | 65 +++ 12 files changed, 1511 insertions(+), 3 deletions(-) create mode 100644 limiter/traffic/cache/limiter_test.go create mode 100644 limiter/traffic/generator_test.go create mode 100644 limiter/traffic/limiter_test.go create mode 100644 limiter/traffic/traffic_test.go create mode 100644 limiter/traffic/wrapper/conn_test.go create mode 100644 limiter/traffic/wrapper/listener_test.go diff --git a/limiter/traffic/cache/limiter.go b/limiter/traffic/cache/limiter.go index c52d685a..70c497de 100644 --- a/limiter/traffic/cache/limiter.go +++ b/limiter/traffic/cache/limiter.go @@ -1,3 +1,6 @@ +// Package limiter provides a cached traffic limiter that wraps a +// TrafficLimiter with in-memory caching to reduce calls to the underlying +// limiter (e.g., plugin-based limiters). package limiter import ( @@ -20,20 +23,27 @@ type options struct { scope string } +// Option configures a cachedTrafficLimiter. type Option func(*options) +// RefreshIntervalOption sets the cache refresh interval. Cached limiters +// are refreshed after this duration. The minimum value is 1 second. func RefreshIntervalOption(interval time.Duration) Option { return func(o *options) { o.refreshInterval = interval } } +// CleanupIntervalOption sets the interval for cleaning up expired cache +// entries. The minimum value is 1 second. func CleanupIntervalOption(interval time.Duration) Option { return func(o *options) { o.cleanupInterval = interval } } +// ScopeOption filters cached limiters by scope. When set, only requests +// with a matching scope return a limiter; non-matching scopes return nil. func ScopeOption(scope string) Option { return func(o *options) { o.scope = scope @@ -47,6 +57,9 @@ type cachedTrafficLimiter struct { options options } +// NewCachedTrafficLimiter wraps a TrafficLimiter with an in-memory cache. +// Cached limiters are refreshed after the configured refresh interval. +// Returns nil if the inner limiter is nil. func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, opts ...Option) traffic.TrafficLimiter { if limiter == nil { return nil diff --git a/limiter/traffic/cache/limiter_test.go b/limiter/traffic/cache/limiter_test.go new file mode 100644 index 00000000..0dcbe3e5 --- /dev/null +++ b/limiter/traffic/cache/limiter_test.go @@ -0,0 +1,175 @@ +package limiter + +import ( + "context" + "testing" + "time" + + "github.com/go-gost/core/limiter" + "github.com/go-gost/core/limiter/traffic" +) + +// --- mock types --- + +type mockLimiter struct { + limit int +} + +func (m *mockLimiter) Wait(ctx context.Context, n int) int { return n } +func (m *mockLimiter) Limit() int { return m.limit } +func (m *mockLimiter) Set(n int) { m.limit = n } + +type mockTrafficLimiter struct { + inCalls int + outCalls int + inLim traffic.Limiter + outLim traffic.Limiter +} + +func (m *mockTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + m.inCalls++ + return m.inLim +} + +func (m *mockTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + m.outCalls++ + return m.outLim +} + +// --- tests --- + +func TestCachedTrafficLimiter_In(t *testing.T) { + inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(50*time.Millisecond)) + if cached == nil { + t.Fatal("NewCachedTrafficLimiter should not return nil") + } + + ctx := context.Background() + v1 := cached.In(ctx, "key1") + if v1 == nil { + t.Fatal("first In() should return a limiter") + } + if v1.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", v1.Limit()) + } + if inner.inCalls != 1 { + t.Fatalf("expected 1 In call, got %d", inner.inCalls) + } + + // Second call: should return cached result without calling inner again. + v2 := cached.In(ctx, "key1") + if v2 == nil { + t.Fatal("cached In() should return a limiter") + } + if inner.inCalls != 1 { + t.Fatalf("cached call should not invoke inner, got %d calls", inner.inCalls) + } +} + +func TestCachedTrafficLimiter_Out(t *testing.T) { + inner := &mockTrafficLimiter{outLim: &mockLimiter{limit: 200}} + cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(50*time.Millisecond)) + + ctx := context.Background() + v1 := cached.Out(ctx, "key1") + if v1 == nil { + t.Fatal("first Out() should return a limiter") + } + if v1.Limit() != 200 { + t.Fatalf("expected limit 200, got %d", v1.Limit()) + } + if inner.outCalls != 1 { + t.Fatalf("expected 1 Out call, got %d", inner.outCalls) + } + + cached.Out(ctx, "key1") + if inner.outCalls != 1 { + t.Fatalf("cached call should not invoke inner, got %d calls", inner.outCalls) + } +} + +func TestCachedTrafficLimiter_Expiration(t *testing.T) { + inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + // Bypass constructor to avoid 1s minimum refresh interval. + cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(10*time.Millisecond)) + if cached == nil { + t.Fatal("NewCachedTrafficLimiter should not return nil") + } + ct := cached.(*cachedTrafficLimiter) + ct.options.refreshInterval = 10 * time.Millisecond + + ctx := context.Background() + ct.In(ctx, "key1") + if inner.inCalls != 1 { + t.Fatalf("expected 1 In call, got %d", inner.inCalls) + } + + // Wait for cache to expire. + time.Sleep(30 * time.Millisecond) + + ct.In(ctx, "key1") + if inner.inCalls != 2 { + t.Fatalf("after expiration, inner In should be called again, got %d calls", inner.inCalls) + } +} + +func TestCachedTrafficLimiter_ScopeFilter(t *testing.T) { + inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + cached := NewCachedTrafficLimiter(inner, + RefreshIntervalOption(50*time.Millisecond), + ScopeOption("service"), + ) + + ctx := context.Background() + + // Non-matching scope: returns nil. + v := cached.In(ctx, "key1", limiter.ScopeOption(limiter.ScopeConn)) + if v != nil { + t.Fatal("non-matching scope should return nil") + } + + // Matching scope: returns limiter. + v = cached.In(ctx, "key1", limiter.ScopeOption(limiter.ScopeService)) + if v == nil { + t.Fatal("matching scope should return limiter") + } +} + +func TestCachedTrafficLimiter_NilInner(t *testing.T) { + cached := NewCachedTrafficLimiter(nil) + if cached != nil { + t.Fatal("NewCachedTrafficLimiter(nil) should return nil") + } +} + +func TestCachedTrafficLimiter_NilNew(t *testing.T) { + inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(10*time.Millisecond)) + if cached == nil { + t.Fatal("NewCachedTrafficLimiter should not return nil") + } + ct := cached.(*cachedTrafficLimiter) + ct.options.refreshInterval = 10 * time.Millisecond + + ctx := context.Background() + v1 := ct.In(ctx, "key1") + if v1 == nil || v1.Limit() != 100 { + t.Fatal("first In() should return limiter with limit 100") + } + + // Now make inner return nil. + inner.inLim = nil + + // Wait for cache to expire. + time.Sleep(30 * time.Millisecond) + + // After expiration: inner returns nil, falls back to cached limiter. + v2 := ct.In(ctx, "key1") + if v2 == nil { + t.Fatal("should fall back to cached limiter when inner returns nil") + } + if v2.Limit() != 100 { + t.Fatalf("expected fallback limit 100, got %d", v2.Limit()) + } +} diff --git a/limiter/traffic/generator_test.go b/limiter/traffic/generator_test.go new file mode 100644 index 00000000..06b0d072 --- /dev/null +++ b/limiter/traffic/generator_test.go @@ -0,0 +1,63 @@ +package traffic + +import ( + "testing" +) + +func TestLimitGenerator_In(t *testing.T) { + g := newLimitGenerator(100, 0) + lim := g.In() + if lim == nil { + t.Fatal("In() should not be nil for positive in") + } + if lim.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", lim.Limit()) + } + + // Each In() call creates a new instance. + lim2 := g.In() + if lim2 == lim { + t.Fatal("each In() call should create a new instance") + } +} + +func TestLimitGenerator_In_Zero(t *testing.T) { + g := newLimitGenerator(0, 200) + if lim := g.In(); lim != nil { + t.Fatal("In() should be nil for zero in") + } +} + +func TestLimitGenerator_Out(t *testing.T) { + g := newLimitGenerator(0, 200) + lim := g.Out() + if lim == nil { + t.Fatal("Out() should not be nil for positive out") + } + if lim.Limit() != 200 { + t.Fatalf("expected limit 200, got %d", lim.Limit()) + } + + // Each Out() call creates a new instance. + lim2 := g.Out() + if lim2 == lim { + t.Fatal("each Out() call should create a new instance") + } +} + +func TestLimitGenerator_Out_Zero(t *testing.T) { + g := newLimitGenerator(100, 0) + if lim := g.Out(); lim != nil { + t.Fatal("Out() should be nil for zero out") + } +} + +func TestLimitGenerator_NilReceiver(t *testing.T) { + var g *limitGenerator + if lim := g.In(); lim != nil { + t.Fatal("nil receiver In() should return nil") + } + if lim := g.Out(); lim != nil { + t.Fatal("nil receiver Out() should return nil") + } +} diff --git a/limiter/traffic/limiter.go b/limiter/traffic/limiter.go index da9dba31..eb9ae12f 100644 --- a/limiter/traffic/limiter.go +++ b/limiter/traffic/limiter.go @@ -14,6 +14,8 @@ type llimiter struct { limiter *rate.Limiter } +// NewLimiter creates a rate-based traffic Limiter with the given rate +// (bytes per second) and burst size. func NewLimiter(r int) limiter.Limiter { return &llimiter{ limiter: rate.NewLimiter(rate.Limit(r), r), diff --git a/limiter/traffic/limiter_test.go b/limiter/traffic/limiter_test.go new file mode 100644 index 00000000..d5cb85e5 --- /dev/null +++ b/limiter/traffic/limiter_test.go @@ -0,0 +1,135 @@ +package traffic + +import ( + "context" + "testing" + + limiter "github.com/go-gost/core/limiter/traffic" +) + +func TestNewLimiter(t *testing.T) { + l := NewLimiter(100) + if l == nil { + t.Fatal("NewLimiter should not be nil") + } + if l.Limit() != 100 { + t.Fatalf("expected limit 100, got %d", l.Limit()) + } +} + +func TestLimiter_Wait(t *testing.T) { + // High burst so Wait does not block. + l := NewLimiter(1000000) + n := l.Wait(context.Background(), 3) + if n != 3 { + t.Fatalf("expected 3, got %d", n) + } +} + +func TestLimiter_Wait_BurstClamped(t *testing.T) { + // Burst is 5, request 100 -> clamped to burst. + l := NewLimiter(5) + n := l.Wait(context.Background(), 100) + if n != 5 { + t.Fatalf("expected burst 5, got %d", n) + } +} + +func TestLimiter_Wait_ContextCancelled(t *testing.T) { + l := NewLimiter(1000) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + n := l.Wait(ctx, 10) + if n != 0 { + t.Fatalf("cancelled context should return 0, got %d", n) + } +} + +func TestLimiter_Set(t *testing.T) { + l := NewLimiter(5) + l.Set(10) + if l.Limit() != 10 { + t.Fatalf("expected limit 10 after Set, got %d", l.Limit()) + } + // Verify burst also updated: Wait should allow up to 10. + n := l.Wait(context.Background(), 10) + if n != 10 { + t.Fatalf("expected burst 10 after Set, got %d", n) + } +} + +func TestLimiter_String(t *testing.T) { + l := NewLimiter(42) + // llimiter implements fmt.Stringer (not part of core interface). + if s := l.(*llimiter).String(); s != "42" { + t.Fatalf("expected '42', got '%s'", s) + } +} + +func TestLimiterGroup_SortedByLimit(t *testing.T) { + g := newLimiterGroup( + NewLimiter(200), + NewLimiter(50), + NewLimiter(100), + ) + if g.Limit() != 50 { + t.Fatalf("expected smallest limit 50, got %d", g.Limit()) + } +} + +func TestLimiterGroup_Wait(t *testing.T) { + // Group: [3, 100000] — smallest limit is 3. + small := NewLimiter(3) + large := NewLimiter(100000) + g := newLimiterGroup(large, small) + + // Wait for 3: small returns 3, large returns 3, min = 3. + n := g.Wait(context.Background(), 3) + if n != 3 { + t.Fatalf("expected 3, got %d", n) + } + + // Wait for 5: small clamps to burst=3, returns 3, large returns 5, min = 3. + n = g.Wait(context.Background(), 5) + if n != 3 { + t.Fatalf("expected min=3, got %d", n) + } +} + +func TestLimiterGroup_Limit(t *testing.T) { + g := newLimiterGroup(NewLimiter(10), NewLimiter(5), NewLimiter(20)) + if g.Limit() != 5 { + t.Fatalf("expected 5, got %d", g.Limit()) + } +} + +func TestLimiterGroup_Set(t *testing.T) { + g := newLimiterGroup(NewLimiter(5)) + g.Set(100) + if g.Limit() != 5 { + t.Fatal("Set should be a no-op on limiterGroup") + } +} + +func TestLimiterGroup_Empty(t *testing.T) { + g := newLimiterGroup() + if g.Limit() != 0 { + t.Fatalf("empty group limit should be 0, got %d", g.Limit()) + } + n := g.Wait(context.Background(), 100) + if n != 100 { + t.Fatalf("empty group Wait should return n unchanged, got %d", n) + } +} + +func TestLimiterGroup_String(t *testing.T) { + g := newLimiterGroup(NewLimiter(5)) + if g.String() != "[5]" { + t.Fatalf("expected '[5]', got '%s'", g.String()) + } +} + +var ( + _ limiter.Limiter = (*llimiter)(nil) + _ limiter.Limiter = (*limiterGroup)(nil) +) diff --git a/limiter/traffic/traffic.go b/limiter/traffic/traffic.go index 21ed5d8d..23504872 100644 --- a/limiter/traffic/traffic.go +++ b/limiter/traffic/traffic.go @@ -1,3 +1,6 @@ +// Package traffic provides a traffic/bandwidth limiter implementation that +// supports per-service, per-connection, per-IP, and per-CIDR rate limiting +// for both ingress (In) and egress (Out) directions. package traffic import ( @@ -20,8 +23,10 @@ import ( ) const ( + // ServiceLimitKey is the special key used for service-level rate limits. ServiceLimitKey = "$" - ConnLimitKey = "$$" + // ConnLimitKey is the special key used for connection-level rate limits. + ConnLimitKey = "$$" ) const ( @@ -38,38 +43,48 @@ type options struct { logger logger.Logger } +// Option configures a TrafficLimiter. type Option func(opts *options) +// LimitsOption sets the static limit rules. Each string is a limit definition +// in the format " [out]" where in/out are byte quantities with +// optional SI/IEC suffixes (e.g., "100B", "1MB", "512KB"). func LimitsOption(limits ...string) Option { return func(opts *options) { opts.limits = limits } } +// ReloadPeriodOption sets the period for reloading limit rules from external +// loaders. A period of zero disables periodic reloading. func ReloadPeriodOption(period time.Duration) Option { return func(opts *options) { opts.period = period } } +// FileLoaderOption sets a file-based loader for limit rules. func FileLoaderOption(fileLoader loader.Loader) Option { return func(opts *options) { opts.fileLoader = fileLoader } } +// RedisLoaderOption sets a Redis-based loader for limit rules. func RedisLoaderOption(redisLoader loader.Loader) Option { return func(opts *options) { opts.redisLoader = redisLoader } } +// HTTPLoaderOption sets an HTTP-based loader for limit rules. func HTTPLoaderOption(httpLoader loader.Loader) Option { return func(opts *options) { opts.httpLoader = httpLoader } } +// LoggerOption sets the logger for the traffic limiter. func LoggerOption(logger logger.Logger) Option { return func(opts *options) { opts.logger = logger @@ -96,6 +111,9 @@ type trafficLimiter struct { cancelFunc context.CancelFunc } +// NewTrafficLimiter creates a new TrafficLimiter with the given options. +// It starts a background goroutine for periodic reload of limit rules from +// external loaders. func NewTrafficLimiter(opts ...Option) traffic.TrafficLimiter { var options options for _, opt := range opts { diff --git a/limiter/traffic/traffic_test.go b/limiter/traffic/traffic_test.go new file mode 100644 index 00000000..e96e0acd --- /dev/null +++ b/limiter/traffic/traffic_test.go @@ -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) +) diff --git a/limiter/traffic/wrapper/conn.go b/limiter/traffic/wrapper/conn.go index 19b85022..cc4d46b7 100644 --- a/limiter/traffic/wrapper/conn.go +++ b/limiter/traffic/wrapper/conn.go @@ -1,3 +1,5 @@ +// Package wrapper provides net.Conn, net.PacketConn, and io.ReadWriter +// wrappers that apply traffic rate limiting to reads and writes. package wrapper import ( @@ -21,7 +23,7 @@ var ( errRateLimited = errors.New("rate limited") ) -// limitConn is a Conn with traffic limiter supported. +// limitConn wraps a net.Conn with traffic rate limiting applied to reads and writes. type limitConn struct { net.Conn rbuf bytes.Buffer @@ -30,6 +32,8 @@ type limitConn struct { key string } +// WrapConn wraps a net.Conn with traffic rate limiting. If tlimiter is nil, +// the original conn is returned unchanged. func WrapConn(c net.Conn, tlimiter traffic.TrafficLimiter, key string, opts ...limiter.Option) net.Conn { if tlimiter == nil { return c @@ -133,6 +137,8 @@ type packetConn struct { key string } +// WrapPacketConn wraps a net.PacketConn with traffic rate limiting. Packets +// exceeding the rate limit are discarded on read or rejected with an error on write. func WrapPacketConn(pc net.PacketConn, lim traffic.TrafficLimiter, key string, opts ...limiter.Option) net.PacketConn { if lim == nil { return pc @@ -190,6 +196,7 @@ type udpConn struct { key string } +// WrapUDPConn wraps a net.PacketConn as a udp.Conn with traffic rate limiting. func WrapUDPConn(pc net.PacketConn, limiter traffic.TrafficLimiter, key string, opts ...limiter.Option) udp.Conn { return &udpConn{ PacketConn: pc, diff --git a/limiter/traffic/wrapper/conn_test.go b/limiter/traffic/wrapper/conn_test.go new file mode 100644 index 00000000..fb4e2f2e --- /dev/null +++ b/limiter/traffic/wrapper/conn_test.go @@ -0,0 +1,448 @@ +package wrapper + +import ( + "context" + "errors" + "io" + "net" + "syscall" + "testing" + "time" + + "github.com/go-gost/core/limiter" + "github.com/go-gost/core/limiter/traffic" + xio "github.com/go-gost/x/internal/io" + ctxutil "github.com/go-gost/x/ctx" +) + +// --- mock types --- + +type mockLimiter struct { + waitFunc func(ctx context.Context, n int) int + limit int +} + +func (m *mockLimiter) Wait(ctx context.Context, n int) int { + if m.waitFunc != nil { + return m.waitFunc(ctx, n) + } + return n +} +func (m *mockLimiter) Limit() int { return m.limit } +func (m *mockLimiter) Set(n int) { m.limit = n } + +type mockTrafficLimiter struct { + inLim traffic.Limiter + outLim traffic.Limiter +} + +func (m *mockTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + return m.inLim +} +func (m *mockTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + return m.outLim +} + +// --- WrapConn tests --- + +func TestWrapConn_NilLimiter(t *testing.T) { + c := &struct{ net.Conn }{} + result := WrapConn(c, nil, "key") + if result != c { + t.Fatal("nil limiter should return original conn") + } +} + +func TestWrapConn_WithLimiter(t *testing.T) { + c := &struct{ net.Conn }{} + tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + result := WrapConn(c, tl, "key") + lc, ok := result.(*limitConn) + if !ok { + t.Fatalf("expected *limitConn, got %T", result) + } + if lc.Conn != c { + t.Fatal("wrapped conn should retain inner conn") + } +} + +func TestLimitConn_Read_NoLimiter(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + tl := &mockTrafficLimiter{inLim: nil} + wrapped := WrapConn(client, tl, "test-key") + + go server.Write([]byte("hello")) + + buf := make([]byte, 10) + n, err := wrapped.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + if n != 5 || string(buf[:5]) != "hello" { + t.Fatalf("expected 'hello', got %q", string(buf[:n])) + } +} + +func TestLimitConn_Read_WithLimiter(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + ml := &mockLimiter{limit: 100} + tl := &mockTrafficLimiter{inLim: ml} + wrapped := WrapConn(client, tl, "test-key") + + go func() { + server.Write([]byte("hello world")) + server.Close() + }() + + buf := make([]byte, 100) + n, err := wrapped.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + if string(buf[:n]) != "hello world" { + t.Fatalf("expected 'hello world', got %q", string(buf[:n])) + } +} + +func TestLimitConn_Read_BufferedData(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + callCount := 0 + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + callCount++ + if callCount == 1 { + return 5 // only allow 5 of 11 bytes + } + return 10 + }, + } + tl := &mockTrafficLimiter{inLim: ml} + wrapped := WrapConn(client, tl, "test-key") + + go func() { + server.Write([]byte("hello world")) + server.Close() + }() + + buf := make([]byte, 100) + n, err := readFull(wrapped, buf, 5) + if err != nil { + t.Fatalf("first read: %v", err) + } + if n != 5 { + t.Fatalf("expected 5 bytes, got %d", n) + } + + // Remaining bytes from rbuf. + n, err = readFull(wrapped, buf[5:], 6) + if err != nil { + t.Fatalf("second read (buffered): %v", err) + } + if n != 6 { + t.Fatalf("expected 6 buffered bytes, got %d", n) + } +} + +func TestLimitConn_Write_NoLimiter(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + tl := &mockTrafficLimiter{outLim: nil} + wrapped := WrapConn(client, tl, "test-key") + + go func() { + buf := make([]byte, 5) + server.Read(buf) + }() + + n, err := wrapped.Write([]byte("hello")) + if err != nil { + t.Fatal(err) + } + if n != 5 { + t.Fatalf("expected 5, got %d", n) + } +} + +func TestLimitConn_Write_ZeroBurst(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + return 0 + }, + } + tl := &mockTrafficLimiter{outLim: ml} + wrapped := WrapConn(client, tl, "test-key") + + n, err := wrapped.Write([]byte("data")) + if err != nil { + t.Fatalf("should not error on zero burst: %v", err) + } + if n != 0 { + t.Fatalf("expected 0 bytes, got %d", n) + } +} + +type connWithSyscall struct { + net.Conn +} + +func (c *connWithSyscall) SyscallConn() (syscall.RawConn, error) { + return nil, nil +} + +func TestLimitConn_SyscallConn(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + c := &connWithSyscall{Conn: client} + tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + wrapped := WrapConn(c, tl, "test-key") + + _, err := wrapped.(*limitConn).SyscallConn() + if err != nil { + t.Fatal("SyscallConn should succeed", err) + } +} + +type connWithCloseRead struct { + net.Conn + closed bool +} + +func (c *connWithCloseRead) CloseRead() error { + c.closed = true + return nil +} + +func TestLimitConn_CloseRead(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + c := &connWithCloseRead{Conn: client} + tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + wrapped := WrapConn(c, tl, "test-key") + + lc := wrapped.(*limitConn) + if err := lc.CloseRead(); err != nil { + t.Fatal("CloseRead should succeed", err) + } + if !c.closed { + t.Fatal("CloseRead should delegate to inner") + } +} + +type connWithCloseWrite struct { + net.Conn + closed bool +} + +func (c *connWithCloseWrite) CloseWrite() error { + c.closed = true + return nil +} + +func TestLimitConn_CloseWrite(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + c := &connWithCloseWrite{Conn: client} + tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}} + wrapped := WrapConn(c, tl, "test-key") + + lc := wrapped.(*limitConn) + if err := lc.CloseWrite(); err != nil { + t.Fatal("CloseWrite should succeed", err) + } + if !c.closed { + t.Fatal("CloseWrite should delegate to inner") + } +} + +// --- WrapReadWriter tests --- + +func TestWrapReadWriter_NilLimiter(t *testing.T) { + rw := &bytesReadWriter{} + result := WrapReadWriter(nil, rw, "key") + if result != rw { + t.Fatal("nil limiter should return original ReadWriter") + } +} + +func TestReadWriter_Read_WithLimiter(t *testing.T) { + data := []byte("hello world") + rw := &bytesReadWriter{buf: append([]byte{}, data...)} + + ml := &mockLimiter{limit: 100} + tl := &mockTrafficLimiter{inLim: ml} + wrapped := WrapReadWriter(tl, rw, "test-key") + + buf := make([]byte, 100) + n, err := wrapped.Read(buf) + if err != nil && err != io.EOF { + t.Fatal(err) + } + if string(buf[:n]) != "hello world" { + t.Fatalf("expected 'hello world', got %q", string(buf[:n])) + } +} + +func TestReadWriter_Write_ZeroBurst(t *testing.T) { + rw := &bytesReadWriter{} + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + return 0 + }, + } + tl := &mockTrafficLimiter{outLim: ml} + wrapped := WrapReadWriter(tl, rw, "test-key") + + n, err := wrapped.Write([]byte("data")) + if err != nil { + t.Fatalf("should not error on zero burst: %v", err) + } + if n != 0 { + t.Fatalf("expected 0 bytes, got %d", n) + } +} + +// --- WrapPacketConn tests --- + +func TestWrapPacketConn_NilLimiter(t *testing.T) { + pc := &struct{ net.PacketConn }{} + result := WrapPacketConn(pc, nil, "key") + if result != pc { + t.Fatal("nil limiter should return original PacketConn") + } +} + +func TestPacketConn_WriteTo_RateLimited(t *testing.T) { + pc := &mockPacketConn{writeBuf: make([]byte, 100)} + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + return 3 // allow only 3, but packet is 4 bytes + }, + } + tl := &mockTrafficLimiter{outLim: ml} + wrapped := WrapPacketConn(pc, tl, "test-key") + + _, err := wrapped.WriteTo([]byte("data"), nil) + if !errors.Is(err, errRateLimited) { + t.Fatalf("expected errRateLimited, got %v", err) + } +} + +// --- WrapUDPConn tests --- + +func TestWrapUDPConn(t *testing.T) { + pc := &mockPacketConn{writeBuf: make([]byte, 100)} + result := WrapUDPConn(pc, &mockTrafficLimiter{}, "key") + if _, ok := result.(*udpConn); !ok { + t.Fatalf("expected *udpConn, got %T", result) + } +} + +func TestUDPConn_Write_RateLimited(t *testing.T) { + pc := &mockPacketConn{writeBuf: make([]byte, 100)} + ml := &mockLimiter{ + limit: 100, + waitFunc: func(ctx context.Context, n int) int { + return 2 // allow fewer bytes than packet + }, + } + tl := &mockTrafficLimiter{outLim: ml} + wrapped := WrapUDPConn(pc, tl, "test-key") + + _, err := wrapped.Write([]byte("dat")) + if !errors.Is(err, errRateLimited) { + t.Fatalf("expected errRateLimited, got %v", err) + } +} + +// --- interface compliance --- + +var ( + _ net.Conn = (*limitConn)(nil) + _ net.Conn = (*struct{ net.Conn })(nil) + _ xio.CloseRead = (*limitConn)(nil) + _ xio.CloseWrite = (*limitConn)(nil) + _ syscall.Conn = (*limitConn)(nil) + _ ctxutil.Context = (*limitConn)(nil) +) + +// --- helpers --- + +type bytesReadWriter struct { + buf []byte +} + +func (rw *bytesReadWriter) Read(b []byte) (int, error) { + if len(rw.buf) == 0 { + return 0, io.EOF + } + n := copy(b, rw.buf) + rw.buf = rw.buf[n:] + return n, nil +} + +func (rw *bytesReadWriter) Write(b []byte) (int, error) { + rw.buf = append(rw.buf, b...) + return len(b), nil +} + +type mockPacketConn struct { + net.PacketConn + writeBuf []byte +} + +func (pc *mockPacketConn) Write(p []byte) (int, error) { + pc.writeBuf = append(pc.writeBuf[:0], p...) + return len(p), nil +} + +func (pc *mockPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) { + pc.writeBuf = append(pc.writeBuf[:0], p...) + return len(p), nil +} + +func (pc *mockPacketConn) ReadFrom(p []byte) (int, net.Addr, error) { + return copy(p, "test"), nil, nil +} + +func (pc *mockPacketConn) Close() error { return nil } +func (pc *mockPacketConn) LocalAddr() net.Addr { return nil } +func (pc *mockPacketConn) SetDeadline(t time.Time) error { return nil } +func (pc *mockPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (pc *mockPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +func readFull(c net.Conn, buf []byte, want int) (int, error) { + total := 0 + for total < want { + n, err := c.Read(buf[total:]) + if err != nil { + return total, err + } + total += n + } + return total, nil +} diff --git a/limiter/traffic/wrapper/io.go b/limiter/traffic/wrapper/io.go index f05122fb..f6898624 100644 --- a/limiter/traffic/wrapper/io.go +++ b/limiter/traffic/wrapper/io.go @@ -9,7 +9,7 @@ import ( "github.com/go-gost/core/limiter/traffic" ) -// readWriter is an io.ReadWriter with traffic limiter supported. +// readWriter wraps an io.ReadWriter with traffic rate limiting applied to reads and writes. type readWriter struct { io.ReadWriter rbuf bytes.Buffer @@ -18,6 +18,8 @@ type readWriter struct { key string } +// WrapReadWriter wraps an io.ReadWriter with traffic rate limiting. If limiter +// is nil, the original ReadWriter is returned unchanged. func WrapReadWriter(limiter traffic.TrafficLimiter, rw io.ReadWriter, key string, opts ...limiter.Option) io.ReadWriter { if limiter == nil { return rw diff --git a/limiter/traffic/wrapper/listener.go b/limiter/traffic/wrapper/listener.go index 8af59e6c..d4e7096f 100644 --- a/limiter/traffic/wrapper/listener.go +++ b/limiter/traffic/wrapper/listener.go @@ -14,6 +14,8 @@ type listener struct { service string } +// WrapListener wraps a net.Listener so that accepted connections are +// automatically wrapped with service-level traffic rate limiting. func WrapListener(service string, ln net.Listener, limiter traffic.TrafficLimiter) net.Listener { if limiter == nil { return ln diff --git a/limiter/traffic/wrapper/listener_test.go b/limiter/traffic/wrapper/listener_test.go new file mode 100644 index 00000000..84c67b1b --- /dev/null +++ b/limiter/traffic/wrapper/listener_test.go @@ -0,0 +1,65 @@ +package wrapper + +import ( + "errors" + "net" + "testing" + + traffic "github.com/go-gost/core/limiter/traffic" +) + +type testListener struct { + connCh chan net.Conn + addr net.Addr +} + +func (l *testListener) Accept() (net.Conn, error) { + c, ok := <-l.connCh + if !ok { + return nil, errors.New("listener closed") + } + return c, nil +} + +func (l *testListener) Close() error { + close(l.connCh) + return nil +} + +func (l *testListener) Addr() net.Addr { return l.addr } + +func TestWrapListener_NilLimiter(t *testing.T) { + ln := &testListener{} + result := WrapListener("svc", ln, nil) + if result != ln { + t.Fatal("nil limiter should return original listener") + } +} + +func TestAccept_WrapsConn(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + ln := &testListener{ + connCh: make(chan net.Conn, 1), + addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080}, + } + ln.connCh <- client + + tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 1000}} + wrappedLn := WrapListener("test-svc", ln, tl) + + conn, err := wrappedLn.Accept() + if err != nil { + t.Fatal(err) + } + + if _, ok := conn.(*limitConn); !ok { + t.Fatalf("accepted conn should be *limitConn, got %T", conn) + } +} + +var ( + _ traffic.TrafficLimiter = (*mockTrafficLimiter)(nil) +)