From 32cfa80f29d1e57930183ea29ce61e38dc9632e5 Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Tue, 24 Jun 2025 05:28:37 +0000 Subject: [PATCH] chore: move InProcNet to testutil --- cli/portforward_test.go | 120 ++++------------------------------------ testutil/net.go | 105 +++++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 108 deletions(-) create mode 100644 testutil/net.go diff --git a/cli/portforward_test.go b/cli/portforward_test.go index e995b31950314..9899bd28cccdf 100644 --- a/cli/portforward_test.go +++ b/cli/portforward_test.go @@ -13,7 +13,6 @@ import ( "github.com/pion/udp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "golang.org/x/xerrors" "github.com/coder/coder/v2/agent" "github.com/coder/coder/v2/agent/agenttest" @@ -161,7 +160,7 @@ func TestPortForward(t *testing.T) { inv.Stdout = pty.Output() inv.Stderr = pty.Output() - iNet := newInProcNet() + iNet := testutil.NewInProcNet() inv.Net = iNet ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -177,10 +176,10 @@ func TestPortForward(t *testing.T) { // sync. dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort) defer dialCtxCancel() - c1, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]}) + c1, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0])) require.NoError(t, err, "open connection 1 to 'local' listener") defer c1.Close() - c2, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]}) + c2, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0])) require.NoError(t, err, "open connection 2 to 'local' listener") defer c2.Close() testDial(t, c2) @@ -218,7 +217,7 @@ func TestPortForward(t *testing.T) { inv.Stdout = pty.Output() inv.Stderr = pty.Output() - iNet := newInProcNet() + iNet := testutil.NewInProcNet() inv.Net = iNet ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -232,10 +231,10 @@ func TestPortForward(t *testing.T) { // then test them out of order. dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort) defer dialCtxCancel() - c1, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[0]}) + c1, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[0])) require.NoError(t, err, "open connection 1 to 'local' listener 1") defer c1.Close() - c2, err := iNet.dial(dialCtx, addr{c.network, c.localAddress[1]}) + c2, err := iNet.Dial(dialCtx, testutil.NewAddr(c.network, c.localAddress[1])) require.NoError(t, err, "open connection 2 to 'local' listener 2") defer c2.Close() testDial(t, c2) @@ -257,7 +256,7 @@ func TestPortForward(t *testing.T) { t.Run("All", func(t *testing.T) { t.Parallel() var ( - dials = []addr{} + dials = []testutil.Addr{} flags = []string{} ) @@ -265,10 +264,7 @@ func TestPortForward(t *testing.T) { for _, c := range cases { p := setupTestListener(t, c.setupRemote(t)) - dials = append(dials, addr{ - network: c.network, - addr: c.localAddress[0], - }) + dials = append(dials, testutil.NewAddr(c.network, c.localAddress[0])) flags = append(flags, fmt.Sprintf(c.flag[0], p)) } @@ -279,7 +275,7 @@ func TestPortForward(t *testing.T) { pty := ptytest.New(t).Attach(inv) inv.Stderr = pty.Output() - iNet := newInProcNet() + iNet := testutil.NewInProcNet() inv.Net = iNet ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() @@ -296,7 +292,7 @@ func TestPortForward(t *testing.T) { ) defer dialCtxCancel() for i, a := range dials { - c, err := iNet.dial(dialCtx, a) + c, err := iNet.Dial(dialCtx, a) require.NoErrorf(t, err, "open connection %v to 'local' listener %v", i+1, i+1) t.Cleanup(func() { _ = c.Close() @@ -340,7 +336,7 @@ func TestPortForward(t *testing.T) { inv.Stdout = pty.Output() inv.Stderr = pty.Output() - iNet := newInProcNet() + iNet := testutil.NewInProcNet() inv.Net = iNet // 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) { // Test IPv4 still works dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort) defer dialCtxCancel() - c1, err := iNet.dial(dialCtx, addr{"tcp", "127.0.0.1:5555"}) + c1, err := iNet.Dial(dialCtx, testutil.NewAddr("tcp", "127.0.0.1:5555")) require.NoError(t, err, "open connection 1 to 'local' listener") defer c1.Close() testDial(t, c1) @@ -473,95 +469,3 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) { assert.NoError(t, err, "write payload") assert.Equal(t, len(payload), n, "payload length does not match") } - -type addr struct { - network string - addr string -} - -func (a addr) Network() string { - return a.network -} - -func (a addr) Address() string { - return a.addr -} - -func (a addr) String() string { - return a.network + "|" + a.addr -} - -type inProcNet struct { - sync.Mutex - - listeners map[addr]*inProcListener -} - -type inProcListener struct { - c chan net.Conn - n *inProcNet - a addr - o sync.Once -} - -func newInProcNet() *inProcNet { - return &inProcNet{listeners: make(map[addr]*inProcListener)} -} - -func (n *inProcNet) Listen(network, address string) (net.Listener, error) { - a := addr{network, address} - n.Lock() - defer n.Unlock() - if _, ok := n.listeners[a]; ok { - return nil, xerrors.New("busy") - } - l := newInProcListener(n, a) - n.listeners[a] = l - return l, nil -} - -func (n *inProcNet) dial(ctx context.Context, a addr) (net.Conn, error) { - n.Lock() - defer n.Unlock() - l, ok := n.listeners[a] - if !ok { - return nil, xerrors.Errorf("nothing listening on %s", a) - } - x, y := net.Pipe() - select { - case <-ctx.Done(): - return nil, ctx.Err() - case l.c <- x: - return y, nil - } -} - -func newInProcListener(n *inProcNet, a addr) *inProcListener { - return &inProcListener{ - c: make(chan net.Conn), - n: n, - a: a, - } -} - -func (l *inProcListener) Accept() (net.Conn, error) { - c, ok := <-l.c - if !ok { - return nil, net.ErrClosed - } - return c, nil -} - -func (l *inProcListener) Close() error { - l.o.Do(func() { - l.n.Lock() - defer l.n.Unlock() - delete(l.n.listeners, l.a) - close(l.c) - }) - return nil -} - -func (l *inProcListener) Addr() net.Addr { - return l.a -} diff --git a/testutil/net.go b/testutil/net.go new file mode 100644 index 0000000000000..b38597c9d9b73 --- /dev/null +++ b/testutil/net.go @@ -0,0 +1,105 @@ +package testutil + +import ( + "context" + "net" + "sync" + + "golang.org/x/xerrors" +) + +type Addr struct { + network string + addr string +} + +func NewAddr(network, addr string) Addr { + return Addr{network, addr} +} + +func (a Addr) Network() string { + return a.network +} + +func (a Addr) Address() string { + return a.addr +} + +func (a Addr) String() string { + return a.network + "|" + a.addr +} + +type InProcNet struct { + sync.Mutex + + listeners map[Addr]*inProcListener +} + +type inProcListener struct { + c chan net.Conn + n *InProcNet + a Addr + o sync.Once +} + +func NewInProcNet() *InProcNet { + return &InProcNet{listeners: make(map[Addr]*inProcListener)} +} + +func (n *InProcNet) Listen(network, address string) (net.Listener, error) { + a := Addr{network, address} + n.Lock() + defer n.Unlock() + if _, ok := n.listeners[a]; ok { + return nil, xerrors.New("busy") + } + l := newInProcListener(n, a) + n.listeners[a] = l + return l, nil +} + +func (n *InProcNet) Dial(ctx context.Context, a Addr) (net.Conn, error) { + n.Lock() + defer n.Unlock() + l, ok := n.listeners[a] + if !ok { + return nil, xerrors.Errorf("nothing listening on %s", a) + } + x, y := net.Pipe() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case l.c <- x: + return y, nil + } +} + +func newInProcListener(n *InProcNet, a Addr) *inProcListener { + return &inProcListener{ + c: make(chan net.Conn), + n: n, + a: a, + } +} + +func (l *inProcListener) Accept() (net.Conn, error) { + c, ok := <-l.c + if !ok { + return nil, net.ErrClosed + } + return c, nil +} + +func (l *inProcListener) Close() error { + l.o.Do(func() { + l.n.Lock() + defer l.n.Unlock() + delete(l.n.listeners, l.a) + close(l.c) + }) + return nil +} + +func (l *inProcListener) Addr() net.Addr { + return l.a +}