fix: use independent pipes in vpn tunnel tests

ethanndickson · ethanndickson · commit 5b19e5e96be1 · 2024-12-13T15:38:12.000Z
diff --git a/vpn/tunnel_internal_test.go b/vpn/tunnel_internal_test.go
@@ -2,6 +2,7 @@ package vpn
 
 import (
 	"context"
+	"io"
 	"net"
 	"net/netip"
 	"net/url"
@@ -485,9 +486,7 @@ func TestTunnel_sendAgentUpdate(t *testing.T) {
 
 //nolint:revive // t takes precedence
 func setupTunnel(t *testing.T, ctx context.Context, client *fakeClient, mClock quartz.Clock) (*Tunnel, *speaker[*ManagerMessage, *TunnelMessage, TunnelMessage]) {
-	mp, tp := net.Pipe()
-	t.Cleanup(func() { _ = mp.Close() })
-	t.Cleanup(func() { _ = tp.Close() })
+	mp, tp := newBidiPipe(t)
 	logger := testutil.Logger(t)
 
 	var tun *Tunnel
@@ -510,3 +509,55 @@ func setupTunnel(t *testing.T, ctx context.Context, client *fakeClient, mClock q
 	mgr.start()
 	return tun, mgr
 }
+
+// In practice, the underlying connection is a pair of dup'd pipes. The tunnel
+// closing the connection (the dup'd FD) should not cause the manager to read EOF.
+// By replicating the behavior of these FDs, we avoid a race where the manager
+// reads EOF before it reads the stop response from the tunnel.
+type bidiPipe struct {
+	rw io.ReadWriteCloser
+	*refcnt
+}
+
+type refcnt struct {
+	count     int
+	mu        sync.Mutex
+	closeFunc func() error
+}
+
+// Not idempotent, but it's only called once per bidiPipe.
+func (b *bidiPipe) Close() error {
+	b.refcnt.mu.Lock()
+	defer b.refcnt.mu.Unlock()
+
+	b.refcnt.count--
+	if b.refcnt.count == 0 {
+		_ = b.refcnt.closeFunc()
+	}
+	return nil
+}
+
+func (b *bidiPipe) Read(p []byte) (n int, err error) {
+	return b.rw.Read(p)
+}
+
+func (b *bidiPipe) Write(p []byte) (n int, err error) {
+	return b.rw.Write(p)
+}
+
+func newBidiPipe(t *testing.T) (io.ReadWriteCloser, io.ReadWriteCloser) {
+	mp, tp := net.Pipe()
+	t.Cleanup(func() { _ = mp.Close() })
+	t.Cleanup(func() { _ = tp.Close() })
+
+	refcnt := &refcnt{
+		count: 2,
+		closeFunc: func() error {
+			_ = tp.Close()
+			_ = mp.Close()
+			return nil
+		},
+	}
+
+	return &bidiPipe{rw: tp, refcnt: refcnt}, &bidiPipe{rw: mp, refcnt: refcnt}
+}
