github.com/sagernet/sing-box@v1.2.7/experimental/trackerconn/conn.go (about)

     1  package trackerconn
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/sagernet/sing/common/atomic"
     7  	"github.com/sagernet/sing/common/buf"
     8  	"github.com/sagernet/sing/common/bufio"
     9  	N "github.com/sagernet/sing/common/network"
    10  )
    11  
    12  func New(conn net.Conn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64) *Conn {
    13  	return &Conn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
    14  }
    15  
    16  func NewHook(conn net.Conn, readCounter func(n int64), writeCounter func(n int64)) *HookConn {
    17  	return &HookConn{bufio.NewExtendedConn(conn), readCounter, writeCounter}
    18  }
    19  
    20  type Conn struct {
    21  	N.ExtendedConn
    22  	readCounter  []*atomic.Int64
    23  	writeCounter []*atomic.Int64
    24  }
    25  
    26  func (c *Conn) Read(p []byte) (n int, err error) {
    27  	n, err = c.ExtendedConn.Read(p)
    28  	for _, counter := range c.readCounter {
    29  		counter.Add(int64(n))
    30  	}
    31  	return n, err
    32  }
    33  
    34  func (c *Conn) ReadBuffer(buffer *buf.Buffer) error {
    35  	err := c.ExtendedConn.ReadBuffer(buffer)
    36  	if err != nil {
    37  		return err
    38  	}
    39  	for _, counter := range c.readCounter {
    40  		counter.Add(int64(buffer.Len()))
    41  	}
    42  	return nil
    43  }
    44  
    45  func (c *Conn) Write(p []byte) (n int, err error) {
    46  	n, err = c.ExtendedConn.Write(p)
    47  	for _, counter := range c.writeCounter {
    48  		counter.Add(int64(n))
    49  	}
    50  	return n, err
    51  }
    52  
    53  func (c *Conn) WriteBuffer(buffer *buf.Buffer) error {
    54  	dataLen := int64(buffer.Len())
    55  	err := c.ExtendedConn.WriteBuffer(buffer)
    56  	if err != nil {
    57  		return err
    58  	}
    59  	for _, counter := range c.writeCounter {
    60  		counter.Add(dataLen)
    61  	}
    62  	return nil
    63  }
    64  
    65  func (c *Conn) Upstream() any {
    66  	return c.ExtendedConn
    67  }
    68  
    69  type HookConn struct {
    70  	N.ExtendedConn
    71  	readCounter  func(n int64)
    72  	writeCounter func(n int64)
    73  }
    74  
    75  func (c *HookConn) Read(p []byte) (n int, err error) {
    76  	n, err = c.ExtendedConn.Read(p)
    77  	c.readCounter(int64(n))
    78  	return n, err
    79  }
    80  
    81  func (c *HookConn) ReadBuffer(buffer *buf.Buffer) error {
    82  	err := c.ExtendedConn.ReadBuffer(buffer)
    83  	if err != nil {
    84  		return err
    85  	}
    86  	c.readCounter(int64(buffer.Len()))
    87  	return nil
    88  }
    89  
    90  func (c *HookConn) Write(p []byte) (n int, err error) {
    91  	n, err = c.ExtendedConn.Write(p)
    92  	c.writeCounter(int64(n))
    93  	return n, err
    94  }
    95  
    96  func (c *HookConn) WriteBuffer(buffer *buf.Buffer) error {
    97  	dataLen := int64(buffer.Len())
    98  	err := c.ExtendedConn.WriteBuffer(buffer)
    99  	if err != nil {
   100  		return err
   101  	}
   102  	c.writeCounter(dataLen)
   103  	return nil
   104  }
   105  
   106  func (c *HookConn) Upstream() any {
   107  	return c.ExtendedConn
   108  }