github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/stats/io.go (about)

     1  package stats
     2  
     3  import (
     4  	"net"
     5  	"runtime"
     6  )
     7  
     8  type StatsConn struct {
     9  	net.Conn
    10  	CounterDown *Counter
    11  	CounterUp   *Counter
    12  }
    13  
    14  func (s StatsConn) Read(b []byte) (n int, err error) {
    15  	n, err = s.Conn.Read(b)
    16  	s.CounterDown.Add(int64(n))
    17  	return
    18  }
    19  
    20  func (s StatsConn) Write(b []byte) (n int, err error) {
    21  	n, err = s.Conn.Write(b)
    22  	s.CounterUp.Add(int64(n))
    23  	return
    24  }
    25  
    26  func NewStatsConn(c net.Conn) StatsConn {
    27  	return StatsConn{
    28  		Conn:        c,
    29  		CounterDown: DefaultManager.GetCounter(DownProxy),
    30  		CounterUp:   DefaultManager.GetCounter(UpProxy),
    31  	}
    32  }
    33  
    34  func NewConn(c net.Conn) net.Conn {
    35  	if runtime.GOOS != "android" {
    36  		return c
    37  	}
    38  	return NewStatsConn(c)
    39  }
    40  
    41  type StatsPacketConn struct {
    42  	net.PacketConn
    43  	CounterDown *Counter
    44  	CounterUp   *Counter
    45  }
    46  
    47  func (pc StatsPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    48  	n, addr, err = pc.PacketConn.ReadFrom(p)
    49  	pc.CounterDown.Add(int64(n))
    50  	return
    51  }
    52  
    53  func (pc StatsPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
    54  	n, err = pc.PacketConn.WriteTo(p, addr)
    55  	pc.CounterUp.Add(int64(n))
    56  	return
    57  }
    58  
    59  func NewStatsPacketConn(pc net.PacketConn) StatsPacketConn {
    60  	return StatsPacketConn{
    61  		PacketConn:  pc,
    62  		CounterDown: DefaultManager.GetCounter(DownProxy),
    63  		CounterUp:   DefaultManager.GetCounter(UpProxy),
    64  	}
    65  }