fix(metrics): add doc comments, fix labels mutation and nil returns, add idempotent Close, add nil guards, add tests

- Fix promMetrics.Gauge/Counter/Observer mutating caller's labels map by
  using maps.Copy instead of direct assignment
- Fix Gauge/Counter/Observer returning nil for unknown metric names
  (now return noop implementations)
- Fix metricService.Close() race with sync.Once for safe concurrent calls
- Fix SetDSCP returning nil instead of errUnsupport on unsupported ops
- Add nil guards to WrapConn, WrapPacketConn, WrapUDPConn, WrapListener
- Add doc comments to all exported symbols
- Add 45 unit tests (metrics, noop, prom, service, wrapper)
This commit is contained in:
ginuerzh
2026-05-24 22:50:41 +08:00
parent 64f9a048f8
commit bc3d12ec2c
12 changed files with 950 additions and 23 deletions
+9
View File
@@ -34,14 +34,19 @@ var (
enabled atomic.Bool 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) { func Enable(b bool) {
enabled.Store(b) enabled.Store(b)
} }
// IsEnabled reports whether metrics collection is enabled.
func IsEnabled() bool { func IsEnabled() bool {
return enabled.Load() 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 { func GetCounter(name metrics.MetricName, labels metrics.Labels) metrics.Counter {
if IsEnabled() { if IsEnabled() {
return defaultMetrics.Counter(name, labels) return defaultMetrics.Counter(name, labels)
@@ -49,6 +54,8 @@ func GetCounter(name metrics.MetricName, labels metrics.Labels) metrics.Counter
return noop.Counter(name, labels) 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 { func GetGauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge {
if IsEnabled() { if IsEnabled() {
return defaultMetrics.Gauge(name, labels) return defaultMetrics.Gauge(name, labels)
@@ -56,6 +63,8 @@ func GetGauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge {
return noop.Gauge(name, labels) 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 { func GetObserver(name metrics.MetricName, labels metrics.Labels) metrics.Observer {
if IsEnabled() { if IsEnabled() {
return defaultMetrics.Observer(name, labels) return defaultMetrics.Observer(name, labels)
+173
View File
@@ -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)
}
}
}
+2
View File
@@ -12,6 +12,8 @@ var (
type noopMetrics struct{} type noopMetrics struct{}
// Noop returns a Metrics implementation whose counters, gauges, and observers
// are all no-ops.
func Noop() metrics.Metrics { func Noop() metrics.Metrics {
return noop return noop
} }
+80
View File
@@ -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)
}
+18 -18
View File
@@ -1,6 +1,7 @@
package metrics package metrics
import ( import (
"maps"
"os" "os"
"github.com/go-gost/core/metrics" "github.com/go-gost/core/metrics"
@@ -14,6 +15,8 @@ type promMetrics struct {
histograms map[metrics.MetricName]*prometheus.HistogramVec 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 { func NewMetrics() metrics.Metrics {
host, _ := os.Hostname() host, _ := os.Hostname()
m := &promMetrics{ m := &promMetrics{
@@ -107,35 +110,32 @@ func NewMetrics() metrics.Metrics {
func (m *promMetrics) Gauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge { func (m *promMetrics) Gauge(name metrics.MetricName, labels metrics.Labels) metrics.Gauge {
v, ok := m.gauges[name] v, ok := m.gauges[name]
if !ok { if !ok {
return nil return nopGauge
} }
if labels == nil { plabels := prometheus.Labels{}
labels = metrics.Labels{} maps.Copy(plabels, labels)
} plabels["host"] = m.host
labels["host"] = m.host return v.With(plabels)
return v.With(prometheus.Labels(labels))
} }
func (m *promMetrics) Counter(name metrics.MetricName, labels metrics.Labels) metrics.Counter { func (m *promMetrics) Counter(name metrics.MetricName, labels metrics.Labels) metrics.Counter {
v, ok := m.counters[name] v, ok := m.counters[name]
if !ok { if !ok {
return nil return nopCounter
} }
if labels == nil { plabels := prometheus.Labels{}
labels = metrics.Labels{} maps.Copy(plabels, labels)
} plabels["host"] = m.host
labels["host"] = m.host return v.With(plabels)
return v.With(prometheus.Labels(labels))
} }
func (m *promMetrics) Observer(name metrics.MetricName, labels metrics.Labels) metrics.Observer { func (m *promMetrics) Observer(name metrics.MetricName, labels metrics.Labels) metrics.Observer {
v, ok := m.histograms[name] v, ok := m.histograms[name]
if !ok { if !ok {
return nil return nopObserver
} }
if labels == nil { plabels := prometheus.Labels{}
labels = metrics.Labels{} maps.Copy(plabels, labels)
} plabels["host"] = m.host
labels["host"] = m.host return v.With(plabels)
return v.With(prometheus.Labels(labels))
} }
+114
View File
@@ -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)
}
+17 -1
View File
@@ -3,6 +3,7 @@ package service
import ( import (
"net" "net"
"net/http" "net/http"
"sync"
"github.com/go-gost/core/auth" "github.com/go-gost/core/auth"
"github.com/go-gost/core/service" "github.com/go-gost/core/service"
@@ -10,6 +11,7 @@ import (
) )
const ( const (
// DefaultPath is the default HTTP path for the metrics endpoint.
DefaultPath = "/metrics" DefaultPath = "/metrics"
) )
@@ -18,14 +20,17 @@ type options struct {
auther auth.Authenticator auther auth.Authenticator
} }
// Option configures the metrics service.
type Option func(*options) type Option func(*options)
// PathOption sets the HTTP path for the metrics endpoint.
func PathOption(path string) Option { func PathOption(path string) Option {
return func(o *options) { return func(o *options) {
o.path = path o.path = path
} }
} }
// AutherOption sets the authenticator for the metrics endpoint.
func AutherOption(auther auth.Authenticator) Option { func AutherOption(auther auth.Authenticator) Option {
return func(o *options) { return func(o *options) {
o.auther = auther o.auther = auther
@@ -36,8 +41,11 @@ type metricService struct {
s *http.Server s *http.Server
ln net.Listener ln net.Listener
cclose chan struct{} 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) { func NewService(network, addr string, opts ...Option) (service.Service, error) {
if network == "" { if network == "" {
network = "tcp" network = "tcp"
@@ -75,18 +83,26 @@ func NewService(network, addr string, opts ...Option) (service.Service, error) {
}, nil }, nil
} }
// Serve starts the metrics HTTP server and blocks until the listener is closed.
func (s *metricService) Serve() error { func (s *metricService) Serve() error {
return s.s.Serve(s.ln) return s.s.Serve(s.ln)
} }
// Addr returns the network address the metrics server is listening on.
func (s *metricService) Addr() net.Addr { func (s *metricService) Addr() net.Addr {
return s.ln.Addr() return s.ln.Addr()
} }
// Close stops the metrics HTTP server. It is safe to call multiple times.
func (s *metricService) Close() error { 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 { func (s *metricService) IsClosed() bool {
select { select {
case <-s.cclose: case <-s.cclose:
+206
View File
@@ -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")
}
}
+11 -1
View File
@@ -26,6 +26,8 @@ type serverConn struct {
clientIP string 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 { func WrapConn(service string, c net.Conn) net.Conn {
if c == nil { if c == nil {
return c return c
@@ -101,6 +103,8 @@ type packetConn struct {
service string 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 { func WrapPacketConn(service string, pc net.PacketConn) net.PacketConn {
if pc == nil { if pc == nil {
return pc return pc
@@ -161,7 +165,13 @@ type udpConn struct {
service string 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 { func WrapUDPConn(service string, pc net.PacketConn) udp.Conn {
if pc == nil {
return nil
}
return &udpConn{ return &udpConn{
PacketConn: pc, PacketConn: pc,
service: service, service: service,
@@ -377,7 +387,7 @@ func (c *udpConn) SetDSCP(n int) error {
if nc, ok := c.PacketConn.(xnet.SetDSCP); ok { if nc, ok := c.PacketConn.(xnet.SetDSCP); ok {
return nc.SetDSCP(n) return nc.SetDSCP(n)
} }
return nil return errUnsupport
} }
func (c *udpConn) Context() context.Context { func (c *udpConn) Context() context.Context {
+230
View File
@@ -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)
}
}
+6
View File
@@ -9,7 +9,13 @@ type listener struct {
net.Listener 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 { func WrapListener(service string, ln net.Listener) net.Listener {
if ln == nil {
return ln
}
return &listener{ return &listener{
service: service, service: service,
Listener: ln, Listener: ln,
+81
View File
@@ -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")
}
}