ref.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. package packet
  2. import (
  3. "net"
  4. "runtime"
  5. "time"
  6. )
  7. type refPacketConn struct {
  8. pc EnhancePacketConn
  9. ref any
  10. }
  11. func (c *refPacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) {
  12. defer runtime.KeepAlive(c.ref)
  13. return c.pc.WaitReadFrom()
  14. }
  15. func (c *refPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
  16. defer runtime.KeepAlive(c.ref)
  17. return c.pc.ReadFrom(p)
  18. }
  19. func (c *refPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
  20. defer runtime.KeepAlive(c.ref)
  21. return c.pc.WriteTo(p, addr)
  22. }
  23. func (c *refPacketConn) Close() error {
  24. defer runtime.KeepAlive(c.ref)
  25. return c.pc.Close()
  26. }
  27. func (c *refPacketConn) LocalAddr() net.Addr {
  28. defer runtime.KeepAlive(c.ref)
  29. return c.pc.LocalAddr()
  30. }
  31. func (c *refPacketConn) SetDeadline(t time.Time) error {
  32. defer runtime.KeepAlive(c.ref)
  33. return c.pc.SetDeadline(t)
  34. }
  35. func (c *refPacketConn) SetReadDeadline(t time.Time) error {
  36. defer runtime.KeepAlive(c.ref)
  37. return c.pc.SetReadDeadline(t)
  38. }
  39. func (c *refPacketConn) SetWriteDeadline(t time.Time) error {
  40. defer runtime.KeepAlive(c.ref)
  41. return c.pc.SetWriteDeadline(t)
  42. }
  43. func (c *refPacketConn) Upstream() any {
  44. return c.pc
  45. }
  46. func (c *refPacketConn) ReaderReplaceable() bool { // Relay() will handle reference
  47. return true
  48. }
  49. func (c *refPacketConn) WriterReplaceable() bool { // Relay() will handle reference
  50. return true
  51. }
  52. func NewRefPacketConn(pc net.PacketConn, ref any) EnhancePacketConn {
  53. rPC := &refPacketConn{pc: NewEnhancePacketConn(pc), ref: ref}
  54. if singPC, isSingPC := pc.(SingPacketConn); isSingPC {
  55. return &refSingPacketConn{
  56. refPacketConn: rPC,
  57. singPacketConn: singPC,
  58. }
  59. }
  60. return rPC
  61. }