v2ray-core/common/buf/io.go

104 lines
2.1 KiB
Go

package buf
import (
"io"
"net"
"syscall"
"time"
)
// Reader extends io.Reader with MultiBuffer.
type Reader interface {
// ReadMultiBuffer reads content from underlying reader, and put it into a MultiBuffer.
ReadMultiBuffer() (MultiBuffer, error)
}
// ErrReadTimeout is an error that happens with IO timeout.
var ErrReadTimeout = newError("IO timeout")
// TimeoutReader is a reader that returns error if Read() operation takes longer than the given timeout.
type TimeoutReader interface {
ReadMultiBufferTimeout(time.Duration) (MultiBuffer, error)
}
// Writer extends io.Writer with MultiBuffer.
type Writer interface {
// WriteMultiBuffer writes a MultiBuffer into underlying writer.
WriteMultiBuffer(MultiBuffer) error
}
// WriteAllBytes ensures all bytes are written into the given writer.
func WriteAllBytes(writer io.Writer, payload []byte) error {
for len(payload) > 0 {
n, err := writer.Write(payload)
if err != nil {
return err
}
payload = payload[n:]
}
return nil
}
func isPacketReader(reader io.Reader) bool {
_, ok := reader.(net.PacketConn)
return ok
}
// NewReader creates a new Reader.
// The Reader instance doesn't take the ownership of reader.
func NewReader(reader io.Reader) Reader {
if mr, ok := reader.(Reader); ok {
return mr
}
if isPacketReader(reader) {
return &PacketReader{
Reader: reader,
}
}
if useReadv {
if sc, ok := reader.(syscall.Conn); ok {
rawConn, err := sc.SyscallConn()
if err != nil {
newError("failed to get sysconn").Base(err).WriteToLog()
} else {
return NewReadVReader(reader, rawConn)
}
}
}
return &SingleReader{
Reader: reader,
}
}
func isPacketWriter(writer io.Writer) bool {
if _, ok := writer.(net.PacketConn); ok {
return true
}
// If the writer doesn't implement syscall.Conn, it is probably not a TCP connection.
if _, ok := writer.(syscall.Conn); !ok {
return true
}
return false
}
// NewWriter creates a new Writer.
func NewWriter(writer io.Writer) Writer {
if mw, ok := writer.(Writer); ok {
return mw
}
if isPacketWriter(writer) {
return &SequentialWriter{
Writer: writer,
}
}
return &BufferToBytesWriter{
Writer: writer,
}
}