-
Notifications
You must be signed in to change notification settings - Fork 8
/
conn.go
229 lines (207 loc) · 5.17 KB
/
conn.go
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
package rwc
import (
"context"
"io"
"net"
"os"
"sync"
"time"
)
// connPktSize is the size of the buffers to use for packets for the Conn.
const connPktSize = 2048
// Conn implements a Conn with a buffered ReadWriteCloser.
type Conn struct {
// ctx is the context
ctx context.Context
// ctxCancel is the context canceler
ctxCancel context.CancelFunc
// rwc is the read-write-closer
rwc io.ReadWriteCloser
// laddr is the local addr
laddr net.Addr
// raddr is the remote addr
raddr net.Addr
ar sync.Pool // packet arena
rd time.Time // read deadline
wd time.Time // write deadline
packetCh chan []byte // packet ch
closeErr error
}
// NewConn constructs a new packet conn and starts the rx pump.
func NewConn(
ctx context.Context,
rwc io.ReadWriteCloser,
laddr, raddr net.Addr,
bufferPacketN int,
) *Conn {
ctx, ctxCancel := context.WithCancel(ctx)
if bufferPacketN <= 0 {
bufferPacketN = 10
}
c := &Conn{
ctx: ctx,
ctxCancel: ctxCancel,
rwc: rwc,
laddr: laddr,
raddr: raddr,
packetCh: make(chan []byte, bufferPacketN),
}
go func() {
_ = c.rxPump()
}()
return c
}
// LocalAddr returns the local network address.
func (p *Conn) LocalAddr() net.Addr {
return p.laddr
}
// RemoteAddr returns the bound remote network address.
func (p *Conn) RemoteAddr() net.Addr {
return p.raddr
}
// Read reads data from the connection.
// Read can be made to time out and return an error after a fixed
// time limit; see SetDeadline and SetReadDeadline.
func (p *Conn) Read(b []byte) (n int, err error) {
deadline := p.rd
ctx := p.ctx
if !deadline.IsZero() {
var ctxCancel context.CancelFunc
ctx, ctxCancel = context.WithDeadline(ctx, deadline)
defer ctxCancel()
}
var pkt []byte
var ok bool
select {
case <-ctx.Done():
if !deadline.IsZero() {
return 0, os.ErrDeadlineExceeded
}
return 0, context.Canceled
case pkt, ok = <-p.packetCh:
if !ok {
err = p.closeErr
if err == nil {
err = io.EOF
}
return 0, err
}
}
pl := len(pkt)
copy(b, pkt)
p.ar.Put(&pkt)
if len(b) < pl {
return len(b), io.ErrShortBuffer
}
return pl, nil
}
// Write writes data to the connection.
func (p *Conn) Write(pkt []byte) (n int, err error) {
if len(pkt) == 0 {
return 0, nil
}
written := 0
for written < len(pkt) {
n, err = p.rwc.Write(pkt[written:])
written += n
if err != nil {
return written, err
}
}
return written, nil
}
// SetDeadline sets the read and write deadlines associated
// with the connection. It is equivalent to calling both
// SetReadDeadline and SetWriteDeadline.
//
// A deadline is an absolute time after which I/O operations
// fail instead of blocking. The deadline applies to all future
// and pending I/O, not just the immediately following call to
// Read or Write. After a deadline has been exceeded, the
// connection can be refreshed by setting a deadline in the future.
//
// If the deadline is exceeded a call to Read or Write or to other
// I/O methods will return an error that wraps os.ErrDeadlineExceeded.
// This can be tested using errors.Is(err, os.ErrDeadlineExceeded).
// The error's Timeout method will return true, but note that there
// are other possible errors for which the Timeout method will
// return true even if the deadline has not been exceeded.
//
// An idle timeout can be implemented by repeatedly extending
// the deadline after successful ReadFrom or WriteTo calls.
//
// A zero value for t means I/O operations will not time out.
func (p *Conn) SetDeadline(t time.Time) error {
p.rd = t
p.wd = t
return nil
}
// SetReadDeadline sets the deadline for future ReadFrom calls
// and any currently-blocked ReadFrom call.
// A zero value for t means ReadFrom will not time out.
func (p *Conn) SetReadDeadline(t time.Time) error {
p.rd = t
return nil
}
// SetWriteDeadline sets the deadline for future WriteTo calls
// and any currently-blocked WriteTo call.
// Even if write times out, it may return n > 0, indicating that
// some of the data was successfully written.
// A zero value for t means WriteTo will not time out.
func (p *Conn) SetWriteDeadline(t time.Time) error {
p.wd = t
return nil
}
// Close closes the connection.
// Any blocked ReadFrom or WriteTo operations will be unblocked and return errors.
func (p *Conn) Close() error {
return p.rwc.Close()
}
// getArenaBuf returns a buf from the packet arena with at least the given size
func (p *Conn) getArenaBuf(size int) []byte {
var buf []byte
bufp := p.ar.Get()
if bufp != nil {
buf = *bufp.(*[]byte)
}
if size != 0 {
if cap(buf) < size {
buf = make([]byte, size)
} else {
buf = buf[:size]
}
} else {
buf = buf[:cap(buf)]
}
return buf
}
// rxPump receives messages from the underlying connection.
func (p *Conn) rxPump() (rerr error) {
defer func() {
p.closeErr = rerr
close(p.packetCh)
}()
for {
select {
case <-p.ctx.Done():
return p.ctx.Err()
default:
}
pktBuf := p.getArenaBuf(int(connPktSize))
n, err := p.rwc.Read(pktBuf)
if n == 0 {
p.ar.Put(&pktBuf)
} else {
select {
case <-p.ctx.Done():
return context.Canceled
case p.packetCh <- pktBuf[:n]:
}
}
if err != nil {
return err
}
}
}
// _ is a type assertion
var _ net.Conn = ((*Conn)(nil))