diff --git a/limiter/conn/conn_test.go b/limiter/conn/conn_test.go new file mode 100644 index 00000000..2ffaa4a1 --- /dev/null +++ b/limiter/conn/conn_test.go @@ -0,0 +1,458 @@ +package conn + +import ( + "context" + "io" + "strings" + "sync" + "testing" + "time" + + limiter "github.com/go-gost/core/limiter/conn" + "github.com/go-gost/core/logger" + xlogger "github.com/go-gost/x/logger" + "github.com/yl2chen/cidranger" +) + +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 +} + +func nopLogger() logger.Logger { + return xlogger.Nop() +} + +func TestParseLine(t *testing.T) { + cl := &connLimiter{} + + tests := []struct { + input string + expected string + }{ + {"192.168.1.1 100", "192.168.1.1 100"}, + {" 192.168.1.1 100 ", "192.168.1.1 100"}, + {"192.168.1.1 100 # comment", "192.168.1.1 100"}, + {"# only comment", ""}, + {" # indented comment", ""}, + {"", ""}, + } + for _, tt := range tests { + got := cl.parseLine(tt.input) + if got != tt.expected { + t.Errorf("parseLine(%q) = %q, want %q", tt.input, got, tt.expected) + } + } +} + +func TestParseLimit(t *testing.T) { + cl := &connLimiter{} + + tests := []struct { + input string + key string + limit int + }{ + {"192.168.1.1 100", "192.168.1.1", 100}, + {"192.168.1.1\t100", "192.168.1.1", 100}, + {"192.168.1.1 100", "192.168.1.1", 100}, + {"$ 50", "$", 50}, + {"$$ 20", "$$", 20}, + {"10.0.0.0/8 500", "10.0.0.0/8", 500}, + {"key only", "key", 0}, + {"key invalid", "key", 0}, + {"", "", 0}, + } + for _, tt := range tests { + key, limit := cl.parseLimit(tt.input) + if key != tt.key || limit != tt.limit { + t.Errorf("parseLimit(%q) = (%q, %d), want (%q, %d)", tt.input, key, limit, tt.key, tt.limit) + } + } +} + +func TestParsePatterns(t *testing.T) { + cl := &connLimiter{} + + patterns, err := cl.parsePatterns(strings.NewReader("192.168.1.1 10\n10.0.0.1 20\n# comment\n\n 172.16.0.1 30 \n")) + if err != nil { + t.Fatal(err) + } + if len(patterns) != 3 { + t.Fatalf("expected 3 patterns, got %d: %v", len(patterns), patterns) + } +} + +func TestParsePatterns_NilReader(t *testing.T) { + cl := &connLimiter{} + patterns, err := cl.parsePatterns(nil) + if err != nil { + t.Fatal(err) + } + if len(patterns) != 0 { + t.Fatal("expected no patterns from nil reader") + } +} + +func TestLoad_NoLoaders(t *testing.T) { + cl := &connLimiter{options: options{}, logger: nopLogger()} + patterns, err := cl.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(patterns) != 0 { + t.Fatalf("expected 0 patterns, got %d", len(patterns)) + } +} + +func TestLoad_FileLister(t *testing.T) { + fl := &fakeLoader{data: "10.0.0.1 10\n# ignored\n10.0.0.2 20"} + cl := &connLimiter{ + options: options{fileLoader: fl}, + logger: nopLogger(), + } + patterns, err := cl.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(patterns) != 2 { + t.Fatalf("expected 2 patterns, got %d: %v", len(patterns), patterns) + } +} + +func TestLoad_RedisLoader(t *testing.T) { + rl := &fakeLoader{data: "10.0.0.1 5"} + cl := &connLimiter{ + options: options{redisLoader: rl}, + logger: nopLogger(), + } + patterns, err := cl.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // redis list items are NOT parseLine-filtered (raw). + if len(patterns) != 1 { + t.Fatalf("expected 1 pattern, got %d", len(patterns)) + } +} + +func TestLoad_HTTPLoader(t *testing.T) { + hl := &fakeLoader{data: "10.0.0.1 5\n10.0.0.2 3"} + cl := &connLimiter{ + options: options{httpLoader: hl}, + logger: nopLogger(), + } + patterns, err := cl.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(patterns) != 2 { + t.Fatalf("expected 2 patterns, got %d", len(patterns)) + } +} + +func TestLoad_FileListerError(t *testing.T) { + fl := &fakeLoader{data: "", listErr: io.ErrUnexpectedEOF} + cl := &connLimiter{ + options: options{fileLoader: fl}, + logger: nopLogger(), + } + patterns, err := cl.load(context.Background()) + if err != nil { + t.Fatal("load should not return error on loader failure") + } + if len(patterns) != 0 { + t.Fatalf("expected 0 patterns, got %d", len(patterns)) + } +} + +func TestReload_StaticLimits(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cl := &connLimiter{ + options: options{ + limits: []string{"192.168.1.1 3", "10.0.0.0/8 10", "$ 100"}, + }, + ipLimits: make(map[string]ConnLimitGenerator), + cidrLimits: cidranger.NewPCTrieRanger(), + limits: make(map[string]limiter.Limiter), + logger: nopLogger(), + cancelFunc: cancel, + } + + if err := cl.reload(ctx); err != nil { + t.Fatal(err) + } + + // Global limit applies to any key. + lim := cl.Limiter("any-key") + if lim == nil { + t.Fatal("global limiter should exist") + } + if lim.Limit() != 100 { + t.Fatalf("expected global limit 100, got %d", lim.Limit()) + } +} + +// newTestConnLimiter creates a connLimiter with sync reload, no goroutine. +func newTestConnLimiter(limits ...string) *connLimiter { + ctx, cancel := context.WithCancel(context.Background()) + cl := &connLimiter{ + options: options{limits: limits}, + ipLimits: make(map[string]ConnLimitGenerator), + cidrLimits: cidranger.NewPCTrieRanger(), + limits: make(map[string]limiter.Limiter), + logger: nopLogger(), + cancelFunc: cancel, + } + _ = cl.reload(ctx) + return cl +} + +func TestReload_ClearsCache(t *testing.T) { + cl := newTestConnLimiter("192.168.1.1 5") + defer cl.Close() + + lim := cl.Limiter("192.168.1.1") + if lim == nil || lim.Limit() != 5 { + t.Fatal("limiter should exist with limit 5") + } + + // Reload clears cache and re-parses static limits. + if err := cl.reload(context.Background()); err != nil { + t.Fatal(err) + } + + lim2 := cl.Limiter("192.168.1.1") + if lim2 == nil { + t.Fatal("limiter should still exist after reload") + } + if lim2.Limit() != 5 { + t.Fatalf("expected limit 5 after reload, got %d", lim2.Limit()) + } +} + +func TestLimiter_NilCacheDoesNotStick(t *testing.T) { + cl := newTestConnLimiter() + defer cl.Close() + + // Lookup for a key with no matching limit should NOT cache nil. + lim := cl.Limiter("10.0.0.1") + if lim != nil { + t.Fatal("should be no limiter for 10.0.0.1") + } + + // Manually add a limit via reload. + cl.options.limits = []string{"10.0.0.1 3"} + if err := cl.reload(context.Background()); err != nil { + t.Fatal(err) + } + + // Should now find the limit because nil was not cached. + lim = cl.Limiter("10.0.0.1") + if lim == nil { + t.Fatal("limiter should exist after reload when nil was not cached") + } + if lim.Limit() != 3 { + t.Fatalf("expected limit 3, got %d", lim.Limit()) + } +} + +func TestLimiter_IPLimitKey(t *testing.T) { + cl := newTestConnLimiter("$$ 5") + defer cl.Close() + + lim := cl.Limiter("10.0.0.1") + if lim == nil { + t.Fatal("IP-level limiter should exist") + } + if lim.Limit() != 5 { + t.Fatalf("expected limit 5, got %d", lim.Limit()) + } +} + +func TestLimiter_CIDRMatch(t *testing.T) { + cl := newTestConnLimiter("10.0.0.0/8 7") + defer cl.Close() + + lim := cl.Limiter("10.1.2.3") + if lim == nil { + t.Fatal("CIDR-based limiter should exist for 10.1.2.3") + } + if lim.Limit() != 7 { + t.Fatalf("expected limit 7, got %d", lim.Limit()) + } +} + +func TestLimiter_CIDRNoMatch(t *testing.T) { + cl := newTestConnLimiter("10.0.0.0/8 7") + defer cl.Close() + + lim := cl.Limiter("192.168.1.1") + if lim != nil { + t.Fatal("CIDR-based limiter should not match 192.168.1.1") + } +} + +func TestLimiter_SpecificIPTakesPrecedence(t *testing.T) { + cl := newTestConnLimiter("10.0.0.0/8 5", "10.1.2.3 2") + defer cl.Close() + + lim := cl.Limiter("10.1.2.3") + if lim == nil { + t.Fatal("specific IP limiter should exist") + } + if lim.Limit() != 2 { + t.Fatalf("expected specific IP limit 2, got %d", lim.Limit()) + } +} + +func TestLimiter_GlobalAndIP(t *testing.T) { + cl := newTestConnLimiter("$$ 5", "$ 10") + defer cl.Close() + + lim := cl.Limiter("10.0.0.1") + if lim == nil { + t.Fatal("should have IP+global limiter") + } + if lim.Limit() != 5 { + t.Fatalf("expected limit 5 (min of 5,10), got %d", lim.Limit()) + } +} + +func TestLimiter_NonIPKeyHasGlobalOnly(t *testing.T) { + cl := newTestConnLimiter("$$ 5", "$ 10") + defer cl.Close() + + lim := cl.Limiter("not-an-ip") + if lim == nil { + t.Fatal("should have limiters") + } + // IPLimitKey ($$=5) acts as fallback for any key when no specific IP match. + // Global ($=10) also applies. Group picks the minimum: 5. + if lim.Limit() != 5 { + t.Fatalf("expected limit 5 (min of IP fallback 5 and global 10), got %d", lim.Limit()) + } +} + +func TestClose(t *testing.T) { + fl := &fakeLoader{} + rl := &fakeLoader{} + hl := &fakeLoader{} + + cl := NewConnLimiter( + FileLoaderOption(fl), + RedisLoaderOption(rl), + HTTPLoaderOption(hl), + ).(*connLimiter) + + cl.Close() + + if !fl.closed { + t.Error("fileLoader was not closed") + } + if !rl.closed { + t.Error("redisLoader was not closed") + } + if !hl.closed { + t.Error("httpLoader was not closed") + } +} + +func TestClose_Idempotent(t *testing.T) { + fl := &fakeLoader{} + cl := NewConnLimiter(FileLoaderOption(fl)).(*connLimiter) + cl.Close() + cl.Close() +} + +func TestLimiter_CachedResult(t *testing.T) { + cl := newTestConnLimiter("192.168.1.1 5") + defer cl.Close() + + lim1 := cl.Limiter("192.168.1.1") + if lim1 == nil { + t.Fatal("first lookup should return a limiter") + } + + lim2 := cl.Limiter("192.168.1.1") + if lim2 != lim1 { + t.Fatal("second lookup should return cached limiter") + } +} + +func TestLimiter_ConcurrentAccess(t *testing.T) { + cl := newTestConnLimiter("192.168.1.1 1000") + defer cl.Close() + + var wg sync.WaitGroup + n := 50 + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = cl.Limiter("192.168.1.1") + }() + } + wg.Wait() +} + +func TestConnLimiter_Interface(t *testing.T) { + var cl limiter.ConnLimiter = NewConnLimiter() + if cl == nil { + t.Fatal("ConnLimiter should not be nil") + } +} + +func TestPeriodReload_ZeroPeriod(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cl := &connLimiter{ + options: options{period: 0}, + logger: nopLogger(), + } + + err := cl.periodReload(ctx) + if err != nil { + t.Fatal(err) + } +} + +func TestPeriodReload_ContextCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cl := &connLimiter{ + options: options{period: 10 * time.Second}, + logger: nopLogger(), + } + cancel() + err := cl.periodReload(ctx) + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } +} diff --git a/limiter/conn/generator_test.go b/limiter/conn/generator_test.go new file mode 100644 index 00000000..a391fbc7 --- /dev/null +++ b/limiter/conn/generator_test.go @@ -0,0 +1,84 @@ +package conn + +import ( + "testing" +) + +func TestConnLimitGenerator(t *testing.T) { + g := NewConnLimitGenerator(5) + lim := g.Limiter() + if lim == nil { + t.Fatal("Limiter() should not be nil for non-zero limit") + } + if lim.Limit() != 5 { + t.Fatalf("expected limit 5, got %d", lim.Limit()) + } + + // Each Limiter() call creates a new instance. + lim2 := g.Limiter() + if lim2 == lim { + t.Fatal("each Limiter() call should create a new instance") + } +} + +func TestConnLimitGenerator_Zero(t *testing.T) { + g := NewConnLimitGenerator(0) + if g == nil { + t.Fatal("generator should not be nil") + } + if lim := g.Limiter(); lim != nil { + t.Fatal("Limiter() should be nil for zero limit") + } +} + +func TestConnLimitGenerator_Negative(t *testing.T) { + g := NewConnLimitGenerator(-5) + if lim := g.Limiter(); lim != nil { + t.Fatal("Limiter() should be nil for negative limit") + } +} + +func TestConnLimitGenerator_NilReceiver(t *testing.T) { + var g *connLimitGenerator + if lim := g.Limiter(); lim != nil { + t.Fatal("nil receiver Limiter() should return nil") + } +} + +func TestConnLimitSingleGenerator(t *testing.T) { + g := NewConnLimitSingleGenerator(5) + lim := g.Limiter() + if lim == nil { + t.Fatal("Limiter() should not be nil") + } + if lim.Limit() != 5 { + t.Fatalf("expected limit 5, got %d", lim.Limit()) + } + + // Same instance every call. + lim2 := g.Limiter() + if lim2 != lim { + t.Fatal("single generator should return same instance") + } +} + +func TestConnLimitSingleGenerator_Zero(t *testing.T) { + g := NewConnLimitSingleGenerator(0) + if lim := g.Limiter(); lim != nil { + t.Fatal("Limiter() should be nil for zero limit") + } +} + +func TestConnLimitSingleGenerator_Negative(t *testing.T) { + g := NewConnLimitSingleGenerator(-1) + if lim := g.Limiter(); lim != nil { + t.Fatal("Limiter() should be nil for negative limit") + } +} + +func TestConnLimitSingleGenerator_NilReceiver(t *testing.T) { + var g *connLimitSingleGenerator + if lim := g.Limiter(); lim != nil { + t.Fatal("nil receiver Limiter() should return nil") + } +} diff --git a/limiter/conn/limiter_test.go b/limiter/conn/limiter_test.go new file mode 100644 index 00000000..7b50319d --- /dev/null +++ b/limiter/conn/limiter_test.go @@ -0,0 +1,155 @@ +package conn + +import ( + "sync" + "sync/atomic" + "testing" + + limiter "github.com/go-gost/core/limiter/conn" +) + +func TestLimiter_Allow(t *testing.T) { + l := NewLimiter(3).(*llimiter) + + if l.Limit() != 3 { + t.Fatalf("expected limit 3, got %d", l.Limit()) + } + + // Basic acquire/release within limit. + for i := 0; i < 3; i++ { + if !l.Allow(1) { + t.Fatalf("Allow(1) #%d should succeed", i) + } + } + if l.Allow(1) { + t.Fatal("Allow(1) should fail after reaching limit") + } + + // Release one and re-acquire. + if !l.Allow(-1) { + t.Fatal("Allow(-1) should always succeed") + } + if !l.Allow(1) { + t.Fatal("Allow(1) should succeed after release") + } +} + +func TestLimiter_AllowZero(t *testing.T) { + l := NewLimiter(1) + if !l.Allow(0) { + t.Fatal("Allow(0) should not fail with no connections") + } +} + +func TestLimiter_AllowNegative(t *testing.T) { + l := NewLimiter(1) + // Allow(-1) should always return true even without prior acquire. + if !l.Allow(-1) { + t.Fatal("Allow(-1) should succeed") + } + // Allow(-1) again — still succeeds. + if !l.Allow(-1) { + t.Fatal("Allow(-1) should succeed again") + } +} + +func TestLimiter_Concurrent(t *testing.T) { + l := NewLimiter(100) + var wg sync.WaitGroup + n := 200 + var success int64 + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if l.Allow(1) { + atomic.AddInt64(&success, 1) + } + }() + } + wg.Wait() + if v := atomic.LoadInt64(&success); v != 100 { + t.Errorf("expected 100 successes, got %d", v) + } +} + +func TestLimiterGroup_Allow(t *testing.T) { + lg := newLimiterGroup(NewLimiter(2), NewLimiter(5), NewLimiter(10)) + + if lg.Limit() != 2 { + t.Fatalf("expected group limit 2 (smallest), got %d", lg.Limit()) + } + + // Each Allow must pass all limiters. + for i := 0; i < 2; i++ { + if !lg.Allow(1) { + t.Fatalf("lg.Allow(1) #%d should succeed", i) + } + } + if lg.Allow(1) { + t.Fatal("lg.Allow(1) should fail after 2 connections (tightest limit=2)") + } +} + +func TestLimiterGroup_AllowRelease(t *testing.T) { + lg := newLimiterGroup(NewLimiter(1), NewLimiter(10)) + + if !lg.Allow(1) { + t.Fatal("Allow(1) should succeed") + } + if lg.Allow(1) { + t.Fatal("Allow(1) should fail") + } + + // Release on group must propagate to all limiters. + if !lg.Allow(-1) { + t.Fatal("Allow(-1) should succeed") + } + + // Should be able to acquire again. + if !lg.Allow(1) { + t.Fatal("Allow(1) should succeed after release") + } +} + +func TestLimiterGroup_AllowRollback(t *testing.T) { + // limiters sorted: [1, 5] (ascending by limit). + lg := newLimiterGroup(NewLimiter(5), NewLimiter(1)) + + // First Allow passes both limiters. + if !lg.Allow(1) { + t.Fatal("first Allow(1) should succeed") + } + + // Second Allow fails on the tightest limiter (limit=1) and rolls back. + if lg.Allow(1) { + t.Fatal("second Allow(1) should fail after limit reached") + } + + // The tightest limiter should still be at 1 (rollback worked). + ll := lg.limiters[0].(*llimiter) + if ll.current != 1 { + t.Errorf("first limiter should be at 1 after rollback, got %d", ll.current) + } +} + +func TestLimiterGroup_BadInput(t *testing.T) { + lg := newLimiterGroup(NewLimiter(-5), NewLimiter(10)) + if lg.Limit() != -5 { + t.Fatalf("expected limit -5 for negative-limit limiter") + } +} + +func TestLimiterGroup_Empty(t *testing.T) { + lg := newLimiterGroup() + if lg.Limit() != 0 { + t.Fatal("empty group limit should be 0") + } + // Empty group returns false (no limiters to consult). + if lg.Allow(1) { + t.Fatal("Allow(1) on empty group should return false") + } +} + +var _ limiter.Limiter = (*llimiter)(nil) +var _ limiter.Limiter = (*limiterGroup)(nil) diff --git a/limiter/conn/wrapper/conn_test.go b/limiter/conn/wrapper/conn_test.go new file mode 100644 index 00000000..269bcf9c --- /dev/null +++ b/limiter/conn/wrapper/conn_test.go @@ -0,0 +1,193 @@ +package wrapper + +import ( + "context" + "io" + "net" + "syscall" + "testing" + "time" + + xio "github.com/go-gost/x/internal/io" +) + +type connWithCloseRead struct { + net.Conn + closeReadCalled bool +} + +func (c *connWithCloseRead) CloseRead() error { + c.closeReadCalled = true + return nil +} + +type connWithCloseWrite struct { + net.Conn + closeWriteCalled bool +} + +func (c *connWithCloseWrite) CloseWrite() error { + c.closeWriteCalled = true + return nil +} + +type connWithSyscall struct { + net.Conn +} + +func (c *connWithSyscall) SyscallConn() (syscall.RawConn, error) { + return nil, nil +} + +type connWithContext struct { + net.Conn +} + +func (c *connWithContext) Context() context.Context { + return context.Background() +} + +func TestWrapConn_NilLimiter(t *testing.T) { + c := &mockConn{} + wc := WrapConn(nil, c) + if wc != c { + t.Fatal("WrapConn should return original conn when limiter is nil") + } +} + +func TestWrapConn_WithLimiter(t *testing.T) { + c := &mockConn{} + lim := &allowLimiter{limit: 5} + wc := WrapConn(lim, c) + + sc, ok := wc.(*serverConn) + if !ok { + t.Fatal("WrapConn should return *serverConn when limiter is not nil") + } + if sc.Conn != c { + t.Fatal("inner conn should be the original") + } + if sc.limiter != lim { + t.Fatal("limiter should be set") + } +} + +func TestServerConn_Close_ReleasesLimiter(t *testing.T) { + c := &mockConn{} + lim := &allowLimiter{limit: 1} + lim.Allow(1) // acquire one + + sc := &serverConn{Conn: c, limiter: lim} + if err := sc.Close(); err != nil { + t.Fatal(err) + } + if !c.closed { + t.Fatal("inner conn should be closed") + } + // After Close, the limiter should have one slot released. + if !lim.Allow(1) { + t.Fatal("limiter should allow after close released a slot") + } +} + +func TestServerConn_CloseRead_Supported(t *testing.T) { + inner := &connWithCloseRead{Conn: &mockConn{}} + sc := &serverConn{Conn: inner} + if err := sc.CloseRead(); err != nil { + t.Fatal(err) + } + if !inner.closeReadCalled { + t.Fatal("CloseRead should delegate to inner conn") + } +} + +func TestServerConn_CloseRead_Unsupported(t *testing.T) { + inner := &mockConn{} + sc := &serverConn{Conn: inner} + err := sc.CloseRead() + if err != xio.ErrUnsupported { + t.Fatalf("expected ErrUnsupported, got %v", err) + } +} + +func TestServerConn_CloseWrite_Supported(t *testing.T) { + inner := &connWithCloseWrite{Conn: &mockConn{}} + sc := &serverConn{Conn: inner} + if err := sc.CloseWrite(); err != nil { + t.Fatal(err) + } + if !inner.closeWriteCalled { + t.Fatal("CloseWrite should delegate to inner conn") + } +} + +func TestServerConn_CloseWrite_Unsupported(t *testing.T) { + inner := &mockConn{} + sc := &serverConn{Conn: inner} + err := sc.CloseWrite() + if err != xio.ErrUnsupported { + t.Fatalf("expected ErrUnsupported, got %v", err) + } +} + +func TestServerConn_SyscallConn_Supported(t *testing.T) { + inner := &connWithSyscall{Conn: &mockConn{}} + sc := &serverConn{Conn: inner} + _, err := sc.SyscallConn() + if err != nil { + t.Fatalf("expected success, got %v", err) + } +} + +func TestServerConn_SyscallConn_Unsupported(t *testing.T) { + inner := &mockConn{} + sc := &serverConn{Conn: inner} + _, err := sc.SyscallConn() + if err != errUnsupport { + t.Fatalf("expected errUnsupport, got %v", err) + } +} + +func TestServerConn_Context_Supported(t *testing.T) { + inner := &connWithContext{Conn: &mockConn{}} + sc := &serverConn{Conn: inner} + ctx := sc.Context() + if ctx == nil { + t.Fatal("Context should not be nil when inner conn supports it") + } +} + +func TestServerConn_Context_Unsupported(t *testing.T) { + inner := &mockConn{} + sc := &serverConn{Conn: inner} + ctx := sc.Context() + if ctx == nil { + t.Fatal("Context should return context.Background(), not nil") + } +} + +func TestServerConn_ReadWrite(t *testing.T) { + inner := &mockConn{} + sc := &serverConn{Conn: inner} + + n, err := sc.Read(nil) + if err != nil { + t.Fatal(err) + } + if n != 0 { + t.Fatalf("expected 0, got %d", n) + } + + n, err = sc.Write([]byte("hello")) + if err != nil { + t.Fatal(err) + } + if n != 5 { + t.Fatalf("expected 5, got %d", n) + } +} + +// Ensure unused imports are fine: +var _ io.ReadWriter = (*serverConn)(nil) +var _ net.Conn = (*serverConn)(nil) +var _ = time.Now diff --git a/limiter/conn/wrapper/listener_test.go b/limiter/conn/wrapper/listener_test.go new file mode 100644 index 00000000..8172ee8c --- /dev/null +++ b/limiter/conn/wrapper/listener_test.go @@ -0,0 +1,147 @@ +package wrapper + +import ( + "errors" + "net" + "strings" + "testing" + "time" + + limiter "github.com/go-gost/core/limiter/conn" +) + +type mockConn struct { + closed bool +} + +func (c *mockConn) Read(b []byte) (int, error) { return 0, nil } +func (c *mockConn) Write(b []byte) (int, error) { return len(b), nil } +func (c *mockConn) Close() error { c.closed = true; return nil } +func (c *mockConn) LocalAddr() net.Addr { return nil } +func (c *mockConn) RemoteAddr() net.Addr { return &mockAddr{addr: "192.168.1.1:1234"} } +func (c *mockConn) SetDeadline(t time.Time) error { return nil } +func (c *mockConn) SetReadDeadline(t time.Time) error { return nil } +func (c *mockConn) SetWriteDeadline(t time.Time) error { return nil } + +type mockAddr struct { + addr string +} + +func (a *mockAddr) Network() string { return "tcp" } +func (a *mockAddr) String() string { return a.addr } + +type mockListener struct { + conns []net.Conn + current int +} + +func (l *mockListener) Accept() (net.Conn, error) { + if l.current >= len(l.conns) { + return nil, errors.New("no more conns") + } + c := l.conns[l.current] + l.current++ + return c, nil +} + +func (l *mockListener) Close() error { return nil } +func (l *mockListener) Addr() net.Addr { return nil } + +type allowLimiter struct { + limit int + count int +} + +func (l *allowLimiter) Allow(n int) bool { + if l.count+n > l.limit { + return false + } + l.count += n + return true +} +func (l *allowLimiter) Limit() int { return l.limit } + +type connLimiter struct { + lims map[string]limiter.Limiter +} + +func (cl *connLimiter) Limiter(key string) limiter.Limiter { + return cl.lims[key] +} + +func TestWrapListener_NilLimiter(t *testing.T) { + inner := &mockListener{} + ln := WrapListener(nil, inner) + if ln != inner { + t.Fatal("WrapListener should return the inner listener when limiter is nil") + } +} + +func TestAccept_NoLimiterForKey(t *testing.T) { + inner := &mockListener{conns: []net.Conn{&mockConn{}}} + cl := &connLimiter{lims: map[string]limiter.Limiter{}} + ln := WrapListener(cl, inner) + + c, err := ln.Accept() + if err != nil { + t.Fatal(err) + } + if c == nil { + t.Fatal("conn should not be nil") + } +} + +func TestAccept_AllowSuccess(t *testing.T) { + inner := &mockListener{conns: []net.Conn{&mockConn{}}} + cl := &connLimiter{lims: map[string]limiter.Limiter{ + "192.168.1.1": &allowLimiter{limit: 1}, + }} + ln := WrapListener(cl, inner) + + c, err := ln.Accept() + if err != nil { + t.Fatal(err) + } + if c == nil { + t.Fatal("conn should not be nil") + } + + // Should be wrapped in serverConn. + if _, ok := c.(*serverConn); !ok { + t.Fatal("conn should be wrapped as *serverConn") + } +} + +func TestAccept_LimitExceeded(t *testing.T) { + mc := &mockConn{} + inner := &mockListener{conns: []net.Conn{mc}} + cl := &connLimiter{lims: map[string]limiter.Limiter{ + "192.168.1.1": &allowLimiter{limit: 0}, // always deny + }} + ln := WrapListener(cl, inner) + + c, err := ln.Accept() + if err == nil { + t.Fatal("expected error when limit exceeded") + } + if !strings.Contains(err.Error(), "limit exceeded") { + t.Fatalf("unexpected error: %v", err) + } + if c != nil { + t.Fatal("conn should be nil when limit exceeded") + } + if !mc.closed { + t.Fatal("conn should be closed when rejected") + } +} + +func TestAccept_ListenerError(t *testing.T) { + inner := &mockListener{} // no conns, Accept returns error + cl := &connLimiter{lims: map[string]limiter.Limiter{}} + ln := WrapListener(cl, inner) + + _, err := ln.Accept() + if err == nil { + t.Fatal("expected error from listener") + } +}