github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/buf/io.go (about)

     1  package buf
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"os"
     7  	"syscall"
     8  	"time"
     9  )
    10  
    11  // Reader extends io.Reader with MultiBuffer.
    12  type Reader interface {
    13  	// ReadMultiBuffer reads content from underlying reader, and put it into a MultiBuffer.
    14  	ReadMultiBuffer() (MultiBuffer, error)
    15  }
    16  
    17  // ErrReadTimeout is an error that happens with IO timeout.
    18  var ErrReadTimeout = newError("IO timeout")
    19  
    20  // TimeoutReader is a reader that returns error if Read() operation takes longer than the given timeout.
    21  type TimeoutReader interface {
    22  	ReadMultiBufferTimeout(time.Duration) (MultiBuffer, error)
    23  }
    24  
    25  // Writer extends io.Writer with MultiBuffer.
    26  type Writer interface {
    27  	// WriteMultiBuffer writes a MultiBuffer into underlying writer.
    28  	// Caller relinquish the ownership of MultiBuffer after calling this method.
    29  	WriteMultiBuffer(MultiBuffer) error
    30  }
    31  
    32  // WriteAllBytes ensures all bytes are written into the given writer.
    33  func WriteAllBytes(writer io.Writer, payload []byte) error {
    34  	for len(payload) > 0 {
    35  		n, err := writer.Write(payload)
    36  		if err != nil {
    37  			return err
    38  		}
    39  		payload = payload[n:]
    40  	}
    41  	return nil
    42  }
    43  
    44  func isPacketReader(reader io.Reader) bool {
    45  	_, ok := reader.(net.PacketConn)
    46  	return ok
    47  }
    48  
    49  // NewReader creates a new Reader.
    50  // The Reader instance doesn't take the ownership of reader.
    51  func NewReader(reader io.Reader) Reader {
    52  	if mr, ok := reader.(Reader); ok {
    53  		return mr
    54  	}
    55  
    56  	if isPacketReader(reader) {
    57  		return &PacketReader{
    58  			Reader: reader,
    59  		}
    60  	}
    61  
    62  	_, isFile := reader.(*os.File)
    63  	if !isFile && useReadv {
    64  		if sc, ok := reader.(syscall.Conn); ok {
    65  			rawConn, err := sc.SyscallConn()
    66  			if err != nil {
    67  				newError("failed to get sysconn").Base(err).WriteToLog()
    68  			} else {
    69  				return NewReadVReader(reader, rawConn)
    70  			}
    71  		}
    72  	}
    73  
    74  	return &SingleReader{
    75  		Reader: reader,
    76  	}
    77  }
    78  
    79  // NewPacketReader creates a new PacketReader based on the given reader.
    80  func NewPacketReader(reader io.Reader) Reader {
    81  	if mr, ok := reader.(Reader); ok {
    82  		return mr
    83  	}
    84  
    85  	return &PacketReader{
    86  		Reader: reader,
    87  	}
    88  }
    89  
    90  func isPacketWriter(writer io.Writer) bool {
    91  	if _, ok := writer.(net.PacketConn); ok {
    92  		return true
    93  	}
    94  
    95  	// If the writer doesn't implement syscall.Conn, it is probably not a TCP connection.
    96  	if _, ok := writer.(syscall.Conn); !ok {
    97  		return true
    98  	}
    99  	return false
   100  }
   101  
   102  // NewWriter creates a new Writer.
   103  func NewWriter(writer io.Writer) Writer {
   104  	if mw, ok := writer.(Writer); ok {
   105  		return mw
   106  	}
   107  
   108  	if isPacketWriter(writer) {
   109  		return &SequentialWriter{
   110  			Writer: writer,
   111  		}
   112  	}
   113  
   114  	return &BufferToBytesWriter{
   115  		Writer: writer,
   116  	}
   117  }