packet.go 1.5 KB

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