github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/components/statistics/conn.go (about)

     1  package statistics
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/protos/statistic"
     8  )
     9  
    10  type connection interface {
    11  	io.Closer
    12  	LocalAddr() net.Addr
    13  	Info() *statistic.Connection
    14  	ID() uint64
    15  }
    16  
    17  var _ connection = (*conn)(nil)
    18  
    19  type conn struct {
    20  	net.Conn
    21  
    22  	info    *statistic.Connection
    23  	manager *Connections
    24  }
    25  
    26  func (s *conn) Close() error {
    27  	s.manager.Remove(s.ID())
    28  	return s.Conn.Close()
    29  }
    30  
    31  func (s *conn) Write(b []byte) (_ int, err error) {
    32  	n, err := s.Conn.Write(b)
    33  	s.manager.Cache.AddUpload(uint64(n))
    34  	return int(n), err
    35  }
    36  
    37  func (s *conn) Read(b []byte) (n int, err error) {
    38  	n, err = s.Conn.Read(b)
    39  	s.manager.Cache.AddDownload(uint64(n))
    40  	return
    41  }
    42  
    43  func (s *conn) Info() *statistic.Connection { return s.info }
    44  func (s *conn) ID() uint64                  { return s.info.GetId() }
    45  
    46  var _ connection = (*packetConn)(nil)
    47  
    48  type packetConn struct {
    49  	net.PacketConn
    50  
    51  	info    *statistic.Connection
    52  	manager *Connections
    53  }
    54  
    55  func (s *packetConn) Info() *statistic.Connection { return s.info }
    56  func (s *packetConn) ID() uint64                  { return s.info.GetId() }
    57  
    58  func (s *packetConn) Close() error {
    59  	s.manager.Remove(s.ID())
    60  	return s.PacketConn.Close()
    61  }
    62  
    63  func (s *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
    64  	n, err = s.PacketConn.WriteTo(p, addr)
    65  	s.manager.Cache.AddUpload(uint64(n))
    66  	return
    67  }
    68  
    69  func (s *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    70  	n, addr, err = s.PacketConn.ReadFrom(p)
    71  	s.manager.Cache.AddDownload(uint64(n))
    72  	return
    73  }