From 689ba36e924fc2a53090614471b317ad8a520964 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sun, 24 May 2026 23:47:48 +0800 Subject: [PATCH] fix(observer): remove dead nil checks from udpConn, add unit tests The WrapUDPConn constructor returns nil when either arg is nil, so c.stats is always non-nil in all udpConn methods. Remove 8 redundant nil guards that were inconsistent with packetConn (which already skipped them). Add 65 unit tests across plugin, stats, and wrapper packages covering nil safety, race conditions, double-close, event marshaling, and reset-traffic semantics. Replace magic numbers in tests with named Kind constants. --- observer/plugin/grpc_test.go | 203 +++++++++++ observer/plugin/http_test.go | 244 +++++++++++++ observer/stats/stats_test.go | 200 +++++++++++ observer/stats/wrapper/conn.go | 32 +- observer/stats/wrapper/conn_test.go | 441 ++++++++++++++++++++++++ observer/stats/wrapper/io_test.go | 90 +++++ observer/stats/wrapper/listener_test.go | 81 +++++ 7 files changed, 1267 insertions(+), 24 deletions(-) create mode 100644 observer/plugin/grpc_test.go create mode 100644 observer/plugin/http_test.go create mode 100644 observer/stats/stats_test.go create mode 100644 observer/stats/wrapper/conn_test.go create mode 100644 observer/stats/wrapper/io_test.go create mode 100644 observer/stats/wrapper/listener_test.go diff --git a/observer/plugin/grpc_test.go b/observer/plugin/grpc_test.go new file mode 100644 index 00000000..0560ae35 --- /dev/null +++ b/observer/plugin/grpc_test.go @@ -0,0 +1,203 @@ +package observer + +import ( + "context" + "io" + "os" + "testing" + + corelogger "github.com/go-gost/core/logger" + "github.com/go-gost/core/observer" + "github.com/go-gost/plugin/observer/proto" + xlogger "github.com/go-gost/x/logger" + xstats "github.com/go-gost/x/observer/stats" + "github.com/go-gost/x/service" + "google.golang.org/grpc" +) + +func TestMain(m *testing.M) { + corelogger.SetDefault(xlogger.Nop()) + os.Exit(m.Run()) +} + +// mockGrpcClient implements proto.ObserverClient for testing. +type mockGrpcClient struct { + observeReply *proto.ObserveReply + observeErr error +} + +func (m *mockGrpcClient) Observe(ctx context.Context, in *proto.ObserveRequest, opts ...grpc.CallOption) (*proto.ObserveReply, error) { + if m.observeErr != nil { + return nil, m.observeErr + } + return m.observeReply, nil +} + +type mockGrpcConn struct { + closed bool +} + +func (m *mockGrpcConn) Invoke(ctx context.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error { + return nil +} + +func (m *mockGrpcConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return nil, nil +} + +func (m *mockGrpcConn) Close() error { + m.closed = true + return nil +} + +func TestNewGRPCPlugin_ReturnsNonNil(t *testing.T) { + // NewGRPCPlugin should always return a non-nil observer. Even if the + // underlying gRPC dial would eventually fail, modern gRPC uses lazy + // connect so the constructor always returns a valid instance. + p := NewGRPCPlugin("test", "127.0.0.1:65535") + if p == nil { + t.Fatal("NewGRPCPlugin should never return nil") + } +} + +func TestGrpcPlugin_Observe_NoOpClient_DoesNotPanic(t *testing.T) { + // When client is nil (e.g., from connection error or deliberate no-op), + // Observe should return nil without panicking. + p := &grpcPlugin{client: nil} + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err != nil { + t.Errorf("no-op observer should return nil on Observe, got %v", err) + } +} + +func TestGrpcPlugin_Observe_NilClient(t *testing.T) { + p := &grpcPlugin{ + conn: nil, + client: nil, + } + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err != nil { + t.Errorf("nil client Observe should return nil, got %v", err) + } +} + +func TestGrpcPlugin_Observe_EmptyEvents(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{}, + } + // Empty events should return nil without calling the client + err := p.Observe(context.Background(), nil) + if err != nil { + t.Errorf("empty events should return nil, got %v", err) + } +} + +func TestGrpcPlugin_Observe_StatusEvent(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{ + observeReply: &proto.ObserveReply{Ok: true}, + }, + } + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{ + Kind: "service", + Service: "test-svc", + State: "running", + Msg: "started", + }, + }) + if err != nil { + t.Errorf("Observe with status event failed: %v", err) + } +} + +func TestGrpcPlugin_Observe_StatsEvent(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{ + observeReply: &proto.ObserveReply{Ok: true}, + }, + } + err := p.Observe(context.Background(), []observer.Event{ + xstats.StatsEvent{ + Kind: "service", + Service: "test-svc", + Client: "client1", + TotalConns: 10, + CurrentConns: 3, + InputBytes: 1024, + OutputBytes: 2048, + TotalErrs: 1, + }, + }) + if err != nil { + t.Errorf("Observe with stats event failed: %v", err) + } +} + +func TestGrpcPlugin_Observe_ServerError(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{ + observeErr: io.ErrUnexpectedEOF, + }, + log: xlogger.Nop(), + } + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error when server fails") + } +} + +func TestGrpcPlugin_Observe_ReplyNotOk(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{ + observeReply: &proto.ObserveReply{Ok: false}, + }, + } + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error when reply is not ok") + } +} + +func TestGrpcPlugin_Observe_NilReply(t *testing.T) { + p := &grpcPlugin{ + client: &mockGrpcClient{ + observeReply: nil, + }, + } + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error when reply is nil") + } +} + +func TestGrpcPlugin_Close_NilConn(t *testing.T) { + p := &grpcPlugin{conn: nil} + if err := p.Close(); err != nil { + t.Errorf("Close with nil conn should return nil, got %v", err) + } +} + +func TestGrpcPlugin_Close_WithConn(t *testing.T) { + conn := &mockGrpcConn{} + p := &grpcPlugin{conn: conn} + if err := p.Close(); err != nil { + t.Errorf("Close failed: %v", err) + } + if !conn.closed { + t.Error("Close should close the underlying connection") + } +} + +// Test that grpcPlugin satisfies observer.Observer +var _ observer.Observer = (*grpcPlugin)(nil) diff --git a/observer/plugin/http_test.go b/observer/plugin/http_test.go new file mode 100644 index 00000000..2e150e17 --- /dev/null +++ b/observer/plugin/http_test.go @@ -0,0 +1,244 @@ +package observer + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/go-gost/core/observer" + "github.com/go-gost/x/internal/plugin" + xstats "github.com/go-gost/x/observer/stats" + "github.com/go-gost/x/service" +) + +func TestNewHTTPPlugin(t *testing.T) { + p := NewHTTPPlugin("test", "localhost:8080") + if p == nil { + t.Fatal("NewHTTPPlugin should not return nil") + } + hp, ok := p.(*httpPlugin) + if !ok { + t.Fatalf("NewHTTPPlugin returned %T, want *httpPlugin", p) + } + if hp.url != "http://localhost:8080" { + t.Errorf("url = %s, want http://localhost:8080", hp.url) + } + if hp.client == nil { + t.Error("client should not be nil") + } +} + +func TestNewHTTPPlugin_NoScheme(t *testing.T) { + p := NewHTTPPlugin("test", "example.com:9999") + hp := p.(*httpPlugin) + if !strings.HasPrefix(hp.url, "http://") { + t.Errorf("expected http:// prefix, got %s", hp.url) + } +} + +func TestNewHTTPPlugin_KeepsHTTPSScheme(t *testing.T) { + p := NewHTTPPlugin("test", "https://example.com/api") + hp := p.(*httpPlugin) + if !strings.HasPrefix(hp.url, "https://") { + t.Errorf("expected https:// prefix, got %s", hp.url) + } +} + +func TestHTTPPlugin_Observe_StatusEvent(t *testing.T) { + var receivedBody json.RawMessage + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + if ct := r.Header.Get("Content-Type"); ct != "application/json" { + t.Errorf("Content-Type = %s, want application/json", ct) + } + json.NewDecoder(r.Body).Decode(&receivedBody) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{ + Kind: "service", + Service: "test-svc", + State: "running", + Msg: "started", + }, + }) + if err != nil { + t.Fatal(err) + } + + // Verify the event was sent correctly + var req observeRequest + if err := json.Unmarshal(receivedBody, &req); err != nil { + t.Fatal(err) + } + if len(req.Events) != 1 { + t.Fatalf("expected 1 event, got %d", len(req.Events)) + } + if req.Events[0].Kind != "service" { + t.Errorf("event kind = %s, want service", req.Events[0].Kind) + } + if req.Events[0].Service != "test-svc" { + t.Errorf("event service = %s, want test-svc", req.Events[0].Service) + } + if req.Events[0].Status == nil { + t.Fatal("status event should have Status field") + } + if req.Events[0].Status.State != "running" { + t.Errorf("state = %s, want running", req.Events[0].Status.State) + } +} + +func TestHTTPPlugin_Observe_StatsEvent(t *testing.T) { + var receivedBody json.RawMessage + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewDecoder(r.Body).Decode(&receivedBody) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + err := p.Observe(context.Background(), []observer.Event{ + xstats.StatsEvent{ + Kind: "traffic", + Service: "svc", + Client: "c1", + TotalConns: 100, + CurrentConns: 5, + InputBytes: 4096, + OutputBytes: 8192, + TotalErrs: 2, + }, + }) + if err != nil { + t.Fatal(err) + } + + var req observeRequest + if err := json.Unmarshal(receivedBody, &req); err != nil { + t.Fatal(err) + } + if len(req.Events) != 1 { + t.Fatalf("expected 1 event, got %d", len(req.Events)) + } + if req.Events[0].Stats == nil { + t.Fatal("stats event should have Stats field") + } + if req.Events[0].Stats.TotalConns != 100 { + t.Errorf("totalConns = %d, want 100", req.Events[0].Stats.TotalConns) + } + if req.Events[0].Stats.CurrentConns != 5 { + t.Errorf("currentConns = %d, want 5", req.Events[0].Stats.CurrentConns) + } + if req.Events[0].Stats.InputBytes != 4096 { + t.Errorf("inputBytes = %d, want 4096", req.Events[0].Stats.InputBytes) + } + if req.Events[0].Stats.OutputBytes != 8192 { + t.Errorf("outputBytes = %d, want 8192", req.Events[0].Stats.OutputBytes) + } + if req.Events[0].Stats.TotalErrs != 2 { + t.Errorf("totalErrs = %d, want 2", req.Events[0].Stats.TotalErrs) + } + if req.Events[0].Client != "c1" { + t.Errorf("client = %s, want c1", req.Events[0].Client) + } +} + +func TestHTTPPlugin_Observe_EmptyEvents(t *testing.T) { + p := NewHTTPPlugin("test", "http://localhost:9999") + err := p.Observe(context.Background(), nil) + if err != nil { + t.Errorf("empty events should return nil, got %v", err) + } +} + +func TestHTTPPlugin_Observe_ServerError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error on non-200 response") + } +} + +func TestHTTPPlugin_Observe_NotOk(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": false}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error when ok is false") + } +} + +func TestHTTPPlugin_Observe_InvalidJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`not json`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL) + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err == nil { + t.Error("Observe should return error on invalid JSON response") + } +} + +func TestHTTPPlugin_CustomHeader(t *testing.T) { + var receivedHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHeader = r.Header.Get("X-Custom") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok": true}`)) + })) + defer srv.Close() + + p := NewHTTPPlugin("test", srv.URL, plugin.HeaderOption(http.Header{ + "X-Custom": []string{"my-value"}, + })) + err := p.Observe(context.Background(), []observer.Event{ + service.ServiceEvent{Kind: "service", Service: "test"}, + }) + if err != nil { + t.Fatal(err) + } + if receivedHeader != "my-value" { + t.Errorf("X-Custom header = %s, want my-value", receivedHeader) + } +} + +func TestHTTPPlugin_Observe_WithClient(t *testing.T) { + // Ensure p.client is properly set by constructor and Observe doesn't + // hit the dead nil check. + p := NewHTTPPlugin("test", "http://example.com") + if p.(*httpPlugin).client == nil { + t.Fatal("client should not be nil after construction") + } +} + +// Test that httpPlugin satisfies observer.Observer +var _ observer.Observer = (*httpPlugin)(nil) diff --git a/observer/stats/stats_test.go b/observer/stats/stats_test.go new file mode 100644 index 00000000..7b6babe4 --- /dev/null +++ b/observer/stats/stats_test.go @@ -0,0 +1,200 @@ +package stats + +import ( + "testing" + + "github.com/go-gost/core/observer" + "github.com/go-gost/core/observer/stats" +) + +func TestStats_Add(t *testing.T) { + s := NewStats(false) + st := s.(*Stats) + + s.Add(stats.KindTotalConns, 1) + if v := st.totalConns.Load(); v != 1 { + t.Errorf("totalConns = %d, want 1", v) + } + + s.Add(stats.KindCurrentConns, 5) + if v := st.currentConns.Load(); v != 5 { + t.Errorf("currentConns = %d, want 5", v) + } + s.Add(stats.KindCurrentConns, -2) + if v := st.currentConns.Load(); v != 3 { + t.Errorf("currentConns after -2 = %d, want 3", v) + } + + s.Add(stats.KindInputBytes, 100) + if v := st.inputBytes.Load(); v != 100 { + t.Errorf("inputBytes = %d, want 100", v) + } + + s.Add(stats.KindOutputBytes, 200) + if v := st.outputBytes.Load(); v != 200 { + t.Errorf("outputBytes = %d, want 200", v) + } + + s.Add(stats.KindTotalErrs, 1) + if v := st.totalErrs.Load(); v != 1 { + t.Errorf("totalErrs = %d, want 1", v) + } +} + +func TestStats_Add_NegativeIgnored(t *testing.T) { + s := NewStats(false) + st := s.(*Stats) + + s.Add(stats.KindTotalConns, -1) // ignored (n <= 0 check) + if v := st.totalConns.Load(); v != 0 { + t.Errorf("totalConns after -1 = %d, want 0", v) + } + + s.Add(stats.KindTotalErrs, -1) // ignored (n <= 0 check) + if v := st.totalErrs.Load(); v != 0 { + t.Errorf("totalErrs after -1 = %d, want 0", v) + } +} + +func TestStats_Add_Nil(t *testing.T) { + var s *Stats + s.Add(stats.KindTotalConns, 1) // must not panic +} + +func TestStats_Get(t *testing.T) { + s := NewStats(false) + st := s.(*Stats) + + st.totalConns.Store(10) + st.currentConns.Store(3) + st.inputBytes.Store(1000) + st.outputBytes.Store(2000) + st.totalErrs.Store(2) + + if v := s.Get(stats.KindTotalConns); v != 10 { + t.Errorf("totalConns = %d, want 10", v) + } + if v := s.Get(stats.KindCurrentConns); v != 3 { + t.Errorf("currentConns = %d, want 3", v) + } + if v := s.Get(stats.KindInputBytes); v != 1000 { + t.Errorf("inputBytes = %d, want 1000", v) + } + if v := s.Get(stats.KindOutputBytes); v != 2000 { + t.Errorf("outputBytes = %d, want 2000", v) + } + if v := s.Get(stats.KindTotalErrs); v != 2 { + t.Errorf("totalErrs = %d, want 2", v) + } +} + +func TestStats_Get_Nil(t *testing.T) { + var s *Stats + if v := s.Get(stats.KindTotalConns); v != 0 { + t.Errorf("nil Get should return 0, got %d", v) + } +} + +func TestStats_Get_ResetTraffic(t *testing.T) { + s := NewStats(true) + st := s.(*Stats) + + st.inputBytes.Store(500) + st.outputBytes.Store(300) + + // First Get should return the value and swap to 0 + if v := s.Get(stats.KindInputBytes); v != 500 { + t.Errorf("first Get inputBytes = %d, want 500", v) + } + if v := st.inputBytes.Load(); v != 0 { + t.Errorf("inputBytes after Get = %d, want 0", v) + } + + // Second Get should return 0 + if v := s.Get(stats.KindInputBytes); v != 0 { + t.Errorf("second Get inputBytes = %d, want 0", v) + } + + if v := s.Get(stats.KindOutputBytes); v != 300 { + t.Errorf("first Get outputBytes = %d, want 300", v) + } + if v := st.outputBytes.Load(); v != 0 { + t.Errorf("outputBytes after Get = %d, want 0", v) + } + + // Non-traffic counters are unaffected by resetTraffic + st.totalConns.Store(42) + if v := s.Get(stats.KindTotalConns); v != 42 { + t.Errorf("totalConns = %d, want 42", v) + } +} + +func TestStats_Reset(t *testing.T) { + s := NewStats(false) + st := s.(*Stats) + + st.totalConns.Store(10) + st.currentConns.Store(5) + st.inputBytes.Store(100) + st.outputBytes.Store(200) + st.totalErrs.Store(3) + st.updated.Store(true) + + s.Reset() + + if v := st.totalConns.Load(); v != 0 { + t.Errorf("totalConns after Reset = %d, want 0", v) + } + if v := st.currentConns.Load(); v != 0 { + t.Errorf("currentConns after Reset = %d, want 0", v) + } + if v := st.inputBytes.Load(); v != 0 { + t.Errorf("inputBytes after Reset = %d, want 0", v) + } + if v := st.outputBytes.Load(); v != 0 { + t.Errorf("outputBytes after Reset = %d, want 0", v) + } + if v := st.totalErrs.Load(); v != 0 { + t.Errorf("totalErrs after Reset = %d, want 0", v) + } + if st.updated.Load() { + t.Error("updated after Reset should be false") + } +} + +func TestStats_IsUpdated(t *testing.T) { + s := NewStats(false) + st := s.(*Stats) + + // Initially false + if s.IsUpdated() { + t.Error("IsUpdated should be false initially") + } + + st.updated.Store(true) + + // Should return true then swap to false + if !s.IsUpdated() { + t.Error("IsUpdated should return true after update") + } + if s.IsUpdated() { + t.Error("IsUpdated should return false on second call") + } +} + +func TestStats_UnknownKind(t *testing.T) { + s := NewStats(false) + + if v := s.Get(stats.Kind(999)); v != 0 { + t.Errorf("unknown kind should return 0, got %d", v) + } + // Add with unknown kind should not panic + s.Add(stats.Kind(999), 42) +} + +func TestStatsEvent_Type(t *testing.T) { + ev := StatsEvent{} + if ev.Type() != observer.EventStats { + t.Errorf("Type = %q, want %q", ev.Type(), observer.EventStats) + } +} diff --git a/observer/stats/wrapper/conn.go b/observer/stats/wrapper/conn.go index 90413916..9d3bf31d 100644 --- a/observer/stats/wrapper/conn.go +++ b/observer/stats/wrapper/conn.go @@ -181,9 +181,7 @@ func (c *udpConn) SetWriteBuffer(n int) error { func (c *udpConn) Read(b []byte) (n int, err error) { if nc, ok := c.PacketConn.(io.Reader); ok { n, err = nc.Read(b) - if c.stats != nil { - c.stats.Add(stats.KindInputBytes, int64(n)) - } + c.stats.Add(stats.KindInputBytes, int64(n)) return } err = errUnsupport @@ -192,18 +190,14 @@ func (c *udpConn) Read(b []byte) (n int, err error) { func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { n, addr, err = c.PacketConn.ReadFrom(p) - if c.stats != nil { - c.stats.Add(stats.KindInputBytes, int64(n)) - } + c.stats.Add(stats.KindInputBytes, int64(n)) return } func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { if nc, ok := c.PacketConn.(udp.ReadUDP); ok { n, addr, err = nc.ReadFromUDP(b) - if c.stats != nil { - c.stats.Add(stats.KindInputBytes, int64(n)) - } + c.stats.Add(stats.KindInputBytes, int64(n)) return } err = errUnsupport @@ -213,9 +207,7 @@ func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) { if nc, ok := c.PacketConn.(udp.ReadUDP); ok { n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob) - if c.stats != nil { - c.stats.Add(stats.KindInputBytes, int64(n)) - } + c.stats.Add(stats.KindInputBytes, int64(n)) return } err = errUnsupport @@ -225,9 +217,7 @@ func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAd func (c *udpConn) Write(b []byte) (n int, err error) { if nc, ok := c.PacketConn.(io.Writer); ok { n, err = nc.Write(b) - if c.stats != nil { - c.stats.Add(stats.KindOutputBytes, int64(n)) - } + c.stats.Add(stats.KindOutputBytes, int64(n)) return } err = errUnsupport @@ -236,18 +226,14 @@ func (c *udpConn) Write(b []byte) (n int, err error) { func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { n, err = c.PacketConn.WriteTo(p, addr) - if c.stats != nil { - c.stats.Add(stats.KindOutputBytes, int64(n)) - } + c.stats.Add(stats.KindOutputBytes, int64(n)) return } func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { if nc, ok := c.PacketConn.(udp.WriteUDP); ok { n, err = nc.WriteToUDP(b, addr) - if c.stats != nil { - c.stats.Add(stats.KindOutputBytes, int64(n)) - } + c.stats.Add(stats.KindOutputBytes, int64(n)) return } err = errUnsupport @@ -257,9 +243,7 @@ func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) { if nc, ok := c.PacketConn.(udp.WriteUDP); ok { n, oobn, err = nc.WriteMsgUDP(b, oob, addr) - if c.stats != nil { - c.stats.Add(stats.KindOutputBytes, int64(n)) - } + c.stats.Add(stats.KindOutputBytes, int64(n)) return } err = errUnsupport diff --git a/observer/stats/wrapper/conn_test.go b/observer/stats/wrapper/conn_test.go new file mode 100644 index 00000000..66acb63d --- /dev/null +++ b/observer/stats/wrapper/conn_test.go @@ -0,0 +1,441 @@ +package wrapper + +import ( + "errors" + "net" + "syscall" + "testing" + "time" + + "github.com/go-gost/core/observer/stats" + "github.com/go-gost/x/internal/net/udp" + ostats "github.com/go-gost/x/observer/stats" +) + +// --- mock types --- + +type fakeConn struct { + net.Conn + readBuf []byte + readPos int + writeBuf []byte + closed bool + syscallFn func() (syscall.RawConn, error) +} + +func (c *fakeConn) Read(b []byte) (int, error) { + if c.readPos >= len(c.readBuf) { + return 0, errors.New("EOF") + } + n := copy(b, c.readBuf[c.readPos:]) + c.readPos += n + return n, nil +} + +func (c *fakeConn) Write(b []byte) (int, error) { + c.writeBuf = append(c.writeBuf, b...) + return len(b), nil +} + +func (c *fakeConn) Close() error { + if c.closed { + return errors.New("already closed") + } + c.closed = true + return nil +} + +type fakePacketConn struct { + readBuf []byte + readPos int + writeBuf []byte +} + +func (c *fakePacketConn) ReadFrom(p []byte) (int, net.Addr, error) { + if c.readPos >= len(c.readBuf) { + return 0, nil, errors.New("EOF") + } + n := copy(p, c.readBuf[c.readPos:]) + c.readPos += n + return n, &net.UDPAddr{IP: net.IPv4(1, 1, 1, 1), Port: 1234}, nil +} + +func (c *fakePacketConn) WriteTo(p []byte, addr net.Addr) (int, error) { + c.writeBuf = append(c.writeBuf, p...) + return len(p), nil +} + +func (c *fakePacketConn) Close() error { return nil } + +func (c *fakePacketConn) LocalAddr() net.Addr { + return &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 9999} +} + +func (c *fakePacketConn) SetDeadline(t time.Time) error { return nil } +func (c *fakePacketConn) SetReadDeadline(t time.Time) error { return nil } +func (c *fakePacketConn) SetWriteDeadline(t time.Time) error { return nil } + +// --- WrapConn tests --- + +func TestWrapConn_Nil(t *testing.T) { + // nil conn + if result := WrapConn(nil, ostats.NewStats(false)); result != nil { + t.Error("WrapConn with nil conn should return nil") + } + // nil stats + fc := &fakeConn{} + if result := WrapConn(fc, nil); result != fc { + t.Error("WrapConn with nil stats should return original conn") + } + // both nil + if result := WrapConn(nil, nil); result != nil { + t.Error("WrapConn with both nil should return nil") + } +} + +func TestWrapConn_IncrementsCounters(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + WrapConn(fc, st) + + if v := st.Get(stats.KindTotalConns); v != 1 { + t.Errorf("totalConns = %d, want 1", v) + } + if v := st.Get(stats.KindCurrentConns); v != 1 { + t.Errorf("currentConns = %d, want 1", v) + } +} + +func TestWrapConn_ReadWrite(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{readBuf: []byte("hello")} + wrapped := WrapConn(fc, st) + + buf := make([]byte, 10) + n, _ := wrapped.Read(buf) + if n != 5 { + t.Fatalf("read n = %d, want 5", n) + } + if v := st.Get(stats.KindInputBytes); v != 5 { + t.Errorf("inputBytes = %d, want 5", v) + } + + n, _ = wrapped.Write([]byte("world")) + if n != 5 { + t.Fatalf("write n = %d, want 5", n) + } + if v := st.Get(stats.KindOutputBytes); v != 5 { + t.Errorf("outputBytes = %d, want 5", v) + } +} + +func TestWrapConn_Close(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + // currentConns should be 1 after wrap + if v := st.Get(stats.KindCurrentConns); v != 1 { + t.Fatalf("currentConns = %d, want 1", v) + } + + wrapped.Close() + + if v := st.Get(stats.KindCurrentConns); v != 0 { + t.Errorf("currentConns after close = %d, want 0", v) + } + if !fc.closed { + t.Error("underlying conn should be closed") + } +} + +func TestWrapConn_DoubleClose(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + wrapped.Close() + wrapped.Close() // must not panic and must not double-decrement + + if v := st.Get(stats.KindCurrentConns); v != 0 { + t.Errorf("currentConns = %d after double close, want 0", v) + } +} + +func TestWrapConn_CloseRace(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + done := make(chan struct{}) + go func() { + wrapped.Close() + done <- struct{}{} + }() + go func() { + wrapped.Close() + done <- struct{}{} + }() + <-done + <-done + + // currentConns should only be decremented once + if v := st.Get(stats.KindCurrentConns); v != 0 { + t.Errorf("currentConns after concurrent close = %d, want 0", v) + } +} + +func TestWrapConn_CloseRead(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + err := wrapped.(*conn).CloseRead() + if err == nil { + t.Error("CloseRead on basic conn should return error") + } +} + +func TestWrapConn_CloseWrite(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + err := wrapped.(*conn).CloseWrite() + if err == nil { + t.Error("CloseWrite on basic conn should return error") + } +} + +func TestWrapConn_SyscallConn(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + _, err := wrapped.(*conn).SyscallConn() + if err == nil { + t.Error("SyscallConn on non-syscall conn should return error") + } +} + +func TestWrapConn_Context(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + wrapped := WrapConn(fc, st) + + if wrapped.(*conn).Context() != nil { + t.Error("Context on basic conn should return nil") + } +} + +// --- WrapPacketConn tests --- + +func TestWrapPacketConn_Nil(t *testing.T) { + // nil pc + if result := WrapPacketConn(nil, ostats.NewStats(false)); result != nil { + t.Error("WrapPacketConn with nil pc should return nil") + } + // nil stats + fpc := &fakePacketConn{} + if result := WrapPacketConn(fpc, nil); result != fpc { + t.Error("WrapPacketConn with nil stats should return original conn") + } + // both nil + if result := WrapPacketConn(nil, nil); result != nil { + t.Error("WrapPacketConn with both nil should return nil") + } +} + +func TestWrapPacketConn_ReadFrom(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{readBuf: []byte("packet-data")} + wrapped := WrapPacketConn(fpc, st) + + buf := make([]byte, 32) + n, addr, err := wrapped.ReadFrom(buf) + if err != nil { + t.Fatal(err) + } + if n != 11 { + t.Fatalf("n = %d, want 11", n) + } + if addr == nil { + t.Fatal("addr should not be nil") + } + if v := st.Get(stats.KindInputBytes); v != 11 { + t.Errorf("inputBytes = %d, want 11", v) + } +} + +func TestWrapPacketConn_WriteTo(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapPacketConn(fpc, st) + + addr := &net.UDPAddr{IP: net.IPv4(2, 2, 2, 2), Port: 5678} + n, err := wrapped.WriteTo([]byte("hello"), addr) + if err != nil { + t.Fatal(err) + } + if n != 5 { + t.Fatalf("n = %d, want 5", n) + } + if v := st.Get(stats.KindOutputBytes); v != 5 { + t.Errorf("outputBytes = %d, want 5", v) + } +} + +func TestWrapPacketConn_Context(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapPacketConn(fpc, st) + + if wrapped.(*packetConn).Context() != nil { + t.Error("Context should return nil for basic packet conn") + } +} + +// --- WrapUDPConn tests --- + +func TestWrapUDPConn_Nil(t *testing.T) { + if result := WrapUDPConn(nil, ostats.NewStats(false)); result != nil { + t.Error("WrapUDPConn with nil pc should return nil") + } + if result := WrapUDPConn(&fakePacketConn{}, nil); result != nil { + t.Error("WrapUDPConn with nil stats should return nil") + } +} + +func TestWrapUDPConn_ReadFrom(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{readBuf: []byte("udp-packet")} + wrapped := WrapUDPConn(fpc, st) + + buf := make([]byte, 32) + n, _, err := wrapped.ReadFrom(buf) + if err != nil { + t.Fatal(err) + } + if n != 10 { + t.Fatalf("n = %d, want 10", n) + } + if v := st.Get(stats.KindInputBytes); v != 10 { + t.Errorf("inputBytes = %d, want 10", v) + } +} + +func TestWrapUDPConn_WriteTo(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + addr := &net.UDPAddr{IP: net.IPv4(3, 3, 3, 3), Port: 4444} + n, err := wrapped.WriteTo([]byte("data"), addr) + if err != nil { + t.Fatal(err) + } + if n != 4 { + t.Fatalf("n = %d, want 4", n) + } + if v := st.Get(stats.KindOutputBytes); v != 4 { + t.Errorf("outputBytes = %d, want 4", v) + } +} + +func TestWrapUDPConn_RemoteAddr(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + // fakePacketConn doesn't implement xnet.RemoteAddr, so this returns nil + if wrapped.RemoteAddr() != nil { + t.Error("RemoteAddr should return nil for basic packet conn") + } +} + +func TestWrapUDPConn_SetReadBuffer(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + err := wrapped.SetReadBuffer(4096) + if err == nil { + t.Error("SetReadBuffer should return error for basic packet conn") + } +} + +func TestWrapUDPConn_SetWriteBuffer(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + err := wrapped.SetWriteBuffer(4096) + if err == nil { + t.Error("SetWriteBuffer should return error for basic packet conn") + } +} + +func TestWrapUDPConn_SetDSCP(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + // fakePacketConn doesn't implement xnet.SetDSCP, so this returns nil + uc := wrapped.(*udpConn) + if err := uc.SetDSCP(46); err != nil { + t.Errorf("SetDSCP should return nil, got %v", err) + } +} + +func TestWrapUDPConn_SyscallConn(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + _, err := wrapped.SyscallConn() + if err == nil { + t.Error("SyscallConn should return error for basic packet conn") + } +} + +func TestWrapUDPConn_Context(t *testing.T) { + st := ostats.NewStats(false) + fpc := &fakePacketConn{} + wrapped := WrapUDPConn(fpc, st) + + uc := wrapped.(*udpConn) + if uc.Context() != nil { + t.Error("Context should return nil for basic packet conn") + } +} + +// Verify the returned types are correct +func TestWrapConn_ReturnType(t *testing.T) { + fc := &fakeConn{} + st := ostats.NewStats(false) + result := WrapConn(fc, st) + if _, ok := result.(*conn); !ok { + t.Errorf("WrapConn returned %T, want *conn", result) + } +} + +func TestWrapPacketConn_ReturnType(t *testing.T) { + fpc := &fakePacketConn{} + st := ostats.NewStats(false) + result := WrapPacketConn(fpc, st) + if _, ok := result.(*packetConn); !ok { + t.Errorf("WrapPacketConn returned %T, want *packetConn", result) + } +} + +func TestWrapUDPConn_ReturnType(t *testing.T) { + fpc := &fakePacketConn{} + st := ostats.NewStats(false) + result := WrapUDPConn(fpc, st) + if _, ok := result.(*udpConn); !ok { + t.Errorf("WrapUDPConn returned %T, want *udpConn", result) + } +} + +// Ensure the wrapper satisfies udp.Conn +var _ udp.Conn = &udpConn{} diff --git a/observer/stats/wrapper/io_test.go b/observer/stats/wrapper/io_test.go new file mode 100644 index 00000000..387df91d --- /dev/null +++ b/observer/stats/wrapper/io_test.go @@ -0,0 +1,90 @@ +package wrapper + +import ( + "testing" + + "github.com/go-gost/core/observer/stats" + ostats "github.com/go-gost/x/observer/stats" +) + +type fakeReadWriter struct { + readBuf []byte + readPos int + writeBuf []byte +} + +func (rw *fakeReadWriter) Read(b []byte) (int, error) { + if rw.readPos >= len(rw.readBuf) { + return 0, nil + } + n := copy(b, rw.readBuf[rw.readPos:]) + rw.readPos += n + return n, nil +} + +func (rw *fakeReadWriter) Write(b []byte) (int, error) { + rw.writeBuf = append(rw.writeBuf, b...) + return len(b), nil +} + +func TestWrapReadWriter_Nil(t *testing.T) { + // nil rw + if result := WrapReadWriter(nil, ostats.NewStats(false)); result != nil { + t.Error("WrapReadWriter with nil rw should return nil") + } + // nil stats + frw := &fakeReadWriter{} + if result := WrapReadWriter(frw, nil); result != frw { + t.Error("WrapReadWriter with nil stats should return original rw") + } + // both nil + if result := WrapReadWriter(nil, nil); result != nil { + t.Error("WrapReadWriter with both nil should return nil") + } +} + +func TestWrapReadWriter_Read(t *testing.T) { + st := ostats.NewStats(false) + frw := &fakeReadWriter{readBuf: []byte("testdata")} + wrapped := WrapReadWriter(frw, st) + + buf := make([]byte, 8) + n, _ := wrapped.Read(buf) + if n != 8 { + t.Fatalf("read n = %d, want 8", n) + } + if st.Get(stats.KindInputBytes) != 8 { + t.Errorf("inputBytes = %d, want 8", st.Get(stats.KindInputBytes)) + } +} + +func TestWrapReadWriter_Write(t *testing.T) { + st := ostats.NewStats(false) + frw := &fakeReadWriter{} + wrapped := WrapReadWriter(frw, st) + + n, _ := wrapped.Write([]byte("hello")) + if n != 5 { + t.Fatalf("write n = %d, want 5", n) + } + if st.Get(stats.KindOutputBytes) != 5 { + t.Errorf("outputBytes = %d, want 5", st.Get(stats.KindOutputBytes)) + } +} + +func TestWrapReadWriter_ReadWrite(t *testing.T) { + st := ostats.NewStats(false) + frw := &fakeReadWriter{readBuf: []byte("abcdefghij")} + wrapped := WrapReadWriter(frw, st) + + buf := make([]byte, 5) + wrapped.Read(buf) // reads 5 bytes + wrapped.Write([]byte("xyz")) // writes 3 bytes + + if st.Get(stats.KindInputBytes) != 5 { + t.Errorf("inputBytes = %d, want 5", st.Get(stats.KindInputBytes)) + } + if st.Get(stats.KindOutputBytes) != 3 { + t.Errorf("outputBytes = %d, want 3", st.Get(stats.KindOutputBytes)) + } +} diff --git a/observer/stats/wrapper/listener_test.go b/observer/stats/wrapper/listener_test.go new file mode 100644 index 00000000..db16b4e1 --- /dev/null +++ b/observer/stats/wrapper/listener_test.go @@ -0,0 +1,81 @@ +package wrapper + +import ( + "net" + "testing" + + "github.com/go-gost/core/observer/stats" + ostats "github.com/go-gost/x/observer/stats" +) + +type fakeListener struct { + net.Listener + acceptConn net.Conn + acceptErr error +} + +func (l *fakeListener) Accept() (net.Conn, error) { + if l.acceptErr != nil { + return nil, l.acceptErr + } + return l.acceptConn, nil +} + +func (l *fakeListener) Close() error { return nil } + +func (l *fakeListener) Addr() net.Addr { + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080} +} + +func TestWrapListener_Nil(t *testing.T) { + // nil ln + if result := WrapListener(nil, ostats.NewStats(false)); result != nil { + t.Error("WrapListener with nil ln should return nil") + } + // nil stats + fl := &fakeListener{} + if result := WrapListener(fl, nil); result != fl { + t.Error("WrapListener with nil stats should return original listener") + } + // both nil + if result := WrapListener(nil, nil); result != nil { + t.Error("WrapListener with both nil should return nil") + } +} + +func TestWrapListener_Accept_WrapsConn(t *testing.T) { + st := ostats.NewStats(false) + fc := &fakeConn{} + fl := &fakeListener{acceptConn: fc} + wrapped := WrapListener(fl, st) + + conn, err := wrapped.Accept() + if err != nil { + t.Fatal(err) + } + if conn == fc { + t.Error("accepted conn should be wrapped, not the raw conn") + } + // The wrapper should track connection count + if st.Get(stats.KindTotalConns) != 1 { + t.Errorf("totalConns = %d, want 1 after accept", st.Get(stats.KindTotalConns)) + } + if st.Get(stats.KindCurrentConns) != 1 { + t.Errorf("currentConns = %d, want 1 after accept", st.Get(stats.KindCurrentConns)) + } +} + +func TestWrapListener_Accept_Error(t *testing.T) { + st := ostats.NewStats(false) + fl := &fakeListener{acceptErr: net.ErrClosed} + wrapped := WrapListener(fl, st) + + _, err := wrapped.Accept() + if err == nil { + t.Fatal("expected error on accept") + } + // connection counters should not be incremented on error + if st.Get(stats.KindTotalConns) != 0 { + t.Errorf("totalConns = %d, want 0 after error", st.Get(stats.KindTotalConns)) + } +}