mirror of
https://github.com/coder/coder.git
synced 2025-07-12 00:14:10 +00:00
chore: move InProcNet to testutil (#18563)
Moves `InProcNet` to `testutil` so that it can be reused by X11 forwarding tests (see up stack PRs).
This commit is contained in:
@ -13,7 +13,6 @@ import (
|
|||||||
"github.com/pion/udp"
|
"github.com/pion/udp"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
"golang.org/x/xerrors"
|
|
||||||
|
|
||||||
"github.com/coder/coder/v2/agent"
|
"github.com/coder/coder/v2/agent"
|
||||||
"github.com/coder/coder/v2/agent/agenttest"
|
"github.com/coder/coder/v2/agent/agenttest"
|
||||||
@ -161,7 +160,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
inv.Stdout = pty.Output()
|
inv.Stdout = pty.Output()
|
||||||
inv.Stderr = pty.Output()
|
inv.Stderr = pty.Output()
|
||||||
|
|
||||||
iNet := newInProcNet()
|
iNet := testutil.NewInProcNet()
|
||||||
inv.Net = iNet
|
inv.Net = iNet
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@ -177,10 +176,10 @@ func TestPortForward(t *testing.T) {
|
|||||||
// sync.
|
// sync.
|
||||||
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
||||||
defer dialCtxCancel()
|
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")
|
require.NoError(t, err, "open connection 1 to 'local' listener")
|
||||||
defer c1.Close()
|
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")
|
require.NoError(t, err, "open connection 2 to 'local' listener")
|
||||||
defer c2.Close()
|
defer c2.Close()
|
||||||
testDial(t, c2)
|
testDial(t, c2)
|
||||||
@ -218,7 +217,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
inv.Stdout = pty.Output()
|
inv.Stdout = pty.Output()
|
||||||
inv.Stderr = pty.Output()
|
inv.Stderr = pty.Output()
|
||||||
|
|
||||||
iNet := newInProcNet()
|
iNet := testutil.NewInProcNet()
|
||||||
inv.Net = iNet
|
inv.Net = iNet
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@ -232,10 +231,10 @@ func TestPortForward(t *testing.T) {
|
|||||||
// then test them out of order.
|
// then test them out of order.
|
||||||
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
||||||
defer dialCtxCancel()
|
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")
|
require.NoError(t, err, "open connection 1 to 'local' listener 1")
|
||||||
defer c1.Close()
|
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")
|
require.NoError(t, err, "open connection 2 to 'local' listener 2")
|
||||||
defer c2.Close()
|
defer c2.Close()
|
||||||
testDial(t, c2)
|
testDial(t, c2)
|
||||||
@ -257,7 +256,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
t.Run("All", func(t *testing.T) {
|
t.Run("All", func(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
var (
|
var (
|
||||||
dials = []addr{}
|
dials = []testutil.Addr{}
|
||||||
flags = []string{}
|
flags = []string{}
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -265,10 +264,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
for _, c := range cases {
|
for _, c := range cases {
|
||||||
p := setupTestListener(t, c.setupRemote(t))
|
p := setupTestListener(t, c.setupRemote(t))
|
||||||
|
|
||||||
dials = append(dials, addr{
|
dials = append(dials, testutil.NewAddr(c.network, c.localAddress[0]))
|
||||||
network: c.network,
|
|
||||||
addr: c.localAddress[0],
|
|
||||||
})
|
|
||||||
flags = append(flags, fmt.Sprintf(c.flag[0], p))
|
flags = append(flags, fmt.Sprintf(c.flag[0], p))
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -279,7 +275,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
pty := ptytest.New(t).Attach(inv)
|
pty := ptytest.New(t).Attach(inv)
|
||||||
inv.Stderr = pty.Output()
|
inv.Stderr = pty.Output()
|
||||||
|
|
||||||
iNet := newInProcNet()
|
iNet := testutil.NewInProcNet()
|
||||||
inv.Net = iNet
|
inv.Net = iNet
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
@ -296,7 +292,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
)
|
)
|
||||||
defer dialCtxCancel()
|
defer dialCtxCancel()
|
||||||
for i, a := range dials {
|
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)
|
require.NoErrorf(t, err, "open connection %v to 'local' listener %v", i+1, i+1)
|
||||||
t.Cleanup(func() {
|
t.Cleanup(func() {
|
||||||
_ = c.Close()
|
_ = c.Close()
|
||||||
@ -340,7 +336,7 @@ func TestPortForward(t *testing.T) {
|
|||||||
inv.Stdout = pty.Output()
|
inv.Stdout = pty.Output()
|
||||||
inv.Stderr = pty.Output()
|
inv.Stderr = pty.Output()
|
||||||
|
|
||||||
iNet := newInProcNet()
|
iNet := testutil.NewInProcNet()
|
||||||
inv.Net = iNet
|
inv.Net = iNet
|
||||||
|
|
||||||
// listen on port 5555 on IPv6 so it's busy when we try to port forward
|
// 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
|
// Test IPv4 still works
|
||||||
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
dialCtx, dialCtxCancel := context.WithTimeout(ctx, testutil.WaitShort)
|
||||||
defer dialCtxCancel()
|
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")
|
require.NoError(t, err, "open connection 1 to 'local' listener")
|
||||||
defer c1.Close()
|
defer c1.Close()
|
||||||
testDial(t, c1)
|
testDial(t, c1)
|
||||||
@ -473,95 +469,3 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
|
|||||||
assert.NoError(t, err, "write payload")
|
assert.NoError(t, err, "write payload")
|
||||||
assert.Equal(t, len(payload), n, "payload length does not match")
|
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
|
|
||||||
}
|
|
||||||
|
105
testutil/net.go
Normal file
105
testutil/net.go
Normal file
@ -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
|
||||||
|
}
|
Reference in New Issue
Block a user