Hunter0x7c7
2022-08-11 a82f9cb69f63aaeba40c024960deda7d75b9fece
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
package buf
 
import (
    "io"
    "net"
    "os"
    "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,
        }
    }
 
    _, isFile := reader.(*os.File)
    if !isFile && 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,
    }
}
 
// NewPacketReader creates a new PacketReader based on the given reader.
func NewPacketReader(reader io.Reader) Reader {
    if mr, ok := reader.(Reader); ok {
        return mr
    }
 
    return &PacketReader{
        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,
    }
}