TUN inbound: Add traffic counters; Metrics: Rely on instance (#6349)

https://github.com/XTLS/Xray-core/pull/6349#issuecomment-4775121300

proxy/tun/handler.go

@@ -17,6 +17,7 @@ import (
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/features/policy"
 	"github.com/xtls/xray-core/features/routing"
+	"github.com/xtls/xray-core/features/stats"
 	"github.com/xtls/xray-core/transport"
 	"github.com/xtls/xray-core/transport/internet"
 	"github.com/xtls/xray-core/transport/internet/stat"
@@ -32,6 +33,8 @@ type Handler struct {
 	dispatcher      routing.Dispatcher
 	tag             string
 	sniffingRequest session.SniffingRequest
+	uplinkCounter   stats.Counter
+	downlinkCounter stats.Counter
 }
 
 // ConnectionHandler interface with the only method that stack is going to push new connections to
@@ -59,6 +62,23 @@ func (t *Handler) Init(ctx context.Context, pm policy.Manager, dispatcher routin
 	t.policyManager = pm
 	t.dispatcher = dispatcher
 
+	if len(t.tag) > 0 && pm.ForSystem().Stats.InboundUplink {
+		statsManager := core.MustFromContext(ctx).GetFeature(stats.ManagerType()).(stats.Manager)
+		name := "inbound>>>" + t.tag + ">>>traffic>>>uplink"
+		c, _ := stats.GetOrRegisterCounter(statsManager, name)
+		if c != nil {
+			t.uplinkCounter = c
+		}
+	}
+	if len(t.tag) > 0 && pm.ForSystem().Stats.InboundDownlink {
+		statsManager := core.MustFromContext(ctx).GetFeature(stats.ManagerType()).(stats.Manager)
+		name := "inbound>>>" + t.tag + ">>>traffic>>>downlink"
+		c, _ := stats.GetOrRegisterCounter(statsManager, name)
+		if c != nil {
+			t.downlinkCounter = c
+		}
+	}
+
 	return nil
 }
 
@@ -151,6 +171,14 @@ func (t *Handler) HandleConnection(conn net.Conn, destination net.Destination) {
 		return
 	}
 	source := net.DestinationFromAddr(remote)
+	if t.uplinkCounter != nil || t.downlinkCounter != nil {
+		conn = &stat.CounterConnection{
+			Connection:   conn,
+			ReadCounter:  t.uplinkCounter,
+			WriteCounter: t.downlinkCounter,
+		}
+	}
+
 	inbound := session.Inbound{
 		Name:          "tun",
 		Tag:           t.tag,

proxy/tun/handler_test.go

@@ -0,0 +1,133 @@
+package tun
+
+import (
+	"bytes"
+	"context"
+	"net"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/xtls/xray-core/common/buf"
+	xnet "github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/features/routing"
+	"github.com/xtls/xray-core/transport"
+)
+
+type testCounter struct {
+	value int64
+}
+
+func (c *testCounter) Value() int64 {
+	return atomic.LoadInt64(&c.value)
+}
+
+func (c *testCounter) Set(value int64) int64 {
+	return atomic.SwapInt64(&c.value, value)
+}
+
+func (c *testCounter) Add(value int64) int64 {
+	return atomic.AddInt64(&c.value, value) - value
+}
+
+type testConn struct {
+	reader *bytes.Reader
+	writer bytes.Buffer
+}
+
+func newTestConn(input []byte) *testConn {
+	return &testConn{reader: bytes.NewReader(input)}
+}
+
+func (c *testConn) Read(payload []byte) (int, error) {
+	return c.reader.Read(payload)
+}
+
+func (c *testConn) Write(payload []byte) (int, error) {
+	return c.writer.Write(payload)
+}
+
+func (c *testConn) Close() error {
+	return nil
+}
+
+func (c *testConn) LocalAddr() net.Addr {
+	return &net.TCPAddr{IP: net.IPv4(10, 0, 0, 1), Port: 1080}
+}
+
+func (c *testConn) RemoteAddr() net.Addr {
+	return &net.TCPAddr{IP: net.IPv4(10, 0, 0, 2), Port: 12345}
+}
+
+func (c *testConn) SetDeadline(time.Time) error {
+	return nil
+}
+
+func (c *testConn) SetReadDeadline(time.Time) error {
+	return nil
+}
+
+func (c *testConn) SetWriteDeadline(time.Time) error {
+	return nil
+}
+
+type testDispatcher struct {
+	writePayload []byte
+	readBytes    int32
+}
+
+func (d *testDispatcher) Type() interface{} {
+	return routing.DispatcherType()
+}
+
+func (d *testDispatcher) Start() error {
+	return nil
+}
+
+func (d *testDispatcher) Close() error {
+	return nil
+}
+
+func (d *testDispatcher) Dispatch(context.Context, xnet.Destination) (*transport.Link, error) {
+	return nil, nil
+}
+
+func (d *testDispatcher) DispatchLink(ctx context.Context, dest xnet.Destination, link *transport.Link) error {
+	mb, err := link.Reader.ReadMultiBuffer()
+	if err != nil {
+		return err
+	}
+	atomic.StoreInt32(&d.readBytes, mb.Len())
+	buf.ReleaseMulti(mb)
+
+	return link.Writer.WriteMultiBuffer(buf.MultiBuffer{buf.FromBytes(d.writePayload)})
+}
+
+func TestHandlerCountsTunConnectionTraffic(t *testing.T) {
+	uplinkCounter := new(testCounter)
+	downlinkCounter := new(testCounter)
+	dispatcher := &testDispatcher{writePayload: []byte("downlink")}
+	conn := newTestConn([]byte("uplink"))
+
+	handler := &Handler{
+		ctx:             context.Background(),
+		config:          &Config{},
+		dispatcher:      dispatcher,
+		uplinkCounter:   uplinkCounter,
+		downlinkCounter: downlinkCounter,
+	}
+	handler.HandleConnection(conn, xnet.TCPDestination(xnet.LocalHostIP, 443))
+
+	if got := uplinkCounter.Value(); got != int64(len("uplink")) {
+		t.Fatalf("unexpected uplink counter: got %d, want %d", got, len("uplink"))
+	}
+	if got := downlinkCounter.Value(); got != int64(len("downlink")) {
+		t.Fatalf("unexpected downlink counter: got %d, want %d", got, len("downlink"))
+	}
+	if got := int(atomic.LoadInt32(&dispatcher.readBytes)); got != len("uplink") {
+		t.Fatalf("dispatcher read unexpected bytes: got %d, want %d", got, len("uplink"))
+	}
+	if got := conn.writer.String(); got != "downlink" {
+		t.Fatalf("connection write mismatch: got %q, want %q", got, "downlink")
+	}
+}