github.com/v2fly/v2ray-core/v4@v4.45.2/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  	WriteMultiBuffer(MultiBuffer) error
    29  }
    30  
    31  // WriteAllBytes ensures all bytes are written into the given writer.
    32  func WriteAllBytes(writer io.Writer, payload []byte) error {
    33  	for len(payload) > 0 {
    34  		n, err := writer.Write(payload)
    35  		if err != nil {
    36  			return err
    37  		}
    38  		payload = payload[n:]
    39  	}
    40  	return nil
    41  }
    42  
    43  func isPacketReader(reader io.Reader) bool {
    44  	_, ok := reader.(net.PacketConn)
    45  	return ok
    46  }
    47  
    48  // NewReader creates a new Reader.
    49  // The Reader instance doesn't take the ownership of reader.
    50  func NewReader(reader io.Reader) Reader {
    51  	if mr, ok := reader.(Reader); ok {
    52  		return mr
    53  	}
    54  
    55  	if isPacketReader(reader) {
    56  		return &PacketReader{
    57  			Reader: reader,
    58  		}
    59  	}
    60  
    61  	_, isFile := reader.(*os.File)
    62  	if !isFile && useReadv {
    63  		if sc, ok := reader.(syscall.Conn); ok {
    64  			rawConn, err := sc.SyscallConn()
    65  			if err != nil {
    66  				newError("failed to get sysconn").Base(err).WriteToLog()
    67  			} else {
    68  				return NewReadVReader(reader, rawConn)
    69  			}
    70  		}
    71  	}
    72  
    73  	return &SingleReader{
    74  		Reader: reader,
    75  	}
    76  }
    77  
    78  // NewPacketReader creates a new PacketReader based on the given reader.
    79  func NewPacketReader(reader io.Reader) Reader {
    80  	if mr, ok := reader.(Reader); ok {
    81  		return mr
    82  	}
    83  
    84  	return &PacketReader{
    85  		Reader: reader,
    86  	}
    87  }
    88  
    89  func isPacketWriter(writer io.Writer) bool {
    90  	if _, ok := writer.(net.PacketConn); ok {
    91  		return true
    92  	}
    93  
    94  	// If the writer doesn't implement syscall.Conn, it is probably not a TCP connection.
    95  	if _, ok := writer.(syscall.Conn); !ok {
    96  		return true
    97  	}
    98  	return false
    99  }
   100  
   101  // NewWriter creates a new Writer.
   102  func NewWriter(writer io.Writer) Writer {
   103  	if mw, ok := writer.(Writer); ok {
   104  		return mw
   105  	}
   106  
   107  	if isPacketWriter(writer) {
   108  		return &SequentialWriter{
   109  			Writer: writer,
   110  		}
   111  	}
   112  
   113  	return &BufferToBytesWriter{
   114  		Writer: writer,
   115  	}
   116  }