1
/*
2
 * Taken from https://github.com/grpc/grpc-go/blob/master/test/bufconn/bufconn.go
3
 *
4
 * Copyright 2017 gRPC authors.
5
 *
6
 * Licensed under the Apache License, Version 2.0 (the "License");
7
 * you may not use this file except in compliance with the License.
8
 * You may obtain a copy of the License at
9
 *
10
 *     http://www.apache.org/licenses/LICENSE-2.0
11
 *
12
 * Unless required by applicable law or agreed to in writing, software
13
 * distributed under the License is distributed on an "AS IS" BASIS,
14
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
 * See the License for the specific language governing permissions and
16
 * limitations under the License.
17
 *
18
 */
19
 
20
// Package bufconn provides a net.Conn implemented by a buffer and related
21
// dialing and listening functionality.
22
package bufconn
23
 
24
import (
25
	"fmt"
26
	"io"
27
	"net"
28
	"sync"
29
	"time"
30
)
31
 
32
// Implementation of net.Error providing timeout
33
type netErrorTimeout struct {
34
	error
35
}
36
 
37
func (e netErrorTimeout) Timeout() bool   { return true }
38
func (e netErrorTimeout) Temporary() bool { return false }
39
 
40
var errTimeout net.Error = netErrorTimeout{error: fmt.Errorf("i/o timeout")}
41
 
42
func BufferedPipe(bufSize int) (net.Conn, net.Conn) {
43
	p1, p2 := newPipe(bufSize), newPipe(bufSize)
44
	return &conn{p1, p2}, &conn{p2, p1}
45
}
46
 
47
type pipe struct {
48
	mu sync.Mutex
49
 
50
	// buf contains the data in the pipe.  It is a ring buffer of fixed capacity,
51
	// with r and w pointing to the offset to read and write, respectively.
52
	//
53
	// Data is read between [r, w) and written to [w, r), wrapping around the end
54
	// of the slice if necessary.
55
	//
56
	// The buffer is empty if r == len(buf), otherwise if r == w, it is full.
57
	//
58
	// w and r are always in the range [0, cap(buf)) and [0, len(buf)].
59
	buf  []byte
60
	w, r int
61
 
62
	wwait sync.Cond
63
	rwait sync.Cond
64
 
65
	// Indicate that a write/read timeout has occurred
66
	wtimedout bool
67
	rtimedout bool
68
 
69
	wtimer *time.Timer
70
	rtimer *time.Timer
71
	wseq   uint64
72
	rseq   uint64
73
 
74
	closed      bool
75
	writeClosed bool
76
}
77
 
78
func newPipe(sz int) *pipe {
79
	p := &pipe{buf: make([]byte, 0, sz)}
80
	p.wwait.L = &p.mu
81
	p.rwait.L = &p.mu
82
 
83
	p.wtimer = time.AfterFunc(0, func() {})
84
	p.rtimer = time.AfterFunc(0, func() {})
85
	return p
86
}
87
 
88
func (p *pipe) empty() bool {
89
	return p.r == len(p.buf)
90
}
91
 
92
func (p *pipe) full() bool {
93
	return p.r < len(p.buf) && p.r == p.w
94
}
95
 
96
func (p *pipe) Read(b []byte) (n int, err error) {
97
	p.mu.Lock()
98
	defer p.mu.Unlock()
99
	// Block until p has data.
100
	for {
101
		if p.closed {
102
			return 0, io.ErrClosedPipe
103
		}
104
		if !p.empty() {
105
			break
106
		}
107
		if p.writeClosed {
108
			return 0, io.EOF
109
		}
110
		if p.rtimedout {
111
			return 0, errTimeout
112
		}
113
 
114
		p.rwait.Wait()
115
	}
116
	wasFull := p.full()
117
 
118
	n = copy(b, p.buf[p.r:len(p.buf)])
119
	p.r += n
120
	if p.r == cap(p.buf) {
121
		p.r = 0
122
		p.buf = p.buf[:p.w]
123
	}
124
 
125
	// Signal a blocked writer, if any
126
	if wasFull {
127
		p.wwait.Signal()
128
	}
129
 
130
	return n, nil
131
}
132
 
133
func (p *pipe) Write(b []byte) (n int, err error) {
134
	p.mu.Lock()
135
	defer p.mu.Unlock()
136
	if p.closed {
137
		return 0, io.ErrClosedPipe
138
	}
139
	for len(b) > 0 {
140
		// Block until p is not full.
141
		for {
142
			if p.closed || p.writeClosed {
143
				return 0, io.ErrClosedPipe
144
			}
145
			if !p.full() {
146
				break
147
			}
148
			if p.wtimedout {
149
				return 0, errTimeout
150
			}
151
 
152
			p.wwait.Wait()
153
		}
154
		wasEmpty := p.empty()
155
 
156
		end := cap(p.buf)
157
		if p.w < p.r {
158
			end = p.r
159
		}
160
		x := copy(p.buf[p.w:end], b)
161
		b = b[x:]
162
		n += x
163
		p.w += x
164
		if p.w > len(p.buf) {
165
			p.buf = p.buf[:p.w]
166
		}
167
		if p.w == cap(p.buf) {
168
			p.w = 0
169
		}
170
 
171
		// Signal a blocked reader, if any.
172
		if wasEmpty {
173
			p.rwait.Signal()
174
		}
175
	}
176
	return n, nil
177
}
178
 
179
func (p *pipe) Close() error {
180
	p.mu.Lock()
181
	defer p.mu.Unlock()
182
	p.closed = true
183
	// Signal all blocked readers and writers to return an error.
184
	p.rwait.Broadcast()
185
	p.wwait.Broadcast()
186
	return nil
187
}
188
 
189
func (p *pipe) closeWrite() error {
190
	p.mu.Lock()
191
	defer p.mu.Unlock()
192
	p.writeClosed = true
193
	// Signal all blocked readers and writers to return an error.
194
	p.rwait.Broadcast()
195
	p.wwait.Broadcast()
196
	return nil
197
}
198
 
199
type conn struct {
200
	io.Reader
201
	io.Writer
202
}
203
 
204
func (c *conn) Close() error {
205
	err1 := c.Reader.(*pipe).Close()
206
	err2 := c.Writer.(*pipe).closeWrite()
207
	if err1 != nil {
208
		return err1
209
	}
210
	return err2
211
}
212
 
213
func (c *conn) SetDeadline(t time.Time) error {
214
	c.SetReadDeadline(t)
215
	c.SetWriteDeadline(t)
216
	return nil
217
}
218
 
219
func (c *conn) SetReadDeadline(t time.Time) error {
220
	p := c.Reader.(*pipe)
221
	p.mu.Lock()
222
	defer p.mu.Unlock()
223
	p.rtimer.Stop()
224
	p.rtimedout = false
225
	p.rseq++
226
	seq := p.rseq
227
	if !t.IsZero() {
228
		p.rtimer = time.AfterFunc(time.Until(t), func() {
229
			p.mu.Lock()
230
			defer p.mu.Unlock()
231
			// Stop cannot prevent a callback that is already running.
232
			if seq != p.rseq {
233
				return
234
			}
235
			p.rtimedout = true
236
			p.rwait.Broadcast()
237
		})
238
	}
239
	return nil
240
}
241
 
242
func (c *conn) SetWriteDeadline(t time.Time) error {
243
	p := c.Writer.(*pipe)
244
	p.mu.Lock()
245
	defer p.mu.Unlock()
246
	p.wtimer.Stop()
247
	p.wtimedout = false
248
	p.wseq++
249
	seq := p.wseq
250
	if !t.IsZero() {
251
		p.wtimer = time.AfterFunc(time.Until(t), func() {
252
			p.mu.Lock()
253
			defer p.mu.Unlock()
254
			if seq != p.wseq {
255
				return
256
			}
257
			p.wtimedout = true
258
			p.wwait.Broadcast()
259
		})
260
	}
261
	return nil
262
}
263
 
264
func (*conn) LocalAddr() net.Addr  { return addr{} }
265
func (*conn) RemoteAddr() net.Addr { return addr{} }
266
 
267
type addr struct{}
268
 
269
func (addr) Network() string { return "bufconn" }
270
func (addr) String() string  { return "bufconn" }