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

     1  package netapi
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  	"net"
     7  	"runtime"
     8  	"time"
     9  
    10  	"github.com/Asutorufa/yuhaiin/pkg/log"
    11  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
    12  )
    13  
    14  type multipleReaderTCPConn struct {
    15  	*net.TCPConn
    16  	mr io.Reader
    17  }
    18  
    19  func (m *multipleReaderTCPConn) Read(b []byte) (int, error) {
    20  	return m.mr.Read(b)
    21  }
    22  
    23  type multipleReaderConn struct {
    24  	net.Conn
    25  	mr io.Reader
    26  }
    27  
    28  func NewMultipleReaderConn(c net.Conn, r io.Reader) net.Conn {
    29  	tc, ok := c.(*net.TCPConn)
    30  	if ok {
    31  		return &multipleReaderTCPConn{tc, r}
    32  	}
    33  
    34  	return &multipleReaderConn{c, r}
    35  }
    36  
    37  func (m *multipleReaderConn) Read(b []byte) (int, error) {
    38  	return m.mr.Read(b)
    39  }
    40  
    41  type prefixBytesConn struct {
    42  	buffers pool.MultipleBytes
    43  	net.Conn
    44  }
    45  
    46  func (p *prefixBytesConn) Close() error {
    47  	err := p.Conn.Close()
    48  	p.buffers.Free()
    49  	return err
    50  }
    51  
    52  func NewPrefixBytesConn(c net.Conn, prefix ...*pool.Bytes) net.Conn {
    53  	if len(prefix) == 0 {
    54  		return c
    55  	}
    56  
    57  	buf := net.Buffers(nil)
    58  
    59  	for _, v := range prefix {
    60  		buf = append(buf, v.Bytes())
    61  	}
    62  
    63  	conn := NewMultipleReaderConn(c, io.MultiReader(&buf, c))
    64  
    65  	return &prefixBytesConn{
    66  		buffers: prefix,
    67  		Conn:    conn,
    68  	}
    69  }
    70  
    71  func MergeBufioReaderConn(c net.Conn, r *bufio.Reader) (net.Conn, error) {
    72  	if r.Buffered() <= 0 {
    73  		return c, nil
    74  	}
    75  
    76  	data, err := r.Peek(r.Buffered())
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  
    81  	return NewPrefixBytesConn(c, pool.GetBytesBuffer(len(data)).Copy(data)), nil
    82  }
    83  
    84  type LogConn struct {
    85  	net.Conn
    86  }
    87  
    88  func (l *LogConn) Write(b []byte) (int, error) {
    89  	return l.Conn.Write(b)
    90  }
    91  
    92  func (l *LogConn) Read(b []byte) (int, error) {
    93  	n, err := l.Conn.Read(b)
    94  	if err != nil {
    95  		log.Error("tls read failed", "err", err)
    96  	}
    97  
    98  	return n, err
    99  }
   100  
   101  func (l *LogConn) SetDeadline(t time.Time) error {
   102  	_, file, line, _ := runtime.Caller(3)
   103  	log.Info("set deadline", "time", t, "line", line, "file", file, "time", t)
   104  
   105  	return l.Conn.SetDeadline(t)
   106  }
   107  
   108  func (l *LogConn) SetReadDeadline(t time.Time) error {
   109  	_, file, line, _ := runtime.Caller(3)
   110  	log.Info("set read deadline", "time", t, "line", line, "file", file, "time", t)
   111  
   112  	return l.Conn.SetReadDeadline(t)
   113  }
   114  
   115  func (l *LogConn) SetWriteDeadline(t time.Time) error {
   116  	_, file, line, _ := runtime.Caller(4)
   117  	log.Info("set write deadline", "time", t, "line", line, "file", file, "time", t)
   118  
   119  	return l.Conn.SetWriteDeadline(t)
   120  }
   121  
   122  func ReadFrom(pc net.PacketConn) (*pool.Bytes, Address, error) {
   123  	b := pool.GetBytesBuffer(pool.MaxSegmentSize)
   124  	_, saddr, err := b.ReadFromPacket(pc)
   125  	if err != nil {
   126  		b.Free()
   127  		return nil, nil, err
   128  	}
   129  
   130  	addr, err := ParseSysAddr(saddr)
   131  	if err != nil {
   132  		b.Free()
   133  		return nil, nil, err
   134  	}
   135  
   136  	return b, addr, nil
   137  }