diff --git a/metrics/metrics.go b/metrics/metrics.go index 3c74dfe4..1f3b802d 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -34,14 +34,19 @@ var ( enabled atomic.Bool ) +// Enable enables or disables metrics collection globally. When disabled, all +// GetCounter, GetGauge, and GetObserver calls return noop implementations. func Enable(b bool) { enabled.Store(b) } +// IsEnabled reports whether metrics collection is enabled. func IsEnabled() bool { return enabled.Load() } +// GetCounter returns a Counter for the given name and labels. When metrics are +// disabled, a noop implementation is returned. func GetCounter(name metrics.MetricName, labels metrics.Labels) metrics.Counter { if IsEnabled() { return defaultMetrics.Counter(name, labels) @@ -49,6 +54,8 @@ func GetCounter(name metrics.MetricName, labels metrics.Labels) metrics.Counter return noop.Counter(name, labels) } +// GetGauge returns a Gauge for the given name and labels. When metrics are +// disabled, a noop implementation is returned. func GetGauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge { if IsEnabled() { return defaultMetrics.Gauge(name, labels) @@ -56,6 +63,8 @@ func GetGauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge { return noop.Gauge(name, labels) } +// GetObserver returns an Observer for the given name and labels. When metrics are +// disabled, a noop implementation is returned. func GetObserver(name metrics.MetricName, labels metrics.Labels) metrics.Observer { if IsEnabled() { return defaultMetrics.Observer(name, labels) diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go new file mode 100644 index 00000000..9c626201 --- /dev/null +++ b/metrics/metrics_test.go @@ -0,0 +1,173 @@ +package metrics + +import ( + "testing" + + "github.com/go-gost/core/metrics" +) + +func TestEnableDisable(t *testing.T) { + // Save and restore original state. + orig := IsEnabled() + defer Enable(orig) + + Enable(false) + if IsEnabled() { + t.Error("IsEnabled should return false after Enable(false)") + } + + Enable(true) + if !IsEnabled() { + t.Error("IsEnabled should return true after Enable(true)") + } + + Enable(false) + if IsEnabled() { + t.Error("IsEnabled should return false after second Enable(false)") + } +} + +func TestGetCounterDisabled(t *testing.T) { + orig := IsEnabled() + Enable(false) + defer Enable(orig) + + c := GetCounter(MetricServiceRequestsCounter, metrics.Labels{"service": "s", "client": "c"}) + if c == nil { + t.Error("GetCounter should not return nil when disabled") + } + // Should be noop — calling methods must not panic. + c.Inc() + c.Add(42) + c.Add(-1) +} + +func TestGetGaugeDisabled(t *testing.T) { + orig := IsEnabled() + Enable(false) + defer Enable(orig) + + // MetricServicesGauge has only ["host"] label (auto-added), so pass no extra labels. + g := GetGauge(MetricServicesGauge, nil) + if g == nil { + t.Error("GetGauge should not return nil when disabled") + } + g.Inc() + g.Dec() + g.Add(1.5) + g.Set(100) + g.Set(-1) +} + +func TestGetObserverDisabled(t *testing.T) { + orig := IsEnabled() + Enable(false) + defer Enable(orig) + + o := GetObserver(MetricServiceRequestsDurationObserver, metrics.Labels{"service": "s"}) + if o == nil { + t.Error("GetObserver should not return nil when disabled") + } + o.Observe(0.5) + o.Observe(0) + o.Observe(-1) +} + +func TestGetCounterEnabled(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + c := GetCounter(MetricServiceRequestsCounter, metrics.Labels{"service": "s", "client": "c"}) + if c == nil { + t.Error("GetCounter should not return nil when enabled") + } + // Methods should not panic. + c.Inc() + c.Add(10) +} + +func TestGetGaugeEnabled(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + // MetricServicesGauge has only ["host"] label (auto-added). + g := GetGauge(MetricServicesGauge, nil) + if g == nil { + t.Error("GetGauge should not return nil when enabled") + } + g.Inc() + g.Dec() + g.Set(5) +} + +func TestGetObserverEnabled(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + o := GetObserver(MetricServiceRequestsDurationObserver, metrics.Labels{"service": "s"}) + if o == nil { + t.Error("GetObserver should not return nil when enabled") + } + o.Observe(0.25) +} + +func TestGetCounterEnabledUnknownName(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + c := GetCounter("nonexistent_metric", nil) + if c == nil { + t.Error("GetCounter should not return nil for unknown metric name") + } + c.Inc() + c.Add(5) +} + +func TestGetGaugeEnabledUnknownName(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + g := GetGauge("nonexistent_metric", nil) + if g == nil { + t.Error("GetGauge should not return nil for unknown metric name") + } + g.Inc() + g.Set(1) +} + +func TestGetObserverEnabledUnknownName(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + o := GetObserver("nonexistent_metric", nil) + if o == nil { + t.Error("GetObserver should not return nil for unknown metric name") + } + o.Observe(1.0) +} + +func TestMetricConstantsNotEmpty(t *testing.T) { + names := []metrics.MetricName{ + MetricServicesGauge, + MetricServiceRequestsCounter, + MetricServiceRequestsInFlightGauge, + MetricServiceRequestsDurationObserver, + MetricServiceTransferInputBytesCounter, + MetricServiceTransferOutputBytesCounter, + MetricNodeConnectDurationObserver, + MetricServiceHandlerErrorsCounter, + MetricChainErrorsCounter, + MetricRecorderRecordsCounter, + } + for i, n := range names { + if n == "" { + t.Errorf("metric constant at index %d is empty", i) + } + } +} diff --git a/metrics/noop.go b/metrics/noop.go index 4ea35d1f..7e750f26 100644 --- a/metrics/noop.go +++ b/metrics/noop.go @@ -12,6 +12,8 @@ var ( type noopMetrics struct{} +// Noop returns a Metrics implementation whose counters, gauges, and observers +// are all no-ops. func Noop() metrics.Metrics { return noop } diff --git a/metrics/noop_test.go b/metrics/noop_test.go new file mode 100644 index 00000000..e12c4e99 --- /dev/null +++ b/metrics/noop_test.go @@ -0,0 +1,80 @@ +package metrics + +import ( + "testing" + + "github.com/go-gost/core/metrics" +) + +func TestNoopReturnsSingleton(t *testing.T) { + n1 := Noop() + n2 := Noop() + if n1 != n2 { + t.Error("Noop() should return the same instance every time") + } +} + +func TestNoopMetricsCounter(t *testing.T) { + c := Noop().Counter("any", nil) + if c == nil { + t.Fatal("noop Counter should not be nil") + } + c.Inc() + c.Add(42) +} + +func TestNoopMetricsGauge(t *testing.T) { + g := Noop().Gauge("any", nil) + if g == nil { + t.Fatal("noop Gauge should not be nil") + } + g.Inc() + g.Dec() + g.Add(1.5) + g.Set(100) +} + +func TestNoopMetricsObserver(t *testing.T) { + o := Noop().Observer("any", nil) + if o == nil { + t.Fatal("noop Observer should not be nil") + } + o.Observe(0.5) +} + +func TestNoopMetricsSameInstances(t *testing.T) { + m := Noop() + // Each call returns the shared singleton noop instances + if m.Counter("a", nil) != m.Counter("b", nil) { + t.Error("noop Counter should return same instance regardless of name") + } + if m.Gauge("a", nil) != m.Gauge("b", nil) { + t.Error("noop Gauge should return same instance regardless of name") + } + if m.Observer("a", nil) != m.Observer("b", nil) { + t.Error("noop Observer should return same instance regardless of name") + } +} + +func TestNoopPackageVarsAreNoopTypes(t *testing.T) { + // The package-level noop variables should be instances of the concrete noop types. + // These are the values returned by GetCounter/Gauge/Observer when disabled. + c := Noop().Counter("x", metrics.Labels{"k": "v"}) + if c == nil { + t.Fatal("noop Counter from Noop() should not be nil") + } + c.Inc() + c.Add(10) + + g := Noop().Gauge("x", metrics.Labels{"k": "v"}) + if g == nil { + t.Fatal("noop Gauge from Noop() should not be nil") + } + g.Set(5) + + o := Noop().Observer("x", metrics.Labels{"k": "v"}) + if o == nil { + t.Fatal("noop Observer from Noop() should not be nil") + } + o.Observe(0.1) +} diff --git a/metrics/prom.go b/metrics/prom.go index 797df6ab..44e5d149 100644 --- a/metrics/prom.go +++ b/metrics/prom.go @@ -1,6 +1,7 @@ package metrics import ( + "maps" "os" "github.com/go-gost/core/metrics" @@ -14,6 +15,8 @@ type promMetrics struct { histograms map[metrics.MetricName]*prometheus.HistogramVec } +// NewMetrics returns a Prometheus-based Metrics implementation. All metrics are +// automatically registered with the default Prometheus registry. func NewMetrics() metrics.Metrics { host, _ := os.Hostname() m := &promMetrics{ @@ -107,35 +110,32 @@ func NewMetrics() metrics.Metrics { func (m *promMetrics) Gauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge { v, ok := m.gauges[name] if !ok { - return nil + return nopGauge } - if labels == nil { - labels = metrics.Labels{} - } - labels["host"] = m.host - return v.With(prometheus.Labels(labels)) + plabels := prometheus.Labels{} + maps.Copy(plabels, labels) + plabels["host"] = m.host + return v.With(plabels) } func (m *promMetrics) Counter(name metrics.MetricName, labels metrics.Labels) metrics.Counter { v, ok := m.counters[name] if !ok { - return nil + return nopCounter } - if labels == nil { - labels = metrics.Labels{} - } - labels["host"] = m.host - return v.With(prometheus.Labels(labels)) + plabels := prometheus.Labels{} + maps.Copy(plabels, labels) + plabels["host"] = m.host + return v.With(plabels) } func (m *promMetrics) Observer(name metrics.MetricName, labels metrics.Labels) metrics.Observer { v, ok := m.histograms[name] if !ok { - return nil + return nopObserver } - if labels == nil { - labels = metrics.Labels{} - } - labels["host"] = m.host - return v.With(prometheus.Labels(labels)) + plabels := prometheus.Labels{} + maps.Copy(plabels, labels) + plabels["host"] = m.host + return v.With(plabels) } diff --git a/metrics/prom_test.go b/metrics/prom_test.go new file mode 100644 index 00000000..446cfe92 --- /dev/null +++ b/metrics/prom_test.go @@ -0,0 +1,114 @@ +package metrics + +import ( + "testing" + + "github.com/go-gost/core/metrics" +) + +func TestPromMetricsGaugeHostLabel(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + // MetricServicesGauge: labels=["host"], host auto-added. + g := GetGauge(MetricServicesGauge, nil) + if g == nil { + t.Fatal("gauge should not be nil") + } + g.Set(3) + g.Add(1) + g.Inc() + g.Dec() +} + +func TestPromMetricsCounterKnownNames(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + tests := []struct { + name metrics.MetricName + labels metrics.Labels + }{ + {MetricServiceRequestsCounter, metrics.Labels{"service": "s", "client": "c"}}, + {MetricServiceTransferInputBytesCounter, metrics.Labels{"service": "s", "client": "c"}}, + {MetricServiceTransferOutputBytesCounter, metrics.Labels{"service": "s", "client": "c"}}, + {MetricServiceHandlerErrorsCounter, metrics.Labels{"service": "s", "client": "c"}}, + {MetricChainErrorsCounter, metrics.Labels{"chain": "ch", "node": "n"}}, + {MetricRecorderRecordsCounter, metrics.Labels{"recorder": "r"}}, + } + + for _, tt := range tests { + c := GetCounter(tt.name, tt.labels) + if c == nil { + t.Errorf("GetCounter(%q) returned nil", tt.name) + continue + } + c.Inc() + c.Add(10) + } +} + +func TestPromMetricsGaugeKnownNames(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + g := GetGauge(MetricServiceRequestsInFlightGauge, metrics.Labels{"service": "s", "client": "c"}) + if g == nil { + t.Fatal("gauge should not be nil") + } + g.Set(1) + g.Inc() + g.Dec() +} + +func TestPromMetricsObserverKnownNames(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + tests := []struct { + name metrics.MetricName + labels metrics.Labels + }{ + {MetricServiceRequestsDurationObserver, metrics.Labels{"service": "s"}}, + {MetricNodeConnectDurationObserver, metrics.Labels{"chain": "ch", "node": "n"}}, + } + + for _, tt := range tests { + o := GetObserver(tt.name, tt.labels) + if o == nil { + t.Errorf("GetObserver(%q) returned nil", tt.name) + continue + } + o.Observe(0.5) + o.Observe(1.0) + } +} + +func TestPromMetricsUnknownNameReturnsNoop(t *testing.T) { + orig := IsEnabled() + Enable(true) + defer Enable(orig) + + // Unknown metric names should return noop implementations (non-nil, safe to call). + c := GetCounter("unknown_counter", nil) + if c == nil { + t.Error("counter for unknown name should not be nil") + } + c.Inc() + + g := GetGauge("unknown_gauge", nil) + if g == nil { + t.Error("gauge for unknown name should not be nil") + } + g.Set(1) + + o := GetObserver("unknown_observer", nil) + if o == nil { + t.Error("observer for unknown name should not be nil") + } + o.Observe(0.1) +} diff --git a/metrics/service/service.go b/metrics/service/service.go index f8faf965..b4f05a96 100644 --- a/metrics/service/service.go +++ b/metrics/service/service.go @@ -3,6 +3,7 @@ package service import ( "net" "net/http" + "sync" "github.com/go-gost/core/auth" "github.com/go-gost/core/service" @@ -10,6 +11,7 @@ import ( ) const ( + // DefaultPath is the default HTTP path for the metrics endpoint. DefaultPath = "/metrics" ) @@ -18,14 +20,17 @@ type options struct { auther auth.Authenticator } +// Option configures the metrics service. type Option func(*options) +// PathOption sets the HTTP path for the metrics endpoint. func PathOption(path string) Option { return func(o *options) { o.path = path } } +// AutherOption sets the authenticator for the metrics endpoint. func AutherOption(auther auth.Authenticator) Option { return func(o *options) { o.auther = auther @@ -33,11 +38,14 @@ func AutherOption(auther auth.Authenticator) Option { } type metricService struct { - s *http.Server - ln net.Listener - cclose chan struct{} + s *http.Server + ln net.Listener + cclose chan struct{} + closeOnce sync.Once + closeErr error } +// NewService creates a metrics Service that exposes Prometheus metrics over HTTP. func NewService(network, addr string, opts ...Option) (service.Service, error) { if network == "" { network = "tcp" @@ -75,18 +83,26 @@ func NewService(network, addr string, opts ...Option) (service.Service, error) { }, nil } +// Serve starts the metrics HTTP server and blocks until the listener is closed. func (s *metricService) Serve() error { return s.s.Serve(s.ln) } +// Addr returns the network address the metrics server is listening on. func (s *metricService) Addr() net.Addr { return s.ln.Addr() } +// Close stops the metrics HTTP server. It is safe to call multiple times. func (s *metricService) Close() error { - return s.s.Close() + s.closeOnce.Do(func() { + close(s.cclose) + s.closeErr = s.s.Close() + }) + return s.closeErr } +// IsClosed reports whether the metrics service has been closed. func (s *metricService) IsClosed() bool { select { case <-s.cclose: diff --git a/metrics/service/service_test.go b/metrics/service/service_test.go new file mode 100644 index 00000000..f070189b --- /dev/null +++ b/metrics/service/service_test.go @@ -0,0 +1,206 @@ +package service + +import ( + "context" + "net" + "net/http" + "testing" + "time" + + "github.com/go-gost/core/auth" +) + +// testAuther is a simple auth.Authenticator for testing. +type testAuther struct { + allowUser, allowPass string +} + +func (a *testAuther) Authenticate(ctx context.Context, user, password string, opts ...auth.Option) (string, bool) { + if user == a.allowUser && password == a.allowPass { + return user, true + } + return "", false +} + +func TestNewServiceDefaultPath(t *testing.T) { + s, err := NewService("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("NewService error: %v", err) + } + defer s.Close() + + if s.Addr() == nil { + t.Error("Addr should not be nil") + } +} + +func TestNewServiceDefaultNetwork(t *testing.T) { + s, err := NewService("", "127.0.0.1:0") + if err != nil { + t.Fatalf("NewService with empty network error: %v", err) + } + defer s.Close() + + if s.Addr() == nil { + t.Error("Addr should not be nil") + } +} + +func TestNewServiceCustomPath(t *testing.T) { + s, err := NewService("tcp", "127.0.0.1:0", PathOption("/custom")) + if err != nil { + t.Fatalf("NewService error: %v", err) + } + defer s.Close() + + // Serve in background for a short time. + go s.Serve() + time.Sleep(50 * time.Millisecond) + + // Default path should 404. + resp, err := http.Get("http://" + s.Addr().String() + "/metrics") + if err != nil { + t.Fatalf("GET /metrics error: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusNotFound { + t.Errorf("GET /metrics: expected 404, got %d", resp.StatusCode) + } + + // Custom path should 200. + resp2, err := http.Get("http://" + s.Addr().String() + "/custom") + if err != nil { + t.Fatalf("GET /custom error: %v", err) + } + resp2.Body.Close() + if resp2.StatusCode != http.StatusOK { + t.Errorf("GET /custom: expected 200, got %d", resp2.StatusCode) + } +} + +func TestNewServiceWithAuth(t *testing.T) { + auther := &testAuther{allowUser: "admin", allowPass: "secret"} + s, err := NewService("tcp", "127.0.0.1:0", AutherOption(auther)) + if err != nil { + t.Fatalf("NewService error: %v", err) + } + defer s.Close() + + go s.Serve() + time.Sleep(50 * time.Millisecond) + + addr := "http://" + s.Addr().String() + "/metrics" + + // No credentials should be 401. + resp, err := http.Get(addr) + if err != nil { + t.Fatalf("GET error: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 without credentials, got %d", resp.StatusCode) + } + + // Wrong credentials should be 401. + req, _ := http.NewRequest("GET", addr, nil) + req.SetBasicAuth("admin", "wrong") + resp2, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("GET with wrong auth error: %v", err) + } + resp2.Body.Close() + if resp2.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 with wrong credentials, got %d", resp2.StatusCode) + } + + // Correct credentials should be 200. + req3, _ := http.NewRequest("GET", addr, nil) + req3.SetBasicAuth("admin", "secret") + resp3, err := http.DefaultClient.Do(req3) + if err != nil { + t.Fatalf("GET with correct auth error: %v", err) + } + resp3.Body.Close() + if resp3.StatusCode != http.StatusOK { + t.Errorf("expected 200 with correct credentials, got %d", resp3.StatusCode) + } +} + +func TestServiceClose(t *testing.T) { + svc, err := NewService("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("NewService error: %v", err) + } + s := svc.(*metricService) + + if s.IsClosed() { + t.Error("IsClosed should be false before Close") + } + + // Serve in background. + go s.Serve() + time.Sleep(50 * time.Millisecond) + + err = s.Close() + if err != nil { + t.Errorf("Close error: %v", err) + } + + if !s.IsClosed() { + t.Error("IsClosed should be true after Close") + } + + // Double close should be safe. + err = s.Close() + if err != nil { + t.Errorf("second Close error: %v", err) + } +} + +func TestServiceServeReturnsOnClose(t *testing.T) { + s, err := NewService("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("NewService error: %v", err) + } + + done := make(chan error, 1) + go func() { + done <- s.Serve() + }() + + time.Sleep(50 * time.Millisecond) + s.Close() + + select { + case err := <-done: + if err != nil && err != http.ErrServerClosed { + t.Errorf("Serve returned unexpected error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return within 2s after Close") + } +} + +func TestNewServiceAddr(t *testing.T) { + s, err := NewService("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("NewService error: %v", err) + } + defer s.Close() + + addr := s.Addr() + if addr == nil { + t.Fatal("Addr should not be nil") + } + if _, ok := addr.(*net.TCPAddr); !ok { + t.Errorf("expected *net.TCPAddr, got %T", addr) + } +} + +func TestNewServiceListenError(t *testing.T) { + // Using an invalid address should cause a listen error. + _, err := NewService("invalid", "127.0.0.1:0") + if err == nil { + t.Error("expected error for invalid network, got nil") + } +} diff --git a/metrics/wrapper/conn.go b/metrics/wrapper/conn.go index fe4ec4a2..7bd19451 100644 --- a/metrics/wrapper/conn.go +++ b/metrics/wrapper/conn.go @@ -26,6 +26,8 @@ type serverConn struct { clientIP string } +// WrapConn wraps a net.Conn with metrics tracking. Bytes read/written are +// counted as service transfer input/output. If c is nil, nil is returned. func WrapConn(service string, c net.Conn) net.Conn { if c == nil { return c @@ -101,6 +103,8 @@ type packetConn struct { service string } +// WrapPacketConn wraps a net.PacketConn with metrics tracking. Bytes read/written +// are counted as service transfer input/output. If pc is nil, nil is returned. func WrapPacketConn(service string, pc net.PacketConn) net.PacketConn { if pc == nil { return pc @@ -161,7 +165,13 @@ type udpConn struct { service string } +// WrapUDPConn wraps a net.PacketConn as a udp.Conn with metrics tracking. Bytes +// read/written are counted as service transfer input/output. If pc is nil, nil +// is returned. func WrapUDPConn(service string, pc net.PacketConn) udp.Conn { + if pc == nil { + return nil + } return &udpConn{ PacketConn: pc, service: service, @@ -377,7 +387,7 @@ func (c *udpConn) SetDSCP(n int) error { if nc, ok := c.PacketConn.(xnet.SetDSCP); ok { return nc.SetDSCP(n) } - return nil + return errUnsupport } func (c *udpConn) Context() context.Context { diff --git a/metrics/wrapper/conn_test.go b/metrics/wrapper/conn_test.go new file mode 100644 index 00000000..3618dd95 --- /dev/null +++ b/metrics/wrapper/conn_test.go @@ -0,0 +1,230 @@ +package wrapper + +import ( + "io" + "net" + "sync" + "testing" + "time" + + xio "github.com/go-gost/x/internal/io" +) + +func TestWrapConnNil(t *testing.T) { + if c := WrapConn("svc", nil); c != nil { + t.Error("WrapConn(nil) should return nil") + } +} + +func TestWrapPacketConnNil(t *testing.T) { + if pc := WrapPacketConn("svc", nil); pc != nil { + t.Error("WrapPacketConn(nil) should return nil") + } +} + +func TestWrapUDPConnNil(t *testing.T) { + if uc := WrapUDPConn("svc", nil); uc != nil { + t.Error("WrapUDPConn(nil) should return nil") + } +} + +func TestWrapConnReadWrite(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + client.Write([]byte("hello")) + }() + + buf := make([]byte, 1024) + n, err := wrapped.Read(buf) + wg.Wait() + + if err != nil { + t.Fatalf("Read error: %v", err) + } + if n != 5 { + t.Errorf("expected 5 bytes, got %d", n) + } + if string(buf[:n]) != "hello" { + t.Errorf("expected 'hello', got %q", string(buf[:n])) + } + + // Write through the wrapped conn. + go func() { + buf := make([]byte, 1024) + n, err := client.Read(buf) + if err != nil { + t.Errorf("client Read error: %v", err) + return + } + if n != 5 || string(buf[:n]) != "world" { + t.Errorf("expected 'world', got %q", string(buf[:n])) + } + }() + + nw, err := wrapped.Write([]byte("world")) + if err != nil { + t.Fatalf("Write error: %v", err) + } + if nw != 5 { + t.Errorf("expected 5 bytes written, got %d", nw) + } +} + +func TestWrapConnSyscallConn(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + // SyscallConn should not panic. The return values depend on the underlying + // connection type. + _, _ = wrapped.(*serverConn).SyscallConn() +} + +func TestWrapConnCloseReadCloseWrite(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + // net.Pipe does not implement CloseRead/CloseWrite, so these should + // return ErrUnsupported. + err := wrapped.(*serverConn).CloseRead() + if err != xio.ErrUnsupported { + t.Errorf("CloseRead: expected ErrUnsupported, got %v", err) + } + + err = wrapped.(*serverConn).CloseWrite() + if err != xio.ErrUnsupported { + t.Errorf("CloseWrite: expected ErrUnsupported, got %v", err) + } +} + +func TestWrapConnContext(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + // net.Pipe does not implement ctx.Context, so Context() returns nil. + ctx := wrapped.(*serverConn).Context() + if ctx != nil { + t.Errorf("expected nil context, got %v", ctx) + } +} + +// mockPacketConn implements net.PacketConn for testing. +type mockPacketConn struct { + readFrom func([]byte) (int, net.Addr, error) + writeTo func([]byte, net.Addr) (int, error) +} + +func (m *mockPacketConn) ReadFrom(p []byte) (int, net.Addr, error) { + return m.readFrom(p) +} + +func (m *mockPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) { + return m.writeTo(p, addr) +} + +func (m *mockPacketConn) Close() error { return nil } +func (m *mockPacketConn) LocalAddr() net.Addr { return &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} } +func (m *mockPacketConn) SetDeadline(t time.Time) error { return nil } +func (m *mockPacketConn) SetReadDeadline(t time.Time) error { return nil } +func (m *mockPacketConn) SetWriteDeadline(t time.Time) error { return nil } + +func TestWrapPacketConnReadFromWriteTo(t *testing.T) { + testAddr := &net.UDPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 1234} + + mock := &mockPacketConn{ + readFrom: func(p []byte) (int, net.Addr, error) { + copy(p, []byte("data")) + return 4, testAddr, nil + }, + writeTo: func(p []byte, addr net.Addr) (int, error) { + return len(p), nil + }, + } + + wrapped := WrapPacketConn("test-svc", mock) + + // ReadFrom + buf := make([]byte, 1024) + n, addr, err := wrapped.ReadFrom(buf) + if err != nil { + t.Fatalf("ReadFrom error: %v", err) + } + if n != 4 { + t.Errorf("expected 4 bytes, got %d", n) + } + if addr.String() != testAddr.String() { + t.Errorf("expected addr %s, got %s", testAddr, addr) + } + + // WriteTo + nw, err := wrapped.WriteTo([]byte("reply"), testAddr) + if err != nil { + t.Fatalf("WriteTo error: %v", err) + } + if nw != 5 { + t.Errorf("expected 5 bytes written, got %d", nw) + } +} + +func TestWrapPacketConnContext(t *testing.T) { + mock := &mockPacketConn{ + readFrom: func(p []byte) (int, net.Addr, error) { return 0, nil, io.EOF }, + writeTo: func(p []byte, addr net.Addr) (int, error) { return len(p), nil }, + } + + wrapped := WrapPacketConn("test-svc", mock) + + // mockPacketConn does not implement ctx.Context + ctx := wrapped.(*packetConn).Context() + if ctx != nil { + t.Errorf("expected nil context, got %v", ctx) + } +} + +func TestWrapConnRemoteAddr(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + // The wrapped conn should expose the underlying remote address. + if wrapped.RemoteAddr() == nil { + t.Error("RemoteAddr should not be nil") + } +} + +func TestWrapConnNilLabels(t *testing.T) { + // When the remote address has no port, SplitHostPort may fail. + // WrapConn should still succeed and the host may be empty. + client, server := net.Pipe() + defer client.Close() + defer server.Close() + + wrapped := WrapConn("test-svc", server) + + // Read should not panic even with potentially malformed addresses. + // net.Pipe addresses are well-formed so this is a smoke test. + buf := make([]byte, 1) + go client.Write([]byte{0x01}) + _, err := wrapped.Read(buf) + if err != nil { + t.Fatalf("Read error: %v", err) + } +} diff --git a/metrics/wrapper/listener.go b/metrics/wrapper/listener.go index ed863037..f32efa7d 100644 --- a/metrics/wrapper/listener.go +++ b/metrics/wrapper/listener.go @@ -9,7 +9,13 @@ type listener struct { net.Listener } +// WrapListener wraps a net.Listener so that all accepted connections are +// automatically wrapped with WrapConn for metrics tracking. If ln is nil, +// nil is returned. func WrapListener(service string, ln net.Listener) net.Listener { + if ln == nil { + return ln + } return &listener{ service: service, Listener: ln, diff --git a/metrics/wrapper/listener_test.go b/metrics/wrapper/listener_test.go new file mode 100644 index 00000000..a6a840a9 --- /dev/null +++ b/metrics/wrapper/listener_test.go @@ -0,0 +1,81 @@ +package wrapper + +import ( + "errors" + "net" + "testing" +) + +func TestWrapListenerNil(t *testing.T) { + if ln := WrapListener("svc", nil); ln != nil { + t.Error("WrapListener(nil) should return nil") + } +} + +// mockListener implements net.Listener for testing. +type mockListener struct { + accept func() (net.Conn, error) +} + +func (l *mockListener) Accept() (net.Conn, error) { return l.accept() } +func (l *mockListener) Close() error { return nil } +func (l *mockListener) Addr() net.Addr { return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080} } + +func TestWrapListenerAccept(t *testing.T) { + client, server := net.Pipe() + + ln := &mockListener{ + accept: func() (net.Conn, error) { + return server, nil + }, + } + + wrapped := WrapListener("test-svc", ln) + + c, err := wrapped.Accept() + if err != nil { + t.Fatalf("Accept error: %v", err) + } + if c == nil { + t.Fatal("accepted connection should not be nil") + } + + // The accepted connection should be a *serverConn (wrapped by WrapConn). + if _, ok := c.(*serverConn); !ok { + t.Errorf("accepted connection should be *serverConn, got %T", c) + } + + client.Close() + c.Close() +} + +func TestWrapListenerAcceptError(t *testing.T) { + testErr := errors.New("accept failed") + ln := &mockListener{ + accept: func() (net.Conn, error) { + return nil, testErr + }, + } + + wrapped := WrapListener("test-svc", ln) + + c, err := wrapped.Accept() + if err != testErr { + t.Errorf("expected error %v, got %v", testErr, err) + } + if c != nil { + t.Error("connection should be nil on error") + } +} + +func TestWrapListenerAddr(t *testing.T) { + ln := &mockListener{ + accept: func() (net.Conn, error) { return nil, errors.New("not used") }, + } + + wrapped := WrapListener("test-svc", ln) + + if wrapped.Addr() == nil { + t.Error("Addr should not be nil") + } +}