github.com/metacubex/mihomo@v1.18.5/common/net/packet/packet.go (about)

     1  package packet
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/metacubex/mihomo/common/pool"
     7  )
     8  
     9  type WaitReadFrom interface {
    10  	WaitReadFrom() (data []byte, put func(), addr net.Addr, err error)
    11  }
    12  
    13  type EnhancePacketConn interface {
    14  	net.PacketConn
    15  	WaitReadFrom
    16  }
    17  
    18  func NewEnhancePacketConn(pc net.PacketConn) EnhancePacketConn {
    19  	if udpConn, isUDPConn := pc.(*net.UDPConn); isUDPConn {
    20  		return &enhanceUDPConn{UDPConn: udpConn}
    21  	}
    22  	if enhancePC, isEnhancePC := pc.(EnhancePacketConn); isEnhancePC {
    23  		return enhancePC
    24  	}
    25  	if singPC, isSingPC := pc.(SingPacketConn); isSingPC {
    26  		return newEnhanceSingPacketConn(singPC)
    27  	}
    28  	return &enhancePacketConn{PacketConn: pc}
    29  }
    30  
    31  type enhancePacketConn struct {
    32  	net.PacketConn
    33  }
    34  
    35  func (c *enhancePacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) {
    36  	return waitReadFrom(c.PacketConn)
    37  }
    38  
    39  func (c *enhancePacketConn) Upstream() any {
    40  	return c.PacketConn
    41  }
    42  
    43  func (c *enhancePacketConn) WriterReplaceable() bool {
    44  	return true
    45  }
    46  
    47  func (c *enhancePacketConn) ReaderReplaceable() bool {
    48  	return true
    49  }
    50  
    51  func (c *enhanceUDPConn) Upstream() any {
    52  	return c.UDPConn
    53  }
    54  
    55  func (c *enhanceUDPConn) WriterReplaceable() bool {
    56  	return true
    57  }
    58  
    59  func (c *enhanceUDPConn) ReaderReplaceable() bool {
    60  	return true
    61  }
    62  
    63  func waitReadFrom(pc net.PacketConn) (data []byte, put func(), addr net.Addr, err error) {
    64  	readBuf := pool.Get(pool.UDPBufferSize)
    65  	put = func() {
    66  		_ = pool.Put(readBuf)
    67  	}
    68  	var readN int
    69  	readN, addr, err = pc.ReadFrom(readBuf)
    70  	if readN > 0 {
    71  		data = readBuf[:readN]
    72  	} else {
    73  		put()
    74  		put = nil
    75  	}
    76  	return
    77  }