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

     1  package deadline
     2  
     3  import (
     4  	"net"
     5  	"os"
     6  	"runtime"
     7  	"time"
     8  
     9  	"github.com/metacubex/mihomo/common/atomic"
    10  	"github.com/metacubex/mihomo/common/net/packet"
    11  )
    12  
    13  type readResult struct {
    14  	data []byte
    15  	addr net.Addr
    16  	err  error
    17  }
    18  
    19  type NetPacketConn struct {
    20  	net.PacketConn
    21  	deadline     atomic.TypedValue[time.Time]
    22  	pipeDeadline pipeDeadline
    23  	disablePipe  atomic.Bool
    24  	inRead       atomic.Bool
    25  	resultCh     chan any
    26  }
    27  
    28  func NewNetPacketConn(pc net.PacketConn) net.PacketConn {
    29  	npc := &NetPacketConn{
    30  		PacketConn:   pc,
    31  		pipeDeadline: makePipeDeadline(),
    32  		resultCh:     make(chan any, 1),
    33  	}
    34  	npc.resultCh <- nil
    35  	if enhancePC, isEnhance := pc.(packet.EnhancePacketConn); isEnhance {
    36  		epc := &EnhancePacketConn{
    37  			NetPacketConn: npc,
    38  			enhancePacketConn: enhancePacketConn{
    39  				netPacketConn:     npc,
    40  				enhancePacketConn: enhancePC,
    41  			},
    42  		}
    43  		if singPC, isSingPC := pc.(packet.SingPacketConn); isSingPC {
    44  			return &EnhanceSingPacketConn{
    45  				EnhancePacketConn: epc,
    46  				singPacketConn: singPacketConn{
    47  					netPacketConn:  npc,
    48  					singPacketConn: singPC,
    49  				},
    50  			}
    51  		}
    52  		return epc
    53  	}
    54  	if singPC, isSingPC := pc.(packet.SingPacketConn); isSingPC {
    55  		return &SingPacketConn{
    56  			NetPacketConn: npc,
    57  			singPacketConn: singPacketConn{
    58  				netPacketConn:  npc,
    59  				singPacketConn: singPC,
    60  			},
    61  		}
    62  	}
    63  	return npc
    64  }
    65  
    66  func (c *NetPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    67  FOR:
    68  	for {
    69  		select {
    70  		case result := <-c.resultCh:
    71  			if result != nil {
    72  				if result, ok := result.(*readResult); ok {
    73  					n = copy(p, result.data)
    74  					addr = result.addr
    75  					err = result.err
    76  					c.resultCh <- nil // finish cache read
    77  					return
    78  				}
    79  				c.resultCh <- result // another type of read
    80  				runtime.Gosched()    // allowing other goroutines to run
    81  				continue FOR
    82  			} else {
    83  				c.resultCh <- nil
    84  				break FOR
    85  			}
    86  		case <-c.pipeDeadline.wait():
    87  			return 0, nil, os.ErrDeadlineExceeded
    88  		}
    89  	}
    90  
    91  	if c.disablePipe.Load() {
    92  		return c.PacketConn.ReadFrom(p)
    93  	} else if c.deadline.Load().IsZero() {
    94  		c.inRead.Store(true)
    95  		defer c.inRead.Store(false)
    96  		n, addr, err = c.PacketConn.ReadFrom(p)
    97  		return
    98  	}
    99  
   100  	<-c.resultCh
   101  	go c.pipeReadFrom(len(p))
   102  
   103  	return c.ReadFrom(p)
   104  }
   105  
   106  func (c *NetPacketConn) pipeReadFrom(size int) {
   107  	buffer := make([]byte, size)
   108  	n, addr, err := c.PacketConn.ReadFrom(buffer)
   109  	buffer = buffer[:n]
   110  	result := &readResult{}
   111  	result.data = buffer
   112  	result.addr = addr
   113  	result.err = err
   114  	c.resultCh <- result
   115  }
   116  
   117  func (c *NetPacketConn) SetReadDeadline(t time.Time) error {
   118  	if c.disablePipe.Load() {
   119  		return c.PacketConn.SetReadDeadline(t)
   120  	} else if c.inRead.Load() {
   121  		c.disablePipe.Store(true)
   122  		return c.PacketConn.SetReadDeadline(t)
   123  	}
   124  	c.deadline.Store(t)
   125  	c.pipeDeadline.set(t)
   126  	return nil
   127  }
   128  
   129  func (c *NetPacketConn) ReaderReplaceable() bool {
   130  	select {
   131  	case result := <-c.resultCh:
   132  		c.resultCh <- result
   133  		if result != nil {
   134  			return false // cache reading
   135  		} else {
   136  			break
   137  		}
   138  	default:
   139  		return false // pipe reading
   140  	}
   141  	return c.disablePipe.Load() || c.deadline.Load().IsZero()
   142  }
   143  
   144  func (c *NetPacketConn) WriterReplaceable() bool {
   145  	return true
   146  }
   147  
   148  func (c *NetPacketConn) Upstream() any {
   149  	return c.PacketConn
   150  }
   151  
   152  func (c *NetPacketConn) NeedAdditionalReadDeadline() bool {
   153  	return false
   154  }