Skip to content

Commit 678df96

Browse files
vercel-eddieclaude
vercel-eddie
and
claude
authored
fix(tunnel): queue chunks that arrive mid-dial instead of dialing again (#14)
* 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * refactor(ioutil): make BufferedReadWriteCloser an interface Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent 4e43b48 commit 678df96

4 files changed

Lines changed: 518 additions & 19 deletions

File tree

‎pkg/ioutil/buffered.go‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
package ioutil
2+
3+
import (
4+
"bytes"
5+
"errors"
6+
"io"
7+
"sync"
8+
)
9+
10+
var (
11+
ErrClosed = errors.New("ioutil: closed")
12+
ErrNotSet = errors.New("ioutil: underlying ReadWriteCloser not set")
13+
ErrAlreadySet = errors.New("ioutil: underlying ReadWriteCloser already set")
14+
)
15+
16+
// BufferedReadWriteCloser is an io.ReadWriteCloser that buffers writes until
17+
// Set supplies the underlying one, then flushes them in order and passes
18+
// everything through.
19+
type BufferedReadWriteCloser interface {
20+
io.ReadWriteCloser
21+
22+
// Set attaches rwc and flushes the buffered writes to it. It closes rwc
23+
// and returns ErrClosed if Close was already called, or the write error
24+
// if the flush fails.
25+
Set(rwc io.ReadWriteCloser) error
26+
}
27+
28+
func NewBufferedReadWriteCloser() BufferedReadWriteCloser {
29+
return &bufferedReadWriteCloser{}
30+
}
31+
32+
type bufferedReadWriteCloser struct {
33+
// mu guards the handoff in Set: a Write that arrives while Set is flushing
34+
// the buffer must land after the flush, on the underlying rwc.
35+
mu sync.Mutex
36+
rwc io.ReadWriteCloser
37+
buf bytes.Buffer
38+
closed bool
39+
}
40+
41+
func (b *bufferedReadWriteCloser) Set(rwc io.ReadWriteCloser) error {
42+
b.mu.Lock()
43+
defer b.mu.Unlock()
44+
if b.closed {
45+
rwc.Close()
46+
return ErrClosed
47+
}
48+
if b.rwc != nil {
49+
return ErrAlreadySet
50+
}
51+
if b.buf.Len() > 0 {
52+
if _, err := rwc.Write(b.buf.Bytes()); err != nil {
53+
rwc.Close()
54+
b.closed = true
55+
b.buf.Reset()
56+
return err
57+
}
58+
b.buf.Reset()
59+
}
60+
b.rwc = rwc
61+
return nil
62+
}
63+
64+
func (b *bufferedReadWriteCloser) Write(p []byte) (int, error) {
65+
b.mu.Lock()
66+
if b.closed {
67+
b.mu.Unlock()
68+
return 0, ErrClosed
69+
}
70+
if b.rwc == nil {
71+
n, err := b.buf.Write(p)
72+
b.mu.Unlock()
73+
return n, err
74+
}
75+
rwc := b.rwc
76+
b.mu.Unlock()
77+
return rwc.Write(p)
78+
}
79+
80+
func (b *bufferedReadWriteCloser) Read(p []byte) (int, error) {
81+
b.mu.Lock()
82+
rwc, closed := b.rwc, b.closed
83+
b.mu.Unlock()
84+
if closed {
85+
return 0, ErrClosed
86+
}
87+
if rwc == nil {
88+
return 0, ErrNotSet
89+
}
90+
return rwc.Read(p)
91+
}
92+
93+
func (b *bufferedReadWriteCloser) Close() error {
94+
b.mu.Lock()
95+
rwc := b.rwc
96+
b.closed = true
97+
b.buf.Reset()
98+
b.mu.Unlock()
99+
if rwc != nil {
100+
return rwc.Close()
101+
}
102+
return nil
103+
}

‎pkg/ioutil/buffered_test.go‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
package ioutil
2+
3+
import (
4+
"bytes"
5+
"errors"
6+
"io"
7+
"testing"
8+
)
9+
10+
type recordingRWC struct {
11+
written bytes.Buffer
12+
closed bool
13+
writeErr error
14+
}
15+
16+
func (r *recordingRWC) Write(p []byte) (int, error) {
17+
if r.writeErr != nil {
18+
return 0, r.writeErr
19+
}
20+
return r.written.Write(p)
21+
}
22+
23+
func (r *recordingRWC) Read(p []byte) (int, error) { return copy(p, "read"), nil }
24+
25+
func (r *recordingRWC) Close() error {
26+
r.closed = true
27+
return nil
28+
}
29+
30+
func TestWritesBeforeSetAreFlushedInOrder(t *testing.T) {
31+
b := NewBufferedReadWriteCloser()
32+
rwc := &recordingRWC{}
33+
34+
for _, chunk := range []string{"one ", "two ", "three"} {
35+
if _, err := b.Write([]byte(chunk)); err != nil {
36+
t.Fatalf("Write before Set: %v", err)
37+
}
38+
}
39+
if _, err := b.Read(make([]byte, 4)); !errors.Is(err, ErrNotSet) {
40+
t.Fatalf("Read before Set: got %v, want ErrNotSet", err)
41+
}
42+
43+
if err := b.Set(rwc); err != nil {
44+
t.Fatalf("Set: %v", err)
45+
}
46+
if got := rwc.written.String(); got != "one two three" {
47+
t.Fatalf("flushed %q, want %q", got, "one two three")
48+
}
49+
50+
if _, err := b.Write([]byte(" four")); err != nil {
51+
t.Fatalf("Write after Set: %v", err)
52+
}
53+
if got := rwc.written.String(); got != "one two three four" {
54+
t.Fatalf("after write-through got %q", got)
55+
}
56+
57+
p := make([]byte, 4)
58+
if n, err := b.Read(p); err != nil || string(p[:n]) != "read" {
59+
t.Fatalf("Read after Set: %q, %v", p[:n], err)
60+
}
61+
62+
if err := b.Set(&recordingRWC{}); !errors.Is(err, ErrAlreadySet) {
63+
t.Fatalf("second Set: got %v, want ErrAlreadySet", err)
64+
}
65+
}
66+
67+
func TestCloseBeforeSetClosesTheLateArrival(t *testing.T) {
68+
b := NewBufferedReadWriteCloser()
69+
_, _ = b.Write([]byte("queued"))
70+
if err := b.Close(); err != nil {
71+
t.Fatalf("Close: %v", err)
72+
}
73+
74+
rwc := &recordingRWC{}
75+
if err := b.Set(rwc); !errors.Is(err, ErrClosed) {
76+
t.Fatalf("Set after Close: got %v, want ErrClosed", err)
77+
}
78+
if !rwc.closed {
79+
t.Fatal("expected Set to close the ReadWriteCloser handed to a closed buffer")
80+
}
81+
if rwc.written.Len() != 0 {
82+
t.Fatalf("expected nothing flushed to a closed buffer's late arrival, got %q", rwc.written.String())
83+
}
84+
if _, err := b.Write([]byte("x")); !errors.Is(err, ErrClosed) {
85+
t.Fatalf("Write after Close: got %v, want ErrClosed", err)
86+
}
87+
}
88+
89+
func TestFlushFailureClosesTheUnderlying(t *testing.T) {
90+
b := NewBufferedReadWriteCloser()
91+
_, _ = b.Write([]byte("queued"))
92+
93+
rwc := &recordingRWC{writeErr: io.ErrShortWrite}
94+
if err := b.Set(rwc); !errors.Is(err, io.ErrShortWrite) {
95+
t.Fatalf("Set with failing flush: got %v, want ErrShortWrite", err)
96+
}
97+
if !rwc.closed {
98+
t.Fatal("expected the underlying to be closed after a failed flush")
99+
}
100+
if _, err := b.Write([]byte("x")); !errors.Is(err, ErrClosed) {
101+
t.Fatalf("Write after failed flush: got %v, want ErrClosed", err)
102+
}
103+
}
104+
105+
func TestCloseAfterSetClosesTheUnderlying(t *testing.T) {
106+
b := NewBufferedReadWriteCloser()
107+
rwc := &recordingRWC{}
108+
if err := b.Set(rwc); err != nil {
109+
t.Fatalf("Set: %v", err)
110+
}
111+
if err := b.Close(); err != nil {
112+
t.Fatalf("Close: %v", err)
113+
}
114+
if !rwc.closed {
115+
t.Fatal("expected Close to close the underlying")
116+
}
117+
}

‎pkg/tunnel/tunnel.go‎

Lines changed: 47 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111

1212
"github.com/puzpuzpuz/xsync/v3"
1313
bridgev1 "github.com/vercel/bridge/api/go/bridge/v1"
14+
"github.com/vercel/bridge/pkg/ioutil"
1415
"github.com/vercel/bridge/pkg/mitm"
1516
"github.com/vercel/bridge/pkg/plumbing"
1617
)
@@ -53,7 +54,7 @@ func New(dialer plumbing.ContextDialer, stream Stream, opts ...Option) Tunnel {
5354
dialer: dialer,
5455
stream: stream,
5556
sendCh: make(chan *bridgev1.TunnelNetworkMessage, 64),
56-
conns: xsync.NewMapOf[string, net.Conn](),
57+
conns: xsync.NewMapOf[string, io.ReadWriteCloser](),
5758
done: make(chan struct{}),
5859
}
5960
for _, o := range opts {
@@ -67,7 +68,7 @@ type tunnelImpl struct {
6768
hijacker mitm.Hijacker
6869
stream Stream
6970
sendCh chan *bridgev1.TunnelNetworkMessage
70-
conns *xsync.MapOf[string, net.Conn]
71+
conns *xsync.MapOf[string, io.ReadWriteCloser]
7172
done chan struct{}
7273
ctx context.Context
7374
cancel context.CancelFunc
@@ -99,8 +100,8 @@ func (t *tunnelImpl) AddConn(conn net.Conn, destOverride string, hostname string
99100
go t.readFromConn(conn, connID, src, dst, hostname)
100101
}
101102

102-
// readFromConn reads from a net.Conn and forwards data to the stream via sendCh.
103-
func (t *tunnelImpl) readFromConn(conn net.Conn, connID string, src, dst *bridgev1.TunnelAddress, hostname string) {
103+
// readFromConn reads from a connection and forwards data to the stream via sendCh.
104+
func (t *tunnelImpl) readFromConn(conn io.ReadCloser, connID string, src, dst *bridgev1.TunnelAddress, hostname string) {
104105
defer func() {
105106
conn.Close()
106107
t.conns.Delete(connID)
@@ -200,8 +201,8 @@ func (t *tunnelImpl) Start(ctx context.Context) {
200201
continue
201202
}
202203

203-
// Unknown connection ID → dial via the configured dialer.
204-
go t.handleNewConn(msg)
204+
// Unknown connection ID → this message opens it.
205+
t.openConn(msg)
205206

206207
case err := <-recvErr:
207208
if err != io.EOF {
@@ -218,15 +219,31 @@ func (t *tunnelImpl) Start(ctx context.Context) {
218219
}()
219220
}
220221

221-
func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) {
222-
dest := msg.GetDest()
223-
if dest == nil {
224-
slog.Info("Tunnel: ignoring message with no dest", "conn_id", msg.GetConnectionId())
222+
// openConn registers a buffered connection for msg's ID before dialing. The
223+
// peer has no explicit open message: the first chunk of a connection is what
224+
// opens it, and the next chunks can arrive before the dial completes. Storing
225+
// the buffer synchronously on the recv pump makes those chunks queue behind
226+
// the first one instead of each dialing a second connection under the same ID
227+
// and splitting the stream between them.
228+
func (t *tunnelImpl) openConn(msg *bridgev1.TunnelNetworkMessage) {
229+
connID := msg.GetConnectionId()
230+
if msg.GetDest() == nil {
231+
slog.Info("Tunnel: ignoring message with no dest", "conn_id", connID)
225232
return
226233
}
227234

235+
buffered := ioutil.NewBufferedReadWriteCloser()
236+
t.conns.Store(connID, buffered)
237+
if data := msg.GetData(); len(data) > 0 {
238+
_, _ = buffered.Write(data)
239+
}
240+
go t.dial(msg, buffered)
241+
}
242+
243+
func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered ioutil.BufferedReadWriteCloser) {
228244
connID := msg.GetConnectionId()
229245
hostname := msg.GetHostname()
246+
dest := msg.GetDest()
230247

231248
// Resolve the connection: hijacker gets first shot, then fall through to the dialer.
232249
var conn net.Conn
@@ -246,6 +263,8 @@ func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) {
246263

247264
if err != nil {
248265
slog.Info("Tunnel: connect failed", "conn_id", connID, "hostname", hostname, "error", err)
266+
t.deleteIfSame(connID, buffered)
267+
buffered.Close()
249268
select {
250269
case t.sendCh <- &bridgev1.TunnelNetworkMessage{
251270
ConnectionId: connID,
@@ -256,20 +275,29 @@ func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) {
256275
return
257276
}
258277

259-
t.conns.Store(connID, conn)
260-
go t.readFromConn(conn, connID, msg.GetDest(), msg.GetSource(), hostname)
278+
if err := buffered.Set(conn); err != nil {
279+
slog.Debug("Tunnel: dropping dialed connection", "connection_id", connID, "error", err)
280+
t.deleteIfSame(connID, buffered)
281+
return
282+
}
261283

262-
if data := msg.GetData(); len(data) > 0 {
263-
if _, err := conn.Write(data); err != nil {
264-
slog.Debug("Failed to write initial data", "connection_id", connID, "error", err)
265-
conn.Close()
266-
t.conns.Delete(connID)
284+
go t.readFromConn(buffered, connID, dest, msg.GetSource(), hostname)
285+
}
286+
287+
// deleteIfSame removes connID only while it still maps to conn, so a connection
288+
// the peer has since reopened under the same ID is left alone.
289+
func (t *tunnelImpl) deleteIfSame(connID string, conn ioutil.BufferedReadWriteCloser) {
290+
t.conns.Compute(connID, func(cur io.ReadWriteCloser, loaded bool) (io.ReadWriteCloser, bool) {
291+
if !loaded {
292+
return nil, true
267293
}
268-
}
294+
b, ok := cur.(ioutil.BufferedReadWriteCloser)
295+
return cur, ok && b == conn
296+
})
269297
}
270298

271299
func (t *tunnelImpl) closeAll() {
272-
t.conns.Range(func(key string, conn net.Conn) bool {
300+
t.conns.Range(func(key string, conn io.ReadWriteCloser) bool {
273301
conn.Close()
274302
t.conns.Delete(key)
275303
return true

0 commit comments

Comments
 (0)