github.com/sagernet/sing@v0.2.6/common/bufio/deadline/packet_reader_fallback.go (about)

     1  package deadline
     2  
     3  import (
     4  	"net"
     5  	"time"
     6  
     7  	"github.com/sagernet/sing/common/atomic"
     8  	"github.com/sagernet/sing/common/buf"
     9  	M "github.com/sagernet/sing/common/metadata"
    10  )
    11  
    12  type fallbackPacketReader struct {
    13  	*packetReader
    14  	disablePipe atomic.Bool
    15  	inRead      atomic.Bool
    16  }
    17  
    18  func NewFallbackPacketReader(timeoutReader TimeoutPacketReader) PacketReader {
    19  	return &fallbackPacketReader{packetReader: NewPacketReader(timeoutReader).(*packetReader)}
    20  }
    21  
    22  func (r *fallbackPacketReader) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    23  	select {
    24  	case result := <-r.result:
    25  		return r.pipeReturnFrom(result, p)
    26  	default:
    27  	}
    28  	if r.disablePipe.Load() {
    29  		return r.TimeoutPacketReader.ReadFrom(p)
    30  	}
    31  	select {
    32  	case <-r.done:
    33  		if r.deadline.Load().IsZero() {
    34  			r.done <- struct{}{}
    35  			r.inRead.Store(true)
    36  			defer r.inRead.Store(false)
    37  			n, addr, err = r.TimeoutPacketReader.ReadFrom(p)
    38  			return
    39  		}
    40  		go r.pipeReadFrom(len(p))
    41  	default:
    42  	}
    43  	return r.readFrom(p)
    44  }
    45  
    46  func (r *fallbackPacketReader) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
    47  	select {
    48  	case result := <-r.result:
    49  		return r.pipeReturnFromBuffer(result, buffer)
    50  	default:
    51  	}
    52  	if r.disablePipe.Load() {
    53  		return r.TimeoutPacketReader.ReadPacket(buffer)
    54  	}
    55  	select {
    56  	case <-r.done:
    57  		if r.deadline.Load().IsZero() {
    58  			r.done <- struct{}{}
    59  			r.inRead.Store(true)
    60  			defer r.inRead.Store(false)
    61  			destination, err = r.TimeoutPacketReader.ReadPacket(buffer)
    62  			return
    63  		}
    64  		go r.pipeReadFromBuffer(buffer.FreeLen())
    65  	default:
    66  	}
    67  	return r.readPacket(buffer)
    68  }
    69  
    70  func (r *fallbackPacketReader) SetReadDeadline(t time.Time) error {
    71  	if r.disablePipe.Load() {
    72  		return r.TimeoutPacketReader.SetReadDeadline(t)
    73  	} else if r.inRead.Load() {
    74  		r.disablePipe.Store(true)
    75  		return r.TimeoutPacketReader.SetReadDeadline(t)
    76  	}
    77  	return r.packetReader.SetReadDeadline(t)
    78  }
    79  
    80  func (r *fallbackPacketReader) ReaderReplaceable() bool {
    81  	return r.disablePipe.Load() || r.packetReader.ReaderReplaceable()
    82  }
    83  
    84  func (r *fallbackPacketReader) UpstreamReader() any {
    85  	return r.packetReader.UpstreamReader()
    86  }