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

     1  package trackerconn
     2  
     3  import (
     4  	"github.com/sagernet/sing/common/atomic"
     5  	"github.com/sagernet/sing/common/buf"
     6  	M "github.com/sagernet/sing/common/metadata"
     7  	N "github.com/sagernet/sing/common/network"
     8  )
     9  
    10  func NewPacket(conn N.PacketConn, readCounter []*atomic.Int64, writeCounter []*atomic.Int64) *PacketConn {
    11  	return &PacketConn{conn, readCounter, writeCounter}
    12  }
    13  
    14  func NewHookPacket(conn N.PacketConn, readCounter func(n int64), writeCounter func(n int64)) *HookPacketConn {
    15  	return &HookPacketConn{conn, readCounter, writeCounter}
    16  }
    17  
    18  type PacketConn struct {
    19  	N.PacketConn
    20  	readCounter  []*atomic.Int64
    21  	writeCounter []*atomic.Int64
    22  }
    23  
    24  func (c *PacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
    25  	destination, err = c.PacketConn.ReadPacket(buffer)
    26  	if err == nil {
    27  		for _, counter := range c.readCounter {
    28  			counter.Add(int64(buffer.Len()))
    29  		}
    30  	}
    31  	return
    32  }
    33  
    34  func (c *PacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
    35  	dataLen := int64(buffer.Len())
    36  	err := c.PacketConn.WritePacket(buffer, destination)
    37  	if err != nil {
    38  		return err
    39  	}
    40  	for _, counter := range c.writeCounter {
    41  		counter.Add(dataLen)
    42  	}
    43  	return nil
    44  }
    45  
    46  func (c *PacketConn) Upstream() any {
    47  	return c.PacketConn
    48  }
    49  
    50  type HookPacketConn struct {
    51  	N.PacketConn
    52  	readCounter  func(n int64)
    53  	writeCounter func(n int64)
    54  }
    55  
    56  func (c *HookPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
    57  	destination, err = c.PacketConn.ReadPacket(buffer)
    58  	if err == nil {
    59  		c.readCounter(int64(buffer.Len()))
    60  	}
    61  	return
    62  }
    63  
    64  func (c *HookPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
    65  	dataLen := int64(buffer.Len())
    66  	err := c.PacketConn.WritePacket(buffer, destination)
    67  	if err != nil {
    68  		return err
    69  	}
    70  	c.writeCounter(dataLen)
    71  	return nil
    72  }
    73  
    74  func (c *HookPacketConn) Upstream() any {
    75  	return c.PacketConn
    76  }