diff --git a/pkg/ioutil/buffered.go b/pkg/ioutil/buffered.go new file mode 100644 index 0000000..57ac405 --- /dev/null +++ b/pkg/ioutil/buffered.go @@ -0,0 +1,103 @@ +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 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 + rwc io.ReadWriteCloser + buf bytes.Buffer + closed bool +} + +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..04c8f9c --- /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) { + b := NewBufferedReadWriteCloser() + 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) { + b := NewBufferedReadWriteCloser() + _, _ = 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) { + b := NewBufferedReadWriteCloser() + _, _ = 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) { + b := NewBufferedReadWriteCloser() + 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/tunnel.go b/pkg/tunnel/tunnel.go index d280daa..4755f22 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" ) @@ -53,7 +54,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 +68,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 +100,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 +201,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 +219,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 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 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 { + slog.Info("Tunnel: ignoring message with no dest", "conn_id", connID) return } + buffered := ioutil.NewBufferedReadWriteCloser() + t.conns.Store(connID, buffered) + if data := msg.GetData(); len(data) > 0 { + _, _ = buffered.Write(data) + } + go t.dial(msg, buffered) +} + +func (t *tunnelImpl) dial(msg *bridgev1.TunnelNetworkMessage, buffered ioutil.BufferedReadWriteCloser) { 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 +263,8 @@ func (t *tunnelImpl) handleNewConn(msg *bridgev1.TunnelNetworkMessage) { if err != nil { slog.Info("Tunnel: connect failed", "conn_id", connID, "hostname", hostname, "error", err) + t.deleteIfSame(connID, buffered) + buffered.Close() select { case t.sendCh <- &bridgev1.TunnelNetworkMessage{ ConnectionId: connID, @@ -256,20 +275,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 := buffered.Set(conn); err != nil { + slog.Debug("Tunnel: dropping dialed connection", "connection_id", connID, "error", err) + t.deleteIfSame(connID, buffered) + 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(buffered, connID, dest, msg.GetSource(), hostname) +} + +// 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 } - } + b, ok := cur.(ioutil.BufferedReadWriteCloser) + return cur, ok && b == conn + }) } 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..6b51fab --- /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 buffered 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 }) +}