naming consitency, and add test

Author: thomasjib

tracker/metrics/conn.go

@@ -20,7 +20,7 @@ type Conn struct {
 // NewConn creates a new Conn instance.
 func NewConn(conn net.Conn, metadata *adapter.InboundContext) net.Conn {
 	attributes := metadataToAttributes(metadata)
-	metrics.Connections.Add(context.Background(), 1, metric.WithAttributes(attributes...))
+	metrics.conns.Add(context.Background(), 1, metric.WithAttributes(attributes...))
 	return &Conn{
 		Conn:       conn,
 		attributes: attributes,
@@ -33,7 +33,7 @@ func (c *Conn) Read(b []byte) (n int, err error) {
 	n, err = c.Conn.Read(b)
 	if n > 0 {
 		attrs := append(c.attributes, attribute.KeyValue{Key: "direction", Value: attribute.StringValue("receive")})
-		metrics.ProxyIO.Add(context.Background(), int64(n), metric.WithAttributes(attrs...))
+		metrics.proxyIO.Add(context.Background(), int64(n), metric.WithAttributes(attrs...))
 	}
 	return
 }
@@ -43,7 +43,7 @@ func (c *Conn) Write(b []byte) (n int, err error) {
 	n, err = c.Conn.Write(b)
 	if n > 0 {
 		attrs := append(c.attributes, attribute.KeyValue{Key: "direction", Value: attribute.StringValue("transmit")})
-		metrics.ProxyIO.Add(context.Background(), int64(n), metric.WithAttributes(attrs...))
+		metrics.proxyIO.Add(context.Background(), int64(n), metric.WithAttributes(attrs...))
 	}
 	return
 }

tracker/metrics/metrics.go

@@ -18,11 +18,10 @@ import (
 )
 
 type metricsManager struct {
-	meter       metric.Meter
-	ProxyIO     metric.Int64Counter
-	Connections metric.Int64Counter
-	conns       metric.Int64UpDownCounter
-	duration    metric.Int64Histogram
+	meter    metric.Meter
+	proxyIO  metric.Int64Counter
+	conns    metric.Int64UpDownCounter
+	duration metric.Int64Histogram
 
 	countryLookup geo.CountryLookup
 }
@@ -32,33 +31,28 @@ var metrics = &metricsManager{}
 func SetupMetricsManager(geolite2CityURL, cityDBFile string) {
 	meter := otel.GetMeterProvider().Meter("lantern-box")
 
+	// Track the number of bytes sent and recieved.
 	pIO, err := meter.Int64Counter("proxy.io", metric.WithUnit("bytes"))
 	if err != nil {
 		pIO = &noop.Int64Counter{}
 	}
 
-	connections, err := meter.Int64Counter("proxy.connections")
-	if err != nil {
-		connections = &noop.Int64Counter{}
-	}
-
 	// Track the number of connections.
-	conns, err := meter.Int64UpDownCounter("sing.connections", metric.WithDescription("Number of connections"))
+	conns, err := meter.Int64UpDownCounter("proxy.connections", metric.WithDescription("Number of connections"))
 	if err != nil {
 		conns = &noop.Int64UpDownCounter{}
 	}
 
 	// Track connection duration.
-	duration, err := meter.Int64Histogram("sing.connection_duration", metric.WithDescription("Connection duration"))
+	duration, err := meter.Int64Histogram("proxy.connection.duration", metric.WithDescription("Connection duration"))
 	if err != nil {
 		duration = &noop.Int64Histogram{}
 	}
 
 	metrics.meter = meter
-	metrics.ProxyIO = pIO
-	metrics.duration = duration
-	metrics.Connections = connections
+	metrics.proxyIO = pIO
 	metrics.conns = conns
+	metrics.duration = duration
 
 	metrics.countryLookup = geo.FromWeb(geolite2CityURL, cityDBFile, 24*time.Hour, cityDBFile, geo.CountryCode)
 	if metrics.countryLookup == nil {

tracker/metrics/packet_conn.go

@@ -22,7 +22,7 @@ type PacketConn struct {
 // NewPacketConn creates a new PacketConn instance.
 func NewPacketConn(conn N.PacketConn, metadata *adapter.InboundContext) N.PacketConn {
 	attributes := metadataToAttributes(metadata)
-	metrics.Connections.Add(context.Background(), 1, metric.WithAttributes(attributes...))
+	metrics.conns.Add(context.Background(), 1, metric.WithAttributes(attributes...))
 
 	return &PacketConn{
 		PacketConn: conn,
@@ -39,7 +39,7 @@ func (c *PacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, er
 	}
 	if buffer.Len() > 0 {
 		attrs := append(c.attributes, attribute.KeyValue{Key: "direction", Value: attribute.StringValue("receive")})
-		metrics.ProxyIO.Add(context.Background(), int64(buffer.Len()), metric.WithAttributes(attrs...))
+		metrics.proxyIO.Add(context.Background(), int64(buffer.Len()), metric.WithAttributes(attrs...))
 	}
 	return dest, nil
 }
@@ -48,7 +48,7 @@ func (c *PacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, er
 func (c *PacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
 	if buffer.Len() > 0 {
 		attrs := append(c.attributes, attribute.KeyValue{Key: "direction", Value: attribute.StringValue("transmit")})
-		metrics.ProxyIO.Add(context.Background(), int64(buffer.Len()), metric.WithAttributes(attrs...))
+		metrics.proxyIO.Add(context.Background(), int64(buffer.Len()), metric.WithAttributes(attrs...))
 	}
 	return c.PacketConn.WritePacket(buffer, destination)
 }

tracker/metrics/tracker_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	sdkotel "go.opentelemetry.io/otel"
 
@@ -15,9 +16,13 @@ import (
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/sagernet/sing/common/buf"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
 )
 
-func TestTracker(t *testing.T) {
+func TestConnTracker(t *testing.T) {
 	var wg sync.WaitGroup
 	reader := metric.NewManualReader()
 	provider := metric.NewMeterProvider(metric.WithReader(reader))
@@ -84,6 +89,81 @@ func TestTracker(t *testing.T) {
 	}
 }
 
+func TestPacketConnTracker(t *testing.T) {
+	var wg sync.WaitGroup
+	reader := metric.NewManualReader()
+	provider := metric.NewMeterProvider(metric.WithReader(reader))
+	sdkotel.SetMeterProvider(provider)
+	SetupMetricsManager("", "")
+
+	metricsTracker, err := NewTracker()
+	assert.NoError(t, err)
+
+	client, server := packetPipe()
+	defer client.Close()
+	defer server.Close()
+
+	ctx := context.Background()
+	serverTracked := metricsTracker.RoutedPacketConnection(ctx, server, adapter.InboundContext{}, nil, nil)
+
+	wg.Add(1)
+	serverReceive := 0
+	go func() {
+		defer wg.Done()
+		buffer := buf.New()
+		_, err = serverTracked.ReadPacket(buffer)
+		assert.NoError(t, err)
+		serverReceive = buffer.Len()
+		t.Logf("Server received '%v'\n", string(buffer.Bytes()))
+	}()
+
+	clientSentMessage := []byte("A client sent a short request...")
+	clientSendBuf := buf.NewSize(len(clientSentMessage))
+	clientSendBuf.Write(clientSentMessage)
+	err = client.WritePacket(clientSendBuf, server.LocalAddr().(M.Socksaddr))
+	assert.NoError(t, err)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		buffer := buf.New()
+		_, err = client.ReadPacket(buffer)
+		assert.NoError(t, err)
+		t.Logf("Client received '%v'\n", string(buffer.Bytes()))
+
+	}()
+
+	serverSentMessage := []byte("...and the server sent a short response.")
+	serverSendBuf := buf.NewSize(len(serverSentMessage))
+	serverSendBuf.Write(serverSentMessage)
+	serverTransmit := serverSendBuf.Len()
+	err = serverTracked.WritePacket(serverSendBuf, client.LocalAddr().(M.Socksaddr))
+	assert.NoError(t, err)
+
+	wg.Wait()
+	var rm metricdata.ResourceMetrics
+	reader.Collect(ctx, &rm)
+
+	ioCounter := extractCountersByAttribute(rm, "proxy.io")
+	results := map[string]int64{}
+	for k, v := range ioCounter {
+		if strings.Contains(k, "direction=transmit") {
+			results["transmit"] += v
+		} else if strings.Contains(k, "direction=receive") {
+			results["receive"] += v
+		}
+	}
+	for k, v := range results {
+		t.Logf("%s: %d\n", k, v)
+	}
+	if results["transmit"] != int64(serverTransmit) {
+		t.Errorf("transmit bytes did not match, got %d, want %d", results["transmit"], serverTransmit)
+	}
+	if results["receive"] != int64(serverReceive) {
+		t.Errorf("receive bytes did not match, got %d, want %d", results["receive"], serverReceive)
+	}
+}
+
 func extractCountersByAttribute(rm metricdata.ResourceMetrics, name string) map[string]int64 {
 	result := make(map[string]int64)
 	for _, sm := range rm.ScopeMetrics {
@@ -99,3 +179,63 @@ func extractCountersByAttribute(rm metricdata.ResourceMetrics, name string) map[
 	}
 	return result
 }
+
+// Mock packet pipe connection for testing
+type packet struct {
+	buffer *buf.Buffer
+	dest   M.Socksaddr
+}
+
+var _ N.PacketConn = (*packetPipeConn)(nil)
+
+type packetPipeConn struct {
+	readCh  chan packet
+	writeCh chan packet
+	once    sync.Once
+	done    chan struct{}
+}
+
+func packetPipe() (N.PacketConn, N.PacketConn) {
+	ch1 := make(chan packet, 16)
+	ch2 := make(chan packet, 16)
+	done1 := make(chan struct{})
+	done2 := make(chan struct{})
+
+	c1 := &packetPipeConn{readCh: ch1, writeCh: ch2, done: done1}
+	c2 := &packetPipeConn{readCh: ch2, writeCh: ch1, done: done2}
+	return c1, c2
+}
+
+func (c *packetPipeConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
+	select {
+	case p := <-c.readCh:
+		_, err = buffer.ReadOnceFrom(p.buffer)
+		p.buffer.Release()
+		return p.dest, err
+	case <-c.done:
+		return M.Socksaddr{}, net.ErrClosed
+	}
+}
+
+func (c *packetPipeConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
+	newBuf := buf.NewSize(buffer.Len())
+	newBuf.Write(buffer.Bytes())
+	buffer.Release()
+	select {
+	case c.writeCh <- packet{buffer: newBuf, dest: destination}:
+		return nil
+	case <-c.done:
+		newBuf.Release()
+		return net.ErrClosed
+	}
+}
+
+func (c *packetPipeConn) Close() error {
+	c.once.Do(func() { close(c.done) })
+	return nil
+}
+
+func (c *packetPipeConn) LocalAddr() net.Addr                { return M.Socksaddr{} }
+func (c *packetPipeConn) SetDeadline(t time.Time) error      { return nil }
+func (c *packetPipeConn) SetReadDeadline(t time.Time) error  { return nil }
+func (c *packetPipeConn) SetWriteDeadline(t time.Time) error { return nil }