Fix v2ray api

experimental/trackerconn/conn.go

@@ -1,7 +1,6 @@
 package trackerconn
 
 import (
-	"io"
 	"net"
 
 	"github.com/sagernet/sing/common/buf"
@@ -11,33 +10,25 @@ import (
 	"go.uber.org/atomic"
 )
 
-func New(conn net.Conn, readCounter *atomic.Int64, writeCounter *atomic.Int64, direct bool) N.ExtendedConn {
-	trackerConn := &Conn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
-	if direct {
-		return (*DirectConn)(trackerConn)
-	} else {
-		return trackerConn
-	}
+func New(conn net.Conn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64, direct bool) *Conn {
+	return &Conn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
 }
 
-func NewHook(conn net.Conn, readCounter func(n int64), writeCounter func(n int64), direct bool) N.ExtendedConn {
-	trackerConn := &HookConn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
-	if direct {
-		return (*DirectHookConn)(trackerConn)
-	} else {
-		return trackerConn
-	}
+func NewHook(conn net.Conn, readCounter func(n int64), writeCounter func(n int64), direct bool) *HookConn {
+	return &HookConn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
 }
 
 type Conn struct {
 	N.ExtendedConn
-	readCounter  *atomic.Int64
-	writeCounter *atomic.Int64
+	readCounter  []*atomic.Int64
+	writeCounter []*atomic.Int64
 }
 
 func (c *Conn) Read(p []byte) (n int, err error) {
 	n, err = c.ExtendedConn.Read(p)
-	c.readCounter.Add(int64(n))
+	for _, counter := range c.readCounter {
+		counter.Add(int64(n))
+	}
 	return n, err
 }
 
@@ -46,13 +37,17 @@ func (c *Conn) ReadBuffer(buffer *buf.Buffer) error {
 	if err != nil {
 		return err
 	}
-	c.readCounter.Add(int64(buffer.Len()))
+	for _, counter := range c.readCounter {
+		counter.Add(int64(buffer.Len()))
+	}
 	return nil
 }
 
 func (c *Conn) Write(p []byte) (n int, err error) {
 	n, err = c.ExtendedConn.Write(p)
-	c.writeCounter.Add(int64(n))
+	for _, counter := range c.writeCounter {
+		counter.Add(int64(n))
+	}
 	return n, err
 }
 
@@ -62,7 +57,9 @@ func (c *Conn) WriteBuffer(buffer *buf.Buffer) error {
 	if err != nil {
 		return err
 	}
-	c.writeCounter.Add(dataLen)
+	for _, counter := range c.writeCounter {
+		counter.Add(dataLen)
+	}
 	return nil
 }
 
@@ -110,51 +107,3 @@ func (c *HookConn) WriteBuffer(buffer *buf.Buffer) error {
 func (c *HookConn) Upstream() any {
 	return c.ExtendedConn
 }
-
-type DirectConn Conn
-
-func (c *DirectConn) WriteTo(w io.Writer) (n int64, err error) {
-	reader := N.UnwrapReader(c.ExtendedConn)
-	if wt, ok := reader.(io.WriterTo); ok {
-		n, err = wt.WriteTo(w)
-		c.readCounter.Add(n)
-		return
-	} else {
-		return bufio.Copy(w, (*Conn)(c))
-	}
-}
-
-func (c *DirectConn) ReadFrom(r io.Reader) (n int64, err error) {
-	writer := N.UnwrapWriter(c.ExtendedConn)
-	if rt, ok := writer.(io.ReaderFrom); ok {
-		n, err = rt.ReadFrom(r)
-		c.writeCounter.Add(n)
-		return
-	} else {
-		return bufio.Copy((*Conn)(c), r)
-	}
-}
-
-type DirectHookConn HookConn
-
-func (c *DirectHookConn) WriteTo(w io.Writer) (n int64, err error) {
-	reader := N.UnwrapReader(c.ExtendedConn)
-	if wt, ok := reader.(io.WriterTo); ok {
-		n, err = wt.WriteTo(w)
-		c.readCounter(n)
-		return
-	} else {
-		return bufio.Copy(w, (*HookConn)(c))
-	}
-}
-
-func (c *DirectHookConn) ReadFrom(r io.Reader) (n int64, err error) {
-	writer := N.UnwrapWriter(c.ExtendedConn)
-	if rt, ok := writer.(io.ReaderFrom); ok {
-		n, err = rt.ReadFrom(r)
-		c.writeCounter(n)
-		return
-	} else {
-		return bufio.Copy((*HookConn)(c), r)
-	}
-}

experimental/trackerconn/packet_conn.go

@@ -8,7 +8,7 @@ import (
 	"go.uber.org/atomic"
 )
 
-func NewPacket(conn N.PacketConn, readCounter *atomic.Int64, writeCounter *atomic.Int64) *PacketConn {
+func NewPacket(conn N.PacketConn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64) *PacketConn {
 	return &PacketConn{conn, readCounter, writeCounter}
 }
 
@@ -18,14 +18,16 @@ func NewHookPacket(conn N.PacketConn, readCounter func(n int64), writeCounter fu
 
 type PacketConn struct {
 	N.PacketConn
-	readCounter  *atomic.Int64
-	writeCounter *atomic.Int64
+	readCounter  []*atomic.Int64
+	writeCounter []*atomic.Int64
 }
 
 func (c *PacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
 	destination, err = c.PacketConn.ReadPacket(buffer)
 	if err == nil {
-		c.readCounter.Add(int64(buffer.Len()))
+		for _, counter := range c.readCounter {
+			counter.Add(int64(buffer.Len()))
+		}
 	}
 	return
 }
@@ -36,7 +38,9 @@ func (c *PacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) er
 	if err != nil {
 		return err
 	}
-	c.writeCounter.Add(dataLen)
+	for _, counter := range c.writeCounter {
+		counter.Add(dataLen)
+	}
 	return nil
 }
 

experimental/v2rayapi/stats.go

@@ -58,8 +58,8 @@ func NewStatsService(options option.V2RayStatsServiceOptions) *StatsService {
 }
 
 func (s *StatsService) RoutedConnection(inbound string, outbound string, conn net.Conn) net.Conn {
-	var readCounter *atomic.Int64
-	var writeCounter *atomic.Int64
+	var readCounter []*atomic.Int64
+	var writeCounter []*atomic.Int64
 	countInbound := inbound != "" && s.inbounds[inbound]
 	countOutbound := outbound != "" && s.outbounds[outbound]
 	if !countInbound && !countOutbound {
@@ -67,20 +67,20 @@ func (s *StatsService) RoutedConnection(inbound string, outbound string, conn ne
 	}
 	s.access.Lock()
 	if countInbound {
-		readCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink", readCounter)
-		writeCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink", writeCounter)
+		readCounter = append(readCounter, s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink"))
+		writeCounter = append(writeCounter, s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink"))
 	}
 	if countOutbound {
-		readCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink", readCounter)
-		writeCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink", writeCounter)
+		readCounter = append(readCounter, s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink"))
+		writeCounter = append(writeCounter, s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink"))
 	}
 	s.access.Unlock()
 	return trackerconn.New(conn, readCounter, writeCounter, s.directIO)
 }
 
 func (s *StatsService) RoutedPacketConnection(inbound string, outbound string, conn N.PacketConn) N.PacketConn {
-	var readCounter *atomic.Int64
-	var writeCounter *atomic.Int64
+	var readCounter []*atomic.Int64
+	var writeCounter []*atomic.Int64
 	countInbound := inbound != "" && s.inbounds[inbound]
 	countOutbound := outbound != "" && s.outbounds[outbound]
 	if !countInbound && !countOutbound {
@@ -88,12 +88,12 @@ func (s *StatsService) RoutedPacketConnection(inbound string, outbound string, c
 	}
 	s.access.Lock()
 	if countInbound {
-		readCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink", readCounter)
-		writeCounter = s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink", writeCounter)
+		readCounter = append(readCounter, s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>uplink"))
+		writeCounter = append(writeCounter, s.loadOrCreateCounter("inbound>>>"+inbound+">>>traffic>>>downlink"))
 	}
 	if countOutbound {
-		readCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink", readCounter)
-		writeCounter = s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink", writeCounter)
+		readCounter = append(readCounter, s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>uplink"))
+		writeCounter = append(writeCounter, s.loadOrCreateCounter("outbound>>>"+outbound+">>>traffic>>>downlink"))
 	}
 	s.access.Unlock()
 	return trackerconn.NewPacket(conn, readCounter, writeCounter)
@@ -119,7 +119,17 @@ func (s *StatsService) QueryStats(ctx context.Context, request *QueryStatsReques
 	var response QueryStatsResponse
 	s.access.Lock()
 	defer s.access.Unlock()
-	if request.Regexp {
+	if len(request.Patterns) == 0 {
+		for name, counter := range s.counters {
+			var value int64
+			if request.Reset_ {
+				value = counter.Swap(0)
+			} else {
+				value = counter.Load()
+			}
+			response.Stat = append(response.Stat, &Stat{Name: name, Value: value})
+		}
+	} else if request.Regexp {
 		matchers := make([]*regexp.Regexp, 0, len(request.Patterns))
 		for _, pattern := range request.Patterns {
 			matcher, err := regexp.Compile(pattern)
@@ -182,13 +192,12 @@ func (s *StatsService) mustEmbedUnimplementedStatsServiceServer() {
 }
 
 //nolint:staticcheck
-func (s *StatsService) loadOrCreateCounter(name string, counter *atomic.Int64) *atomic.Int64 {
+func (s *StatsService) loadOrCreateCounter(name string) *atomic.Int64 {
 	counter, loaded := s.counters[name]
-	if !loaded {
-		if counter == nil {
-			counter = atomic.NewInt64(0)
-		}
-		s.counters[name] = counter
+	if loaded {
+		return counter
 	}
+	counter = atomic.NewInt64(0)
+	s.counters[name] = counter
 	return counter
 }