8000 chore: move InProcNet to testutil · coder/coder@2665ffd · GitHub
[go: up one dir, main page]
Skip to content

Commit 2665ffd

Browse files
spikecurtisspikecurtis
committed
chore: move InProcNet to testutil
1 parent 6bebfd0 commit 2665ffd

File tree

2 files changed

+117
-108
lines changed

2 files changed

+117
-108
lines changed

‎cli/portforward_test.go

Lines changed: 12 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import (
1313
"github.com/pion/udp"
1414
"github.com/stretchr/testify/assert"
1515
"github.com/stretchr/testify/require"
16-
"golang.org/x/xerrors"
1716

1817
"github.com/coder/coder/v2/agent"
1918
"github.com/coder/coder/v2/agent/agenttest"
@@ -161,7 +160,7 @@ func TestPortForward(t *testing.T) {
161160
inv.Stdout = pty.Output()
162161
inv.Stderr = pty.Output()
163162

164-
iNet := newInProcNet()
163+
iNet := testutil.NewInProcNet()
165164
inv.Net = iNet
166165
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
167166
defer cancel()
@@ -177,10 +176,10 @@ func TestPortForward(t *testing.T) {
177176
// sync.
178177
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
179178
defer dialCtxCancel()
180-
c1, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]})
179+
c1, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0]))
181180
require.NoError(t, err, "open connection 1 to 'local' listener")
182181
defer c1.Close()
183-
c2, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]})
182+
c2, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0]))
184183
require.NoError(t, err, "open connection 2 to 'local' listener")
185184
defer c2.Close()
186185
testDial(t, c2)
@@ -218,7 +217,7 @@ func TestPortForward(t *testing.T) {
218217
inv.Stdout = pty.Output()
219218
inv.Stderr = pty.Output()
220219

221-
iNet := newInProcNet()
220+
iNet := testutil.NewInProcNet()
222221
inv.Net = iNet
223222
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
224223
defer cancel()
@@ -232,10 +231,10 @@ func TestPortForward(t *testing.T) {
232231
// then test them out of order.
233232
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
234233
defer dialCtxCancel()
235-
c1, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]})
234+
c1, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0]))
236235
require.NoError(t, err, "open connection 1 to 'local' listener 1")
237236
defer c1.Close()
238-
c2, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[1]})
237+
c2, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[1]))
239238
require.NoError(t, err, "open connection 2 to 'local' listener 2")
240239
defer c2.Close()
241240
testDial(t, c2)
@@ -257,18 +256,15 @@ func TestPortForward(t *testing.T) {
257256
t.Run("All", func(t *testing.T) {
258257
t.Parallel()
259258
var (
260-
dials = []addr{}
259+
dials = []testutil.Addr{}
261260
flags = []string{}
262261
)
263262

264263
// Start listeners and populate arrays with the cases.
265264
for _, c := range cases {
266265
p := setupTestListener(t, c.setupRemote(t))
267266

268-
dials = append(dials, addr{
269-
network: c.network,
270-
addr: c.localAddress[0],
271-
})
267+
dials = append(dials, testutil.NewAddr(c.network, c.localAddress[0]))
272268
flags = append(flags, fmt.Sprintf(c.flag[0], p))
273269
}
274270

@@ -279,7 +275,7 @@ func TestPortForward(t *testing.T) {
279275
pty := ptytest.New(t).Attach(inv)
280276
inv.Stderr = pty.Output()
281277

282-
iNet := newInProcNet()
278+
iNet := testutil.NewInProcNet()
283279
inv.Net = iNet
284280
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
285281
defer cancel()
@@ -296,7 +292,7 @@ func TestPortForward(t *testing.T) {
296292
)
297293
defer dialCtxCancel()
298294
for i, a := range dials {
299-
c, err := iNet.dial(dialCtx, a)
295+
c, err := iNet.Dial(dialCtx, a)
300296
require.NoErrorf(t, err, "open connection %v to 'local' listener %v", i+1, i+1)
301297
t.Cleanup(func() {
302298
_ = c.Close()
@@ -340,7 +336,7 @@ func TestPortForward(t *testing.T) {
340336
inv.Stdout = pty.Output()
341337
inv.Stderr = pty.Output()
342338

343-
iNet := newInProcNet()
339+
iNet := testutil.NewInProcNet()
344340
inv.Net = iNet
345341

346342
// listen on port 5555 on IPv6 so it's busy when we try to port forward
@@ -361,7 +357,7 @@ func TestPortForward(t *testing.T) {
361357
// Test IPv4 still works
362358
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
363359
defer dialCtxCancel()
364-
c1, err := iNet.dial(dialCtx, addr{"tcp", "127.0.0.1:5555"})
360+
c1, err := iNet.Dial(dialCtx, testutil.NewAddr("tcp", "127.0.0.1:5555"))
365361
require.NoError(t, err, "open connection 1 to 'local' listener")
366362
defer c1.Close()
367363
testDial(t, c1)
@@ -473,95 +469,3 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
473469
assert.NoError(t, err, "write payload")
474470
assert.Equal(t, len(payload), n, "payload length does not match")
475471
}
476-
477-
type addr struct {
478-
network string
479-
addr string
480-
}
481-
482-
func (a addr) Network() string {
483-
return a.network
484-
}
485-
486-
func (a addr) Address() string {
487-
return a.addr
488-
}
489-
490-
func (a addr) String() string {
491-
return a.network + "|" + a.addr
492-
}
493-
494-
type inProcNet struct {
495-
sync.Mutex
496-
497-
listeners map[addr]*inProcListener
498-
}
499-
500-
type inProcListener struct {
501-
c chan net.Conn
502-
n *inProcNet
503-
a addr
504-
o sync.Once
505-
}
506-
507-
func newInProcNet() *inProcNet {
508-
return &inProcNet{listeners: make(map[addr]*inProcListener)}
509-
}
510-
511-
func (n *inProcNet) Listen(network, address string) (net.Listener, error) {
512-
a := addr{network, address}
513-
n.Lock()
514-
defer n.Unlock()
515-
if _, ok := n.listeners[a]; ok {
516-
return nil, xerrors.New("busy")
517-
}
518-
l := newInProcListener(n, a)
519-
n.listeners[a] = l
520-
return l, nil
521-
}
522-
523-
func (n *inProcNet) dial(ctx context.Context, a addr) (net.Conn, error) {
524-
n.Lock()
525-
defer n.Unlock()
526-
l, ok := n.listeners[a]
527-
if !ok {
528-
return nil, xerrors.Errorf("nothing listening on %s", a)
529-
}
530-
x, y := net.Pipe()
531-
select {
532-
case <-ctx.Done():
533-
return nil, ctx.Err()
534-
case l.c <- x:
535-
return y, nil
536-
}
537-
}
538-
539-
func newInProcListener(n *inProcNet, a addr) *inProcListener {
540-
return &inProcListener{
541-
c: make(chan net.Conn),
542-
n: n,
543-
a: a,
544-
}
545-
}
546-
547-
func (l *inProcListener) Accept() (net.Conn, error) {
548-
c, ok := <-l.c
549-
if !ok {
550-
return nil, net.ErrClosed
551-
}
552-
return c, nil
553-
}
554-
555-
func (l *inProcListener) Close() error {
556-
l.o.Do(func() {
557-
l.n.Lock()
558-
defer l.n.Unlock()
559-
delete(l.n.listeners, l.a)
560-
close(l.c)
561-
})
562-
return nil
563-
}
564-
565-
func (l *inProcListener) Addr() net.Addr {
566-
return l.a
567-
}

‎testutil/net.go

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package testutil
2+
3+
import (
4+
"context"
5+
"net"
6+
"sync"
7+
8+
"golang.org/x/xerrors"
9+
)
10+
11+
type Addr struct {
12+
network string
13+
addr string
14+
}
15+
16+
func NewAddr(network, addr string) Addr {
17+
return Addr{network, addr}
18+
}
19+
20+
func (a Addr) Network() string {
21+
return a.network
22+
}
23+
24+
func (a Addr) Address() string {
25+
return a.addr
26+
}
27+
28+
func (a Addr) String() string {
29+
return a.network + "|" + a.addr
30+
}
31+
32+
type InProcNet struct {
33+
sync.Mutex
34+
35+
listeners map[Addr]*inProcListener
36+
}
37+
38+
type inProcListener struct {
39+
c chan net.Conn
40+
n *InProcNet
41+
a Addr
42+
o sync.Once
43+
}
44+
45+
func NewInProcNet() *InProcNet {
46+
return &InProcNet{listeners: make(map[Addr]*inProcListener)}
47+
}
48+
49+
func (n *InProcNet) Listen(network, address string) (net.Listener, error) {
50+
a := Addr{network, address}
51+
n.Lock()
52+
defer n.Unlock()
53+
if _, ok := n.listeners[a]; ok {
54+
return nil, xerrors.New("busy")
55+
}
56+
l := newInProcListener(n, a)
57+
n.listeners[a] = l
58+
return l, nil
59+
}
60+
61+
func (n *InProcNet) Dial(ctx context.Context, a Addr) (net.Conn, error) {
62+
n.Lock()
63+
defer n.Unlock()
64+
l, ok := n.listeners[a]
65+
if !ok {
66+
return nil, xerrors.Errorf("nothing listening on %s", a)
67+
}
68+
x, y := net.Pipe()
69+
select {
70+
case <-ctx.Done():
71+
return nil, ctx.Err()
72+
case l.c <- x:
73+
return y, nil
74+
}
75+
}
76+
77+
func newInProcListener(n *InProcNet, a Addr) *inProcListener {
78+
return &inProcListener{
79+
c: make(chan net.Conn),
80+
n: n,
81+
a: a,
82+
}
83+
}
84+
85+
func (l *inProcListener) Accept() (net.Conn, error) {
86+
c, ok := <-l.c
87+
if !ok {
88+
return nil, net.ErrClosed
89+
}
90+
return c, nil
91+
}
92+
93+
func (l *inProcListener) Close() error {
94+
l.o.Do(func() {
95+
l.n.Lock()
96+
defer l.n.Unlock()
97+
delete(l.n.listeners, l.a)
98+
close(l.c)
99+
})
100+
return nil
101+
}
102+
103+
func (l *inProcListener) Addr() net.Addr {
104+
return l.a
105+
}

0 commit comments

Comments
 (0)
0