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

     1  package deadline
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"time"
     7  
     8  	"github.com/sagernet/sing/common/atomic"
     9  	"github.com/sagernet/sing/common/buf"
    10  	"github.com/sagernet/sing/common/bufio"
    11  	N "github.com/sagernet/sing/common/network"
    12  )
    13  
    14  type TimeoutReader interface {
    15  	io.Reader
    16  	SetReadDeadline(t time.Time) error
    17  }
    18  
    19  type Reader interface {
    20  	N.ExtendedReader
    21  	TimeoutReader
    22  	N.WithUpstreamReader
    23  	N.ReaderWithUpstream
    24  }
    25  
    26  type reader struct {
    27  	N.ExtendedReader
    28  	timeoutReader TimeoutReader
    29  	deadline      atomic.TypedValue[time.Time]
    30  	pipeDeadline  pipeDeadline
    31  	result        chan *readResult
    32  	done          chan struct{}
    33  }
    34  
    35  type readResult struct {
    36  	buffer *buf.Buffer
    37  	err    error
    38  }
    39  
    40  func NewReader(timeoutReader TimeoutReader) Reader {
    41  	return &reader{
    42  		ExtendedReader: bufio.NewExtendedReader(timeoutReader),
    43  		timeoutReader:  timeoutReader,
    44  		pipeDeadline:   makePipeDeadline(),
    45  		result:         make(chan *readResult, 1),
    46  		done:           makeFilledChan(),
    47  	}
    48  }
    49  
    50  func (r *reader) Read(p []byte) (n int, err error) {
    51  	select {
    52  	case result := <-r.result:
    53  		return r.pipeReturn(result, p)
    54  	default:
    55  	}
    56  	select {
    57  	case <-r.done:
    58  		go r.pipeRead(len(p))
    59  	default:
    60  	}
    61  	return r.read(p)
    62  }
    63  
    64  func (r *reader) read(p []byte) (n int, err error) {
    65  	select {
    66  	case result := <-r.result:
    67  		return r.pipeReturn(result, p)
    68  	case <-r.pipeDeadline.wait():
    69  		return 0, os.ErrDeadlineExceeded
    70  	}
    71  }
    72  
    73  func (r *reader) pipeReturn(result *readResult, p []byte) (n int, err error) {
    74  	n = copy(p, result.buffer.Bytes())
    75  	result.buffer.Advance(n)
    76  	if result.buffer.IsEmpty() {
    77  		result.buffer.Release()
    78  		err = result.err
    79  	} else {
    80  		r.result <- result
    81  	}
    82  	return
    83  }
    84  
    85  func (r *reader) pipeRead(pLen int) {
    86  	buffer := buf.NewSize(pLen)
    87  	_, err := buffer.ReadOnceFrom(r.ExtendedReader)
    88  	r.result <- &readResult{
    89  		buffer: buffer,
    90  		err:    err,
    91  	}
    92  	r.done <- struct{}{}
    93  }
    94  
    95  func (r *reader) ReadBuffer(buffer *buf.Buffer) error {
    96  	select {
    97  	case result := <-r.result:
    98  		return r.pipeReturnBuffer(result, buffer)
    99  	default:
   100  	}
   101  	select {
   102  	case <-r.done:
   103  		go r.pipeReadBuffer(buffer.FreeLen())
   104  	default:
   105  	}
   106  	return r.readBuffer(buffer)
   107  }
   108  
   109  func (r *reader) readBuffer(buffer *buf.Buffer) error {
   110  	select {
   111  	case result := <-r.result:
   112  		return r.pipeReturnBuffer(result, buffer)
   113  	case <-r.pipeDeadline.wait():
   114  		return os.ErrDeadlineExceeded
   115  	}
   116  }
   117  
   118  func (r *reader) pipeReturnBuffer(result *readResult, buffer *buf.Buffer) error {
   119  	n, _ := buffer.Write(result.buffer.Bytes())
   120  	result.buffer.Advance(n)
   121  	if !result.buffer.IsEmpty() {
   122  		r.result <- result
   123  		return nil
   124  	} else {
   125  		result.buffer.Release()
   126  		return result.err
   127  	}
   128  }
   129  
   130  func (r *reader) pipeReadBuffer(pLen int) {
   131  	cacheBuffer := buf.NewSize(pLen)
   132  	err := r.ExtendedReader.ReadBuffer(cacheBuffer)
   133  	r.result <- &readResult{
   134  		buffer: cacheBuffer,
   135  		err:    err,
   136  	}
   137  	r.done <- struct{}{}
   138  }
   139  
   140  func (r *reader) SetReadDeadline(t time.Time) error {
   141  	r.deadline.Store(t)
   142  	r.pipeDeadline.set(t)
   143  	return nil
   144  }
   145  
   146  func (r *reader) ReaderReplaceable() bool {
   147  	select {
   148  	case <-r.done:
   149  		r.done <- struct{}{}
   150  	default:
   151  		return false
   152  	}
   153  	select {
   154  	case result := <-r.result:
   155  		r.result <- result
   156  		return false
   157  	default:
   158  	}
   159  	return r.deadline.Load().IsZero()
   160  }
   161  
   162  func (r *reader) UpstreamReader() any {
   163  	return r.ExtendedReader
   164  }