From 898be13487945eb27b2ecfcd3c7ee396401d267d Mon Sep 17 00:00:00 2001 From: edwardowens Date: Tue, 8 Sep 2026 13:08:25 -0700 Subject: [PATCH 1/3] fix(tunnel): queue chunks that arrive mid-dial instead of dialing again The tunnel has no explicit open message: the first data chunk for an unknown connection ID is what opens it, and the receiving side only registered the ID after its dial returned. Any further chunk for that ID arriving during the dial (Node's HTTP/2 client sends the preface and the first request frames ~1-2ms apart) spawned another dial for the same connection, splitting the stream across two sockets. The peer saw a preface on one socket and headers on the other and simply never answered, so the caller hung until its own deadline instead of getting an error. Register a buffering placeholder for the ID synchronously on the recv pump, dial off the pump, then swap the real connection in and flush the queued chunks in order. A peer error that lands while the dial is in flight now also closes the dialed socket instead of leaking it. Co-Authored-By: Claude Fable 5 --- pkg/tunnel/pending.go | 87 +++++++++++++ pkg/tunnel/tunnel.go | 65 +++++++--- pkg/tunnel/tunnel_test.go | 251 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 384 insertions(+), 19 deletions(-) create mode 100644 pkg/tunnel/pending.go create mode 100644 pkg/tunnel/tunnel_test.go diff --git a/pkg/tunnel/pending.go b/pkg/tunnel/pending.go new file mode 100644 index 0000000..bf8f904 --- /dev/null +++ b/pkg/tunnel/pending.go @@ -0,0 +1,87 @@ +package tunnel + +import ( + "errors" + "io" + "net" + "sync" +) + +var errNotResolved = errors.New("tunnel: connection not yet dialed") + +// pendingConn stands in for a tunneled connection whose dial has not completed. +// The recv pump registers it under the connection ID before dialing so that +// chunks arriving while the dial is in flight queue here, in order, instead of +// each spawning another dial for the same connection. +type pendingConn struct { + mu sync.Mutex + conn net.Conn + queued [][]byte + closed bool +} + +var _ io.ReadWriteCloser = (*pendingConn)(nil) + +// resolve attaches the dialed connection and flushes the queued chunks in +// order. It closes conn and reports the reason if the pending connection was +// already closed or a queued chunk fails to write. +func (p *pendingConn) resolve(conn net.Conn) error { + p.mu.Lock() + defer p.mu.Unlock() + if p.closed { + conn.Close() + return net.ErrClosed + } + for _, b := range p.queued { + if _, err := conn.Write(b); err != nil { + conn.Close() + p.closed = true + p.queued = nil + return err + } + } + p.queued = nil + p.conn = conn + return nil +} + +func (p *pendingConn) Write(b []byte) (int, error) { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return 0, net.ErrClosed + } + if p.conn == nil { + p.queued = append(p.queued, append([]byte(nil), b...)) + p.mu.Unlock() + return len(b), nil + } + conn := p.conn + p.mu.Unlock() + return conn.Write(b) +} + +func (p *pendingConn) Read(b []byte) (int, error) { + p.mu.Lock() + conn, closed := p.conn, p.closed + p.mu.Unlock() + if closed { + return 0, net.ErrClosed + } + if conn == nil { + return 0, errNotResolved + } + return conn.Read(b) +} + +func (p *pendingConn) Close() error { + p.mu.Lock() + conn := p.conn + p.closed = true + p.queued = nil + p.mu.Unlock() + if conn != nil { + return conn.Close() + } + return nil +} diff --git a/pkg/tunnel/tunnel.go b/pkg/tunnel/tunnel.go index d280daa..d42ec96 100644 --- a/pkg/tunnel/tunnel.go +++ b/pkg/tunnel/tunnel.go @@ -53,7 +53,7 @@ func New(dialer plumbing.ContextDialer, stream Stream, opts ...Option) Tunnel { dialer: dialer, stream: stream, sendCh: make(chan *bridgev1.TunnelNetworkMessage, 64), - conns: xsync.NewMapOf[string, net.Conn](), + conns: xsync.NewMapOf[string, io.ReadWriteCloser](), done: make(chan struct{}), } for _, o := range opts { @@ -67,7 +67,7 @@ type tunnelImpl struct { hijacker mitm.Hijacker stream Stream sendCh chan *bridgev1.TunnelNetworkMessage - conns *xsync.MapOf[string, net.Conn] + conns *xsync.MapOf[string, io.ReadWriteCloser] done chan struct{} ctx context.Context cancel context.CancelFunc @@ -99,8 +99,8 @@ func (t *tunnelImpl) AddConn(conn net.Conn, destOverride string, hostname string go t.readFromConn(conn, connID, src, dst, hostname) } -// readFromConn reads from a net.Conn and forwards data to the stream via sendCh. -func (t *tunnelImpl) readFromConn(conn net.Conn, connID string, src, dst *bridgev1.TunnelAddress, hostname string) { +// readFromConn reads from a connection and forwards data to the stream via sendCh. +func (t *tunnelImpl) readFromConn(conn io.ReadCloser, connID string, src, dst *bridgev1.TunnelAddress, hostname string) { defer func() { conn.Close() t.conns.Delete(connID) @@ -200,8 +200,8 @@ func (t *tunnelImpl) Start(ctx context.Context) { continue } - // Unknown connection ID → dial via the configured dialer. - go t.handleNewConn(msg) + // Unknown connection ID → this message opens it. + t.openConn(msg) case err := <-recvErr: if err != io.EOF { @@ -218,15 +218,31 @@ func (t *tunnelImpl) Start(ctx context.Context) { }() } -func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) { - dest := msg.GetDest() - if dest == nil { - slog.Info("Tunnel: ignoring message with no dest", "conn_id", msg.GetConnectionId()) +// openConn registers a pending connection for msg's ID before dialing. The +// peer has no explicit open message: the first chunk of a connection is what +// opens it, and the next chunks can arrive before the dial completes. Storing +// the placeholder synchronously on the recv pump makes those chunks queue +// behind the first one instead of each dialing a second connection under the +// same ID and splitting the stream between them. +func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) { + connID := msg.GetConnectionId() + if msg.GetDest() == nil { + slog.Info("Tunnel: ignoring message with no dest", "conn_id", connID) return } + pending := &pendingConn{} + t.conns.Store(connID, pending) + if data := msg.GetData(); len(data) > 0 { + _, _ = pending.Write(data) + } + go t.dialPending(msg, pending) +} + +func (t *tunnelImpl) dialPending(msg *bridgev1.TunnelNetworkMessage, pending *pendingConn) { connID := msg.GetConnectionId() hostname := msg.GetHostname() + dest := msg.GetDest() // Resolve the connection: hijacker gets first shot, then fall through to the dialer. var conn net.Conn @@ -246,6 +262,8 @@ func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) { if err != nil { slog.Info("Tunnel: connect failed", "conn_id", connID, "hostname", hostname, "error", err) + t.deletePending(connID, pending) + pending.Close() select { case t.sendCh <- &bridgev1.TunnelNetworkMessage{ ConnectionId: connID, @@ -256,20 +274,29 @@ func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) { return } - t.conns.Store(connID, conn) - go t.readFromConn(conn, connID, msg.GetDest(), msg.GetSource(), hostname) + if err := pending.resolve(conn); err != nil { + slog.Debug("Tunnel: dropping dialed connection", "connection_id", connID, "error", err) + t.deletePending(connID, pending) + return + } - if data := msg.GetData(); len(data) > 0 { - if _, err := conn.Write(data); err != nil { - slog.Debug("Failed to write initial data", "connection_id", connID, "error", err) - conn.Close() - t.conns.Delete(connID) + go t.readFromConn(pending, connID, dest, msg.GetSource(), hostname) +} + +// deletePending removes connID only while it still maps to pending, so a +// connection the peer has since reopened under the same ID is left alone. +func (t *tunnelImpl) deletePending(connID string, pending *pendingConn) { + t.conns.Compute(connID, func(cur io.ReadWriteCloser, loaded bool) (io.ReadWriteCloser, bool) { + if !loaded { + return nil, true } - } + p, ok := cur.(*pendingConn) + return cur, ok && p == pending + }) } func (t *tunnelImpl) closeAll() { - t.conns.Range(func(key string, conn net.Conn) bool { + t.conns.Range(func(key string, conn io.ReadWriteCloser) bool { conn.Close() t.conns.Delete(key) return true diff --git a/pkg/tunnel/tunnel_test.go b/pkg/tunnel/tunnel_test.go new file mode 100644 index 0000000..b0c01db --- /dev/null +++ b/pkg/tunnel/tunnel_test.go @@ -0,0 +1,251 @@ +package tunnel + +import ( + "context" + "errors" + "io" + "net" + "sync" + "sync/atomic" + "testing" + "time" + + bridgev1 "github.com/vercel/bridge/api/go/bridge/v1" +) + +const testTimeout = 5 * time.Second + +type fakeStream struct { + in chan *bridgev1.TunnelNetworkMessage + out chan *bridgev1.TunnelNetworkMessage +} + +func newFakeStream() *fakeStream { + return &fakeStream{ + in: make(chan *bridgev1.TunnelNetworkMessage), + out: make(chan *bridgev1.TunnelNetworkMessage, 64), + } +} + +func (s *fakeStream) Send(msg *bridgev1.TunnelNetworkMessage) error { + s.out <- msg + return nil +} + +func (s *fakeStream) Recv() (*bridgev1.TunnelNetworkMessage, error) { + msg, ok := <-s.in + if !ok { + return nil, io.EOF + } + return msg, nil +} + +// fakeConn records writes and blocks reads until closed. +type fakeConn struct { + mu sync.Mutex + written []byte + done chan struct{} + once sync.Once +} + +func newFakeConn() *fakeConn { + return &fakeConn{done: make(chan struct{})} +} + +func (c *fakeConn) Write(b []byte) (int, error) { + c.mu.Lock() + c.written = append(c.written, b...) + c.mu.Unlock() + return len(b), nil +} + +func (c *fakeConn) Read([]byte) (int, error) { + <-c.done + return 0, io.EOF +} + +func (c *fakeConn) Close() error { + c.once.Do(func() { close(c.done) }) + return nil +} + +func (c *fakeConn) Written() string { + c.mu.Lock() + defer c.mu.Unlock() + return string(c.written) +} + +func (c *fakeConn) LocalAddr() net.Addr { return &net.TCPAddr{} } +func (c *fakeConn) RemoteAddr() net.Addr { return &net.TCPAddr{} } +func (c *fakeConn) SetDeadline(time.Time) error { return nil } +func (c *fakeConn) SetReadDeadline(time.Time) error { return nil } +func (c *fakeConn) SetWriteDeadline(time.Time) error { return nil } + +// blockingDialer counts dials and parks each one until release is closed. +type blockingDialer struct { + dials atomic.Int32 + release chan struct{} + err error + mu sync.Mutex + conns []*fakeConn +} + +func newBlockingDialer() *blockingDialer { + return &blockingDialer{release: make(chan struct{})} +} + +func (d *blockingDialer) DialContext(ctx context.Context, _, _ string) (net.Conn, error) { + d.dials.Add(1) + select { + case <-d.release: + case <-ctx.Done(): + return nil, ctx.Err() + } + if d.err != nil { + return nil, d.err + } + conn := newFakeConn() + d.mu.Lock() + d.conns = append(d.conns, conn) + d.mu.Unlock() + return conn, nil +} + +func (d *blockingDialer) connCount() int { + d.mu.Lock() + defer d.mu.Unlock() + return len(d.conns) +} + +func (d *blockingDialer) conn(t *testing.T, i int) *fakeConn { + t.Helper() + d.mu.Lock() + defer d.mu.Unlock() + if len(d.conns) <= i { + t.Fatalf("expected at least %d dialed conns, got %d", i+1, len(d.conns)) + } + return d.conns[i] +} + +func dataMsg(connID string, data string) *bridgev1.TunnelNetworkMessage { + return &bridgev1.TunnelNetworkMessage{ + ConnectionId: connID, + Source: &bridgev1.TunnelAddress{Ip: "172.17.0.3", Port: 55988}, + Dest: &bridgev1.TunnelAddress{Ip: "172.20.143.45", Port: 80}, + Data: []byte(data), + } +} + +func waitFor(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(testTimeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(time.Millisecond) + } + t.Fatalf("timed out waiting for %s", what) +} + +func startTunnel(t *testing.T, dialer *blockingDialer) (*fakeStream, Tunnel) { + t.Helper() + stream := newFakeStream() + tun := New(dialer, stream) + ctx, cancel := context.WithCancel(context.Background()) + tun.Start(ctx) + t.Cleanup(func() { + cancel() + tun.Close() + }) + return stream, tun +} + +func TestChunksArrivingDuringDialShareOneConnection(t *testing.T) { + dialer := newBlockingDialer() + stream, _ := startTunnel(t, dialer) + const connID = "172.17.0.3:55988->172.20.143.45:80" + + stream.in <- dataMsg(connID, "PRI * HTTP/2.0\r\n") + stream.in <- dataMsg(connID, "HEADERS") + waitFor(t, "first dial", func() bool { return dialer.dials.Load() >= 1 }) + stream.in <- dataMsg(connID, "DATA") + + close(dialer.release) + + waitFor(t, "dial to return", func() bool { return dialer.connCount() >= 1 }) + if got := dialer.dials.Load(); got != 1 { + t.Fatalf("expected exactly one dial for the connection, got %d", got) + } + conn := dialer.conn(t, 0) + waitFor(t, "queued chunks to flush", func() bool { + return conn.Written() == "PRI * HTTP/2.0\r\nHEADERSDATA" + }) + + stream.in <- dataMsg(connID, "MORE") + waitFor(t, "post-dial chunk to write through", func() bool { + return conn.Written() == "PRI * HTTP/2.0\r\nHEADERSDATAMORE" + }) +} + +func TestDialFailureReportsErrorAndForgetsConnection(t *testing.T) { + dialer := newBlockingDialer() + dialer.err = errors.New("dial tcp: connection refused") + close(dialer.release) + stream, _ := startTunnel(t, dialer) + const connID = "172.17.0.3:55988->172.20.143.45:80" + + stream.in <- dataMsg(connID, "hello") + + select { + case msg := <-stream.out: + if msg.GetConnectionId() != connID || msg.GetError() == "" { + t.Fatalf("expected error reply for %s, got %+v", connID, msg) + } + case <-time.After(testTimeout): + t.Fatal("timed out waiting for the connect-failed reply") + } + + stream.in <- dataMsg(connID, "hello again") + waitFor(t, "a fresh dial after the failed one was forgotten", func() bool { + return dialer.dials.Load() == 2 + }) +} + +func TestPeerErrorDuringDialClosesTheDialedConnection(t *testing.T) { + dialer := newBlockingDialer() + stream, tun := startTunnel(t, dialer) + const connID = "172.17.0.3:55988->172.20.143.45:80" + + stream.in <- dataMsg(connID, "hello") + waitFor(t, "first dial", func() bool { return dialer.dials.Load() == 1 }) + stream.in <- &bridgev1.TunnelNetworkMessage{ConnectionId: connID, Error: "peer closed"} + waitFor(t, "the peer error to drop the pending connection", func() bool { + _, loaded := tun.(*tunnelImpl).conns.Load(connID) + return !loaded + }) + + close(dialer.release) + + waitFor(t, "dial to return", func() bool { return dialer.connCount() == 1 }) + conn := dialer.conn(t, 0) + select { + case <-conn.done: + case <-time.After(testTimeout): + t.Fatal("timed out waiting for the dialed conn to be closed") + } + if got := conn.Written(); got != "" { + t.Fatalf("expected nothing written to a connection closed before dial completed, got %q", got) + } +} + +func TestMessageWithoutDestIsIgnored(t *testing.T) { + dialer := newBlockingDialer() + close(dialer.release) + stream, _ := startTunnel(t, dialer) + + stream.in <- &bridgev1.TunnelNetworkMessage{ConnectionId: "unknown", Error: "late reply"} + stream.in <- dataMsg("172.17.0.3:1->172.20.143.45:80", "hello") + + waitFor(t, "the well-formed message to dial", func() bool { return dialer.dials.Load() == 1 }) +} From ca6d94dc2f7f94933f5ca559b07aa36a504f1b47 Mon Sep 17 00:00:00 2001 From: edwardowens Date: Tue, 8 Sep 2026 13:39:48 -0700 Subject: [PATCH 2/3] refactor(tunnel): extract ioutil.BufferedReadWriteCloser Replace the tunnel-private pendingConn with a reusable BufferedReadWriteCloser in pkg/ioutil: an io.ReadWriteCloser that buffers writes until Set attaches the underlying one, then flushes them as a single write and passes everything through. Co-Authored-By: Claude Fable 5 --- pkg/ioutil/buffered.go | 95 +++++++++++++++++++++++++++++ pkg/ioutil/buffered_test.go | 117 ++++++++++++++++++++++++++++++++++++ pkg/tunnel/pending.go | 87 --------------------------- pkg/tunnel/tunnel.go | 39 ++++++------ pkg/tunnel/tunnel_test.go | 2 +- 5 files changed, 233 insertions(+), 107 deletions(-) create mode 100644 pkg/ioutil/buffered.go create mode 100644 pkg/ioutil/buffered_test.go delete mode 100644 pkg/tunnel/pending.go diff --git a/pkg/ioutil/buffered.go b/pkg/ioutil/buffered.go new file mode 100644 index 0000000..f6e936b --- /dev/null +++ b/pkg/ioutil/buffered.go @@ -0,0 +1,95 @@ +package ioutil + +import ( + "bytes" + "errors" + "io" + "sync" +) + +var ( + ErrClosed = errors.New("ioutil: closed") + ErrNotSet = errors.New("ioutil: underlying ReadWriteCloser not set") + ErrAlreadySet = errors.New("ioutil: underlying ReadWriteCloser already set") +) + +// BufferedReadWriteCloser buffers writes until Set supplies the underlying +// io.ReadWriteCloser, then flushes them in order and passes everything +// through. The zero value is ready to use. +type BufferedReadWriteCloser struct { + // mu guards the handoff in Set: a Write that arrives while Set is flushing + // the buffer must land after the flush, on the underlying rwc. + mu sync.Mutex + rwc io.ReadWriteCloser + buf bytes.Buffer + closed bool +} + +var _ io.ReadWriteCloser = (*BufferedReadWriteCloser)(nil) + +// Set attaches rwc and flushes the buffered writes to it. It closes rwc and +// returns ErrClosed if Close was already called, or the write error if the +// flush fails. +func (b *BufferedReadWriteCloser) Set(rwc io.ReadWriteCloser) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + rwc.Close() + return ErrClosed + } + if b.rwc != nil { + return ErrAlreadySet + } + if b.buf.Len() > 0 { + if _, err := rwc.Write(b.buf.Bytes()); err != nil { + rwc.Close() + b.closed = true + b.buf.Reset() + return err + } + b.buf.Reset() + } + b.rwc = rwc + return nil +} + +func (b *BufferedReadWriteCloser) Write(p []byte) (int, error) { + b.mu.Lock() + if b.closed { + b.mu.Unlock() + return 0, ErrClosed + } + if b.rwc == nil { + n, err := b.buf.Write(p) + b.mu.Unlock() + return n, err + } + rwc := b.rwc + b.mu.Unlock() + return rwc.Write(p) +} + +func (b *BufferedReadWriteCloser) Read(p []byte) (int, error) { + b.mu.Lock() + rwc, closed := b.rwc, b.closed + b.mu.Unlock() + if closed { + return 0, ErrClosed + } + if rwc == nil { + return 0, ErrNotSet + } + return rwc.Read(p) +} + +func (b *BufferedReadWriteCloser) Close() error { + b.mu.Lock() + rwc := b.rwc + b.closed = true + b.buf.Reset() + b.mu.Unlock() + if rwc != nil { + return rwc.Close() + } + return nil +} diff --git a/pkg/ioutil/buffered_test.go b/pkg/ioutil/buffered_test.go new file mode 100644 index 0000000..4a45966 --- /dev/null +++ b/pkg/ioutil/buffered_test.go @@ -0,0 +1,117 @@ +package ioutil + +import ( + "bytes" + "errors" + "io" + "testing" +) + +type recordingRWC struct { + written bytes.Buffer + closed bool + writeErr error +} + +func (r *recordingRWC) Write(p []byte) (int, error) { + if r.writeErr != nil { + return 0, r.writeErr + } + return r.written.Write(p) +} + +func (r *recordingRWC) Read(p []byte) (int, error) { return copy(p, "read"), nil } + +func (r *recordingRWC) Close() error { + r.closed = true + return nil +} + +func TestWritesBeforeSetAreFlushedInOrder(t *testing.T) { + var b BufferedReadWriteCloser + rwc := &recordingRWC{} + + for _, chunk := range []string{"one ", "two ", "three"} { + if _, err := b.Write([]byte(chunk)); err != nil { + t.Fatalf("Write before Set: %v", err) + } + } + if _, err := b.Read(make([]byte, 4)); !errors.Is(err, ErrNotSet) { + t.Fatalf("Read before Set: got %v, want ErrNotSet", err) + } + + if err := b.Set(rwc); err != nil { + t.Fatalf("Set: %v", err) + } + if got := rwc.written.String(); got != "one two three" { + t.Fatalf("flushed %q, want %q", got, "one two three") + } + + if _, err := b.Write([]byte(" four")); err != nil { + t.Fatalf("Write after Set: %v", err) + } + if got := rwc.written.String(); got != "one two three four" { + t.Fatalf("after write-through got %q", got) + } + + p := make([]byte, 4) + if n, err := b.Read(p); err != nil || string(p[:n]) != "read" { + t.Fatalf("Read after Set: %q, %v", p[:n], err) + } + + if err := b.Set(&recordingRWC{}); !errors.Is(err, ErrAlreadySet) { + t.Fatalf("second Set: got %v, want ErrAlreadySet", err) + } +} + +func TestCloseBeforeSetClosesTheLateArrival(t *testing.T) { + var b BufferedReadWriteCloser + _, _ = b.Write([]byte("queued")) + if err := b.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + rwc := &recordingRWC{} + if err := b.Set(rwc); !errors.Is(err, ErrClosed) { + t.Fatalf("Set after Close: got %v, want ErrClosed", err) + } + if !rwc.closed { + t.Fatal("expected Set to close the ReadWriteCloser handed to a closed buffer") + } + if rwc.written.Len() != 0 { + t.Fatalf("expected nothing flushed to a closed buffer's late arrival, got %q", rwc.written.String()) + } + if _, err := b.Write([]byte("x")); !errors.Is(err, ErrClosed) { + t.Fatalf("Write after Close: got %v, want ErrClosed", err) + } +} + +func TestFlushFailureClosesTheUnderlying(t *testing.T) { + var b BufferedReadWriteCloser + _, _ = b.Write([]byte("queued")) + + rwc := &recordingRWC{writeErr: io.ErrShortWrite} + if err := b.Set(rwc); !errors.Is(err, io.ErrShortWrite) { + t.Fatalf("Set with failing flush: got %v, want ErrShortWrite", err) + } + if !rwc.closed { + t.Fatal("expected the underlying to be closed after a failed flush") + } + if _, err := b.Write([]byte("x")); !errors.Is(err, ErrClosed) { + t.Fatalf("Write after failed flush: got %v, want ErrClosed", err) + } +} + +func TestCloseAfterSetClosesTheUnderlying(t *testing.T) { + var b BufferedReadWriteCloser + rwc := &recordingRWC{} + if err := b.Set(rwc); err != nil { + t.Fatalf("Set: %v", err) + } + if err := b.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !rwc.closed { + t.Fatal("expected Close to close the underlying") + } +} diff --git a/pkg/tunnel/pending.go b/pkg/tunnel/pending.go deleted file mode 100644 index bf8f904..0000000 --- a/pkg/tunnel/pending.go +++ /dev/null @@ -1,87 +0,0 @@ -package tunnel - -import ( - "errors" - "io" - "net" - "sync" -) - -var errNotResolved = errors.New("tunnel: connection not yet dialed") - -// pendingConn stands in for a tunneled connection whose dial has not completed. -// The recv pump registers it under the connection ID before dialing so that -// chunks arriving while the dial is in flight queue here, in order, instead of -// each spawning another dial for the same connection. -type pendingConn struct { - mu sync.Mutex - conn net.Conn - queued [][]byte - closed bool -} - -var _ io.ReadWriteCloser = (*pendingConn)(nil) - -// resolve attaches the dialed connection and flushes the queued chunks in -// order. It closes conn and reports the reason if the pending connection was -// already closed or a queued chunk fails to write. -func (p *pendingConn) resolve(conn net.Conn) error { - p.mu.Lock() - defer p.mu.Unlock() - if p.closed { - conn.Close() - return net.ErrClosed - } - for _, b := range p.queued { - if _, err := conn.Write(b); err != nil { - conn.Close() - p.closed = true - p.queued = nil - return err - } - } - p.queued = nil - p.conn = conn - return nil -} - -func (p *pendingConn) Write(b []byte) (int, error) { - p.mu.Lock() - if p.closed { - p.mu.Unlock() - return 0, net.ErrClosed - } - if p.conn == nil { - p.queued = append(p.queued, append([]byte(nil), b...)) - p.mu.Unlock() - return len(b), nil - } - conn := p.conn - p.mu.Unlock() - return conn.Write(b) -} - -func (p *pendingConn) Read(b []byte) (int, error) { - p.mu.Lock() - conn, closed := p.conn, p.closed - p.mu.Unlock() - if closed { - return 0, net.ErrClosed - } - if conn == nil { - return 0, errNotResolved - } - return conn.Read(b) -} - -func (p *pendingConn) Close() error { - p.mu.Lock() - conn := p.conn - p.closed = true - p.queued = nil - p.mu.Unlock() - if conn != nil { - return conn.Close() - } - return nil -} diff --git a/pkg/tunnel/tunnel.go b/pkg/tunnel/tunnel.go index d42ec96..1659631 100644 --- a/pkg/tunnel/tunnel.go +++ b/pkg/tunnel/tunnel.go @@ -11,6 +11,7 @@ import ( "github.com/puzpuzpuz/xsync/v3" bridgev1 "github.com/vercel/bridge/api/go/bridge/v1" + "github.com/vercel/bridge/pkg/ioutil" "github.com/vercel/bridge/pkg/mitm" "github.com/vercel/bridge/pkg/plumbing" ) @@ -218,12 +219,12 @@ func (t *tunnelImpl) Start(ctx context.Context) { }() } -// openConn registers a pending connection for msg's ID before dialing. The +// openConn registers a buffered connection for msg's ID before dialing. The // peer has no explicit open message: the first chunk of a connection is what // opens it, and the next chunks can arrive before the dial completes. Storing -// the placeholder synchronously on the recv pump makes those chunks queue -// behind the first one instead of each dialing a second connection under the -// same ID and splitting the stream between them. +// the buffer synchronously on the recv pump makes those chunks queue behind +// the first one instead of each dialing a second connection under the same ID +// and splitting the stream between them. func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) { connID := msg.GetConnectionId() if msg.GetDest() == nil { @@ -231,15 +232,15 @@ func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) { return } - pending := &pendingConn{} - t.conns.Store(connID, pending) + buffered := &ioutil.BufferedReadWriteCloser{} + t.conns.Store(connID, buffered) if data := msg.GetData(); len(data) > 0 { - _, _ = pending.Write(data) + _, _ = buffered.Write(data) } - go t.dialPending(msg, pending) + go t.dial(msg, buffered) } -func (t *tunnelImpl) dialPending(msg *bridgev1.TunnelNetworkMessage, pending *pendingConn) { +func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered *ioutil.BufferedReadWriteCloser) { connID := msg.GetConnectionId() hostname := msg.GetHostname() dest := msg.GetDest() @@ -262,8 +263,8 @@ func (t *tunnelImpl) dialPending(msg *bridgev1.TunnelNetworkMessage, pending *pe if err != nil { slog.Info("Tunnel: connect failed", "conn_id", connID, "hostname", hostname, "error", err) - t.deletePending(connID, pending) - pending.Close() + t.deleteIfSame(connID, buffered) + buffered.Close() select { case t.sendCh <- &bridgev1.TunnelNetworkMessage{ ConnectionId: connID, @@ -274,24 +275,24 @@ func (t *tunnelImpl) dialPending(msg *bridgev1.TunnelNetworkMessage, pending *pe return } - if err := pending.resolve(conn); err != nil { + if err := buffered.Set(conn); err != nil { slog.Debug("Tunnel: dropping dialed connection", "connection_id", connID, "error", err) - t.deletePending(connID, pending) + t.deleteIfSame(connID, buffered) return } - go t.readFromConn(pending, connID, dest, msg.GetSource(), hostname) + go t.readFromConn(buffered, connID, dest, msg.GetSource(), hostname) } -// deletePending removes connID only while it still maps to pending, so a -// connection the peer has since reopened under the same ID is left alone. -func (t *tunnelImpl) deletePending(connID string, pending *pendingConn) { +// deleteIfSame removes connID only while it still maps to conn, so a connection +// the peer has since reopened under the same ID is left alone. +func (t *tunnelImpl) deleteIfSame(connID string, conn *ioutil.BufferedReadWriteCloser) { t.conns.Compute(connID, func(cur io.ReadWriteCloser, loaded bool) (io.ReadWriteCloser, bool) { if !loaded { return nil, true } - p, ok := cur.(*pendingConn) - return cur, ok && p == pending + b, ok := cur.(*ioutil.BufferedReadWriteCloser) + return cur, ok && b == conn }) } diff --git a/pkg/tunnel/tunnel_test.go b/pkg/tunnel/tunnel_test.go index b0c01db..6b51fab 100644 --- a/pkg/tunnel/tunnel_test.go +++ b/pkg/tunnel/tunnel_test.go @@ -220,7 +220,7 @@ func TestPeerErrorDuringDialClosesTheDialedConnection(t *testing.T) { stream.in <- dataMsg(connID, "hello") waitFor(t, "first dial", func() bool { return dialer.dials.Load() == 1 }) stream.in <- &bridgev1.TunnelNetworkMessage{ConnectionId: connID, Error: "peer closed"} - waitFor(t, "the peer error to drop the pending connection", func() bool { + waitFor(t, "the peer error to drop the buffered connection", func() bool { _, loaded := tun.(*tunnelImpl).conns.Load(connID) return !loaded }) From e84a57dc1592ecd2ccb5739f4cd31c8f4d0fdf83 Mon Sep 17 00:00:00 2001 From: edwardowens Date: Tue, 8 Sep 2026 13:45:27 -0700 Subject: [PATCH 3/3] refactor(ioutil): make BufferedReadWriteCloser an interface Co-Authored-By: Claude Fable 5 --- pkg/ioutil/buffered.go | 34 +++++++++++++++++++++------------- pkg/ioutil/buffered_test.go | 8 ++++---- pkg/tunnel/tunnel.go | 8 ++++---- 3 files changed, 29 insertions(+), 21 deletions(-) diff --git a/pkg/ioutil/buffered.go b/pkg/ioutil/buffered.go index f6e936b..57ac405 100644 --- a/pkg/ioutil/buffered.go +++ b/pkg/ioutil/buffered.go @@ -13,10 +13,23 @@ var ( ErrAlreadySet = errors.New("ioutil: underlying ReadWriteCloser already set") ) -// BufferedReadWriteCloser buffers writes until Set supplies the underlying -// io.ReadWriteCloser, then flushes them in order and passes everything -// through. The zero value is ready to use. -type BufferedReadWriteCloser struct { +// BufferedReadWriteCloser is an io.ReadWriteCloser that buffers writes until +// Set supplies the underlying one, then flushes them in order and passes +// everything through. +type BufferedReadWriteCloser interface { + io.ReadWriteCloser + + // Set attaches rwc and flushes the buffered writes to it. It closes rwc + // and returns ErrClosed if Close was already called, or the write error + // if the flush fails. + Set(rwc io.ReadWriteCloser) error +} + +func NewBufferedReadWriteCloser() BufferedReadWriteCloser { + return &bufferedReadWriteCloser{} +} + +type bufferedReadWriteCloser struct { // mu guards the handoff in Set: a Write that arrives while Set is flushing // the buffer must land after the flush, on the underlying rwc. mu sync.Mutex @@ -25,12 +38,7 @@ type BufferedReadWriteCloser struct { closed bool } -var _ io.ReadWriteCloser = (*BufferedReadWriteCloser)(nil) - -// Set attaches rwc and flushes the buffered writes to it. It closes rwc and -// returns ErrClosed if Close was already called, or the write error if the -// flush fails. -func (b *BufferedReadWriteCloser) Set(rwc io.ReadWriteCloser) error { +func (b *bufferedReadWriteCloser) Set(rwc io.ReadWriteCloser) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -53,7 +61,7 @@ func (b *BufferedReadWriteCloser) Set(rwc io.ReadWriteCloser) error { return nil } -func (b *BufferedReadWriteCloser) Write(p []byte) (int, error) { +func (b *bufferedReadWriteCloser) Write(p []byte) (int, error) { b.mu.Lock() if b.closed { b.mu.Unlock() @@ -69,7 +77,7 @@ func (b *BufferedReadWriteCloser) Write(p []byte) (int, error) { return rwc.Write(p) } -func (b *BufferedReadWriteCloser) Read(p []byte) (int, error) { +func (b *bufferedReadWriteCloser) Read(p []byte) (int, error) { b.mu.Lock() rwc, closed := b.rwc, b.closed b.mu.Unlock() @@ -82,7 +90,7 @@ func (b *BufferedReadWriteCloser) Read(p []byte) (int, error) { return rwc.Read(p) } -func (b *BufferedReadWriteCloser) Close() error { +func (b *bufferedReadWriteCloser) Close() error { b.mu.Lock() rwc := b.rwc b.closed = true diff --git a/pkg/ioutil/buffered_test.go b/pkg/ioutil/buffered_test.go index 4a45966..04c8f9c 100644 --- a/pkg/ioutil/buffered_test.go +++ b/pkg/ioutil/buffered_test.go @@ -28,7 +28,7 @@ func (r *recordingRWC) Close() error { } func TestWritesBeforeSetAreFlushedInOrder(t *testing.T) { - var b BufferedReadWriteCloser + b := NewBufferedReadWriteCloser() rwc := &recordingRWC{} for _, chunk := range []string{"one ", "two ", "three"} { @@ -65,7 +65,7 @@ func TestWritesBeforeSetAreFlushedInOrder(t *testing.T) { } func TestCloseBeforeSetClosesTheLateArrival(t *testing.T) { - var b BufferedReadWriteCloser + b := NewBufferedReadWriteCloser() _, _ = b.Write([]byte("queued")) if err := b.Close(); err != nil { t.Fatalf("Close: %v", err) @@ -87,7 +87,7 @@ func TestCloseBeforeSetClosesTheLateArrival(t *testing.T) { } func TestFlushFailureClosesTheUnderlying(t *testing.T) { - var b BufferedReadWriteCloser + b := NewBufferedReadWriteCloser() _, _ = b.Write([]byte("queued")) rwc := &recordingRWC{writeErr: io.ErrShortWrite} @@ -103,7 +103,7 @@ func TestFlushFailureClosesTheUnderlying(t *testing.T) { } func TestCloseAfterSetClosesTheUnderlying(t *testing.T) { - var b BufferedReadWriteCloser + b := NewBufferedReadWriteCloser() rwc := &recordingRWC{} if err := b.Set(rwc); err != nil { t.Fatalf("Set: %v", err) diff --git a/pkg/tunnel/tunnel.go b/pkg/tunnel/tunnel.go index 1659631..4755f22 100644 --- a/pkg/tunnel/tunnel.go +++ b/pkg/tunnel/tunnel.go @@ -232,7 +232,7 @@ func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) { return } - buffered := &ioutil.BufferedReadWriteCloser{} + buffered := ioutil.NewBufferedReadWriteCloser() t.conns.Store(connID, buffered) if data := msg.GetData(); len(data) > 0 { _, _ = buffered.Write(data) @@ -240,7 +240,7 @@ func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) { go t.dial(msg, buffered) } -func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered *ioutil.BufferedReadWriteCloser) { +func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered ioutil.BufferedReadWriteCloser) { connID := msg.GetConnectionId() hostname := msg.GetHostname() dest := msg.GetDest() @@ -286,12 +286,12 @@ func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered *ioutil.B // deleteIfSame removes connID only while it still maps to conn, so a connection // the peer has since reopened under the same ID is left alone. -func (t *tunnelImpl) deleteIfSame(connID string, conn *ioutil.BufferedReadWriteCloser) { +func (t *tunnelImpl) deleteIfSame(connID string, conn ioutil.BufferedReadWriteCloser) { t.conns.Compute(connID, func(cur io.ReadWriteCloser, loaded bool) (io.ReadWriteCloser, bool) { if !loaded { return nil, true } - b, ok := cur.(*ioutil.BufferedReadWriteCloser) + b, ok := cur.(ioutil.BufferedReadWriteCloser) return cur, ok && b == conn }) }