save

Author: nekohasekai

common/bufio/packet_reactor.go

@@ -15,7 +15,7 @@ import (
 )
 
 const (
-	batchReadTimeout = 250 * time.Millisecond
+	batchReadTimeout = 5 *time.Second
 )
 
 const (
@@ -141,12 +141,14 @@ type reactorStream struct {
 	destination  N.PacketWriter
 	originSource N.PacketReader
 
-	pushable      N.PacketPushable
-	pollable      N.PacketPollable
-	options       N.ReadWaitOptions
-	readWaiter    N.PacketReadWaiter
-	readCounters  []N.CountFunc
-	writeCounters []N.CountFunc
+	pushable         N.PacketPushable
+	pushableCallback func()
+	pollable         N.PacketPollable
+	deadlineSetter   deadlineSetter
+	options          N.ReadWaitOptions
+	readWaiter       N.PacketReadWaiter
+	readCounters     []N.CountFunc
+	writeCounters    []N.CountFunc
 
 	state atomic.Int32
 }
@@ -245,11 +247,27 @@ func (r *PacketReactor) registerStream(stream *reactorStream) {
 		return
 	}
 
+	// Check if deadline setter is available and functional
+	if setter, ok := stream.source.(deadlineSetter); ok {
+		err := setter.SetReadDeadline(time.Time{})
+		if err != nil {
+			r.logger.Trace("packet stream: SetReadDeadline not supported, using legacy copy")
+			go stream.runLegacyCopy()
+			return
+		}
+		stream.deadlineSetter = setter
+	} else {
+		r.logger.Trace("packet stream: no deadline setter, using legacy copy")
+		go stream.runLegacyCopy()
+		return
+	}
+
 	if stream.pushable != nil {
 		r.logger.Trace("packet stream: using pushable mode")
-		stream.pushable.SetOnDataReady(func() {
+		stream.pushableCallback = func() {
 			go stream.runActiveLoop(nil)
-		})
+		}
+		stream.pushable.SetOnDataReady(stream.pushableCallback)
 		if stream.pushable.HasPendingData() {
 			go stream.runActiveLoop(nil)
 		}
@@ -293,6 +311,10 @@ func (s *reactorStream) runActiveLoop(firstPacket *N.PacketBuffer) {
 		return
 	}
 
+	if s.pushable != nil {
+		s.pushable.SetOnDataReady(nil)
+	}
+
 	notFirstTime := false
 
 	if firstPacket != nil {
@@ -309,8 +331,12 @@ func (s *reactorStream) runActiveLoop(firstPacket *N.PacketBuffer) {
 			return
 		}
 
-		if setter, ok := s.source.(interface{ SetReadDeadline(time.Time) error }); ok {
-			setter.SetReadDeadline(time.Now().Add(batchReadTimeout))
+		deadlineErr := s.deadlineSetter.SetReadDeadline(time.Now().Add(batchReadTimeout))
+		if deadlineErr != nil {
+			s.connection.reactor.logger.Trace("packet stream: SetReadDeadline failed, switching to legacy copy")
+			s.state.Store(stateIdle)
+			go s.runLegacyCopy()
+			return
 		}
 
 		var (
@@ -341,16 +367,16 @@ func (s *reactorStream) runActiveLoop(firstPacket *N.PacketBuffer) {
 
 		if err != nil {
 			if E.IsTimeout(err) {
-				if setter, ok := s.source.(interface{ SetReadDeadline(time.Time) error }); ok {
-					setter.SetReadDeadline(time.Time{})
-				}
+				s.deadlineSetter.SetReadDeadline(time.Time{})
 				if !s.state.CompareAndSwap(stateActive, stateIdle) {
 					return
 				}
 				s.connection.reactor.logger.Trace("packet stream: timeout, returning to idle pool")
 				if s.pushable != nil {
+					s.pushable.SetOnDataReady(s.pushableCallback)
 					if s.pushable.HasPendingData() {
 						if s.state.CompareAndSwap(stateIdle, stateActive) {
+							s.pushable.SetOnDataReady(nil)
 							continue
 						}
 					}

common/bufio/stream_reactor.go

@@ -16,7 +16,7 @@ import (
 )
 
 const (
-	streamBatchReadTimeout = 250 * time.Millisecond
+	streamBatchReadTimeout = 5*time.Second
 )
 
 func CreateStreamPollable(reader io.Reader) (N.StreamPollable, bool) {
@@ -88,18 +88,23 @@ type streamConnection struct {
 	err       error
 }
 
+type deadlineSetter interface {
+	SetReadDeadline(time.Time) error
+}
+
 type streamDirection struct {
 	connection *streamConnection
 
 	source       io.Reader
 	destination  io.Writer
 	originSource net.Conn
 
-	pollable      N.StreamPollable
-	options       N.ReadWaitOptions
-	readWaiter    N.ReadWaiter
-	readCounters  []N.CountFunc
-	writeCounters []N.CountFunc
+	pollable       N.StreamPollable
+	deadlineSetter deadlineSetter
+	options        N.ReadWaitOptions
+	readWaiter     N.ReadWaiter
+	readCounters   []N.CountFunc
+	writeCounters  []N.CountFunc
 
 	isUpload bool
 	state    atomic.Int32
@@ -221,6 +226,21 @@ func (r *StreamReactor) registerDirection(direction *streamDirection) {
 		return
 	}
 
+	// Check if deadline setter is available and functional
+	if setter, ok := direction.originSource.(deadlineSetter); ok {
+		err := setter.SetReadDeadline(time.Time{})
+		if err != nil {
+			r.logger.Trace("stream direction: SetReadDeadline not supported, using legacy copy")
+			go direction.runLegacyCopy()
+			return
+		}
+		direction.deadlineSetter = setter
+	} else {
+		r.logger.Trace("stream direction: no deadline setter, using legacy copy")
+		go direction.runLegacyCopy()
+		return
+	}
+
 	// Check if there's buffered data that needs processing first
 	if direction.pollable != nil && direction.pollable.Buffered() > 0 {
 		r.logger.Trace("stream direction: has buffered data, starting active loop")
@@ -261,8 +281,12 @@ func (d *streamDirection) runActiveLoop() {
 		}
 
 		// Set batch read timeout
-		if setter, ok := d.originSource.(interface{ SetReadDeadline(time.Time) error }); ok {
-			setter.SetReadDeadline(time.Now().Add(streamBatchReadTimeout))
+		deadlineErr := d.deadlineSetter.SetReadDeadline(time.Now().Add(streamBatchReadTimeout))
+		if deadlineErr != nil {
+			d.connection.reactor.logger.Trace("stream direction: SetReadDeadline failed, switching to legacy copy")
+			d.state.Store(stateIdle)
+			go d.runLegacyCopy()
+			return
 		}
 
 		var (
@@ -284,9 +308,7 @@ func (d *streamDirection) runActiveLoop() {
 		if err != nil {
 			if E.IsTimeout(err) {
 				// Timeout: check buffer and return to pool
-				if setter, ok := d.originSource.(interface{ SetReadDeadline(time.Time) error }); ok {
-					setter.SetReadDeadline(time.Time{})
-				}
+				d.deadlineSetter.SetReadDeadline(time.Time{})
 				if d.state.CompareAndSwap(stateActive, stateIdle) {
 					d.connection.reactor.logger.Trace("stream direction: timeout, returning to idle pool")
 					d.returnToPool()

common/udpnat2/conn.go

@@ -40,31 +40,45 @@ type natConn struct {
 	queueMutex  sync.Mutex
 	onDataReady func()
 
+	deadlineMutex sync.Mutex
+	deadlineTimer *time.Timer
+	deadlineChan  chan struct{}
+	dataSignal    chan struct{}
+
 	closeOnce sync.Once
 	doneChan  chan struct{}
 }
 
-func (c *natConn) ReadPacket(buffer *buf.Buffer) (addr M.Socksaddr, err error) {
-	select {
-	case <-c.doneChan:
-		return M.Socksaddr{}, io.ErrClosedPipe
-	default:
-	}
+func (c *natConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
+	for {
+		select {
+		case <-c.doneChan:
+			return M.Socksaddr{}, io.ErrClosedPipe
+		default:
+		}
 
-	c.queueMutex.Lock()
-	if len(c.dataQueue) == 0 {
+		c.queueMutex.Lock()
+		if len(c.dataQueue) > 0 {
+			packet := c.dataQueue[0]
+			c.dataQueue = c.dataQueue[1:]
+			c.queueMutex.Unlock()
+			_, err = buffer.ReadOnceFrom(packet.Buffer)
+			destination = packet.Destination
+			packet.Buffer.Release()
+			N.PutPacketBuffer(packet)
+			return
+		}
 		c.queueMutex.Unlock()
-		return M.Socksaddr{}, os.ErrDeadlineExceeded
-	}
-	packet := c.dataQueue[0]
-	c.dataQueue = c.dataQueue[1:]
-	c.queueMutex.Unlock()
 
-	_, err = buffer.ReadOnceFrom(packet.Buffer)
-	destination := packet.Destination
-	packet.Buffer.Release()
-	N.PutPacketBuffer(packet)
-	return destination, err
+		select {
+		case <-c.doneChan:
+			return M.Socksaddr{}, io.ErrClosedPipe
+		case <-c.waitDeadline():
+			return M.Socksaddr{}, os.ErrDeadlineExceeded
+		case <-c.dataSignal:
+			continue
+		}
+	}
 }
 
 func (c *natConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
@@ -79,25 +93,34 @@ func (c *natConn) InitializeReadWaiter(options N.ReadWaitOptions) (needCopy bool
 }
 
 func (c *natConn) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
-	select {
-	case <-c.doneChan:
-		return nil, M.Socksaddr{}, io.ErrClosedPipe
-	default:
-	}
+	for {
+		select {
+		case <-c.doneChan:
+			return nil, M.Socksaddr{}, io.ErrClosedPipe
+		default:
+		}
 
-	c.queueMutex.Lock()
-	if len(c.dataQueue) == 0 {
+		c.queueMutex.Lock()
+		if len(c.dataQueue) > 0 {
+			packet := c.dataQueue[0]
+			c.dataQueue = c.dataQueue[1:]
+			c.queueMutex.Unlock()
+			buffer = c.readWaitOptions.Copy(packet.Buffer)
+			destination = packet.Destination
+			N.PutPacketBuffer(packet)
+			return
+		}
 		c.queueMutex.Unlock()
-		return nil, M.Socksaddr{}, os.ErrDeadlineExceeded
-	}
-	packet := c.dataQueue[0]
-	c.dataQueue = c.dataQueue[1:]
-	c.queueMutex.Unlock()
 
-	buffer = c.readWaitOptions.Copy(packet.Buffer)
-	destination = packet.Destination
-	N.PutPacketBuffer(packet)
-	return
+		select {
+		case <-c.doneChan:
+			return nil, M.Socksaddr{}, io.ErrClosedPipe
+		case <-c.waitDeadline():
+			return nil, M.Socksaddr{}, os.ErrDeadlineExceeded
+		case <-c.dataSignal:
+			continue
+		}
+	}
 }
 
 func (c *natConn) SetHandler(handler N.UDPHandlerEx) {
@@ -141,6 +164,11 @@ func (c *natConn) PushPacket(packet *N.PacketBuffer) {
 	callback := c.onDataReady
 	c.queueMutex.Unlock()
 
+	select {
+	case c.dataSignal <- struct{}{}:
+	default:
+	}
+
 	if callback != nil {
 		callback()
 	}
@@ -187,9 +215,52 @@ func (c *natConn) SetDeadline(t time.Time) error {
 }
 
 func (c *natConn) SetReadDeadline(t time.Time) error {
+	c.deadlineMutex.Lock()
+	defer c.deadlineMutex.Unlock()
+
+	if c.deadlineTimer != nil && !c.deadlineTimer.Stop() {
+		<-c.deadlineChan
+	}
+	c.deadlineTimer = nil
+
+	if t.IsZero() {
+		if isClosedChan(c.deadlineChan) {
+			c.deadlineChan = make(chan struct{})
+		}
+		return nil
+	}
+
+	if duration := time.Until(t); duration > 0 {
+		if isClosedChan(c.deadlineChan) {
+			c.deadlineChan = make(chan struct{})
+		}
+		c.deadlineTimer = time.AfterFunc(duration, func() {
+			close(c.deadlineChan)
+		})
+		return nil
+	}
+
+	if !isClosedChan(c.deadlineChan) {
+		close(c.deadlineChan)
+	}
 	return nil
 }
 
+func (c *natConn) waitDeadline() chan struct{} {
+	c.deadlineMutex.Lock()
+	defer c.deadlineMutex.Unlock()
+	return c.deadlineChan
+}
+
+func isClosedChan(channel <-chan struct{}) bool {
+	select {
+	case <-channel:
+		return true
+	default:
+		return false
+	}
+}
+
 func (c *natConn) SetWriteDeadline(t time.Time) error {
 	return os.ErrInvalid
 }

common/udpnat2/service.go

@@ -56,10 +56,12 @@ func (s *Service) NewPacket(bufferSlices [][]byte, source M.Socksaddr, destinati
 			return nil, false
 		}
 		newConn := &natConn{
-			cache:     s.cache,
-			writer:    writer,
-			localAddr: source,
-			doneChan:  make(chan struct{}),
+			cache:        s.cache,
+			writer:       writer,
+			localAddr:    source,
+			deadlineChan: make(chan struct{}),
+			dataSignal:   make(chan struct{}, 1),
+			doneChan:     make(chan struct{}),
 		}
 		go s.handler.NewPacketConnectionEx(ctx, newConn, source, destination, onClose)
 		return newConn, true