1package buf 2 3import ( 4 "io" 5 "net" 6 "os" 7 "syscall" 8 "time" 9) 10 11// Reader extends io.Reader with MultiBuffer. 12type 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. 18var ErrReadTimeout = newError("IO timeout") 19 20// TimeoutReader is a reader that returns error if Read() operation takes longer than the given timeout. 21type TimeoutReader interface { 22 ReadMultiBufferTimeout(time.Duration) (MultiBuffer, error) 23} 24 25// Writer extends io.Writer with MultiBuffer. 26type 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. 32func 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 43func 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. 50func 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. 79func 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 89func 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. 102func 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} 117