diff --git a/core/listener/forwardproxy/connect_hang_test.go b/core/listener/forwardproxy/connect_hang_test.go new file mode 100644 index 000000000..7a6803667 --- /dev/null +++ b/core/listener/forwardproxy/connect_hang_test.go @@ -0,0 +1,247 @@ +package forwardproxy + +import ( + "bufio" + "io" + "net" + "strings" + "testing" + "time" + + "github.com/rossoctl/cortex/core/session" +) + +// A server-first protocol (SSH, SMTP, FTP) sends a banner before the client +// writes anything. Through CONNECT with the TLS bridge on, the proxy peeks +// five bytes of the client to classify TLS vs opaque — the client sends +// nothing until it sees the banner, so both sides wait on each other until +// the sniff deadline elapses. +func TestHandleConnect_ServerFirstBannerArrivesPromptly(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + go func() { + c, err := ln.Accept() + if err != nil { + return + } + _, _ = c.Write([]byte("SSH-2.0-banner\r\n")) + }() + target := ln.Addr().String() + + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + s := connectServer(t, store, bridgeEngine(t, portOf(target), nil)) + proxyAddr, _ := connectProxy(t, s) + + raw, br, resp := sendConnect(t, proxyAddr, target) + if resp.StatusCode != 200 { + t.Fatalf("CONNECT = %d", resp.StatusCode) + } + // Pre-fix the sniff peeks the client with no deadline and the client sends + // nothing, so the tunnel never starts and the banner never flows. Post-fix + // the peek is bound by connectSniffTimeout, after which the opaque path + // takes over. + _ = raw.SetReadDeadline(time.Now().Add(connectSniffTimeout + 2*time.Second)) + line, err := br.ReadString('\n') + if err != nil { + t.Fatalf("read banner: %v", err) + } + if !strings.HasPrefix(line, "SSH-2.0-") { + t.Fatalf("got %q, want an SSH banner", line) + } +} + +// Half-close: a client that sends its request, calls CloseWrite, then reads +// the reply must receive it. Today tunnel closes both halves as soon as the +// client→upstream copy sees EOF. +func TestTunnel_HalfCloseRelaysTheReply(t *testing.T) { + // Two localhost TCP pairs so CloseWrite is available on both sides. + cl, clPeer := tcpPair(t) + up, upPeer := tcpPair(t) + + // Upstream: read to EOF, then write a reply, then close. + upDone := make(chan struct{}) + go func() { + defer close(upDone) + _, _ = io.ReadAll(upPeer) + _, _ = upPeer.Write([]byte("REPLY")) + _ = upPeer.Close() + }() + + // Client: write the request, half-close send, read reply. + reply := make(chan []byte, 1) + go func() { + _, _ = clPeer.Write([]byte("hello")) + _ = clPeer.(*net.TCPConn).CloseWrite() + b, _ := io.ReadAll(clPeer) + reply <- b + }() + + done := make(chan struct{}) + go func() { + tunnel(cl, up) + close(done) + }() + + select { + case got := <-reply: + if string(got) != "REPLY" { + t.Fatalf("reply = %q, want the upstream's answer", got) + } + case <-time.After(2 * time.Second): + t.Fatal("client never got the reply; half-close tore the tunnel down") + } + <-done + <-upDone +} + +// A wrapped client conn (as the TLS-bridge CONNECT path uses) must still +// deliver the half-close through tunnel, since peekedConn alone embeds +// net.Conn and would otherwise miss the CloseWrite type assertion. +func TestTunnel_HalfCloseThroughPeekedConn(t *testing.T) { + cl, clPeer := tcpPair(t) + up, upPeer := tcpPair(t) + + upDone := make(chan struct{}) + go func() { + defer close(upDone) + _, _ = io.ReadAll(upPeer) + _, _ = upPeer.Write([]byte("REPLY")) + _ = upPeer.Close() + }() + + reply := make(chan []byte, 1) + go func() { + _, _ = clPeer.Write([]byte("hello")) + _ = clPeer.(*net.TCPConn).CloseWrite() + b, _ := io.ReadAll(clPeer) + reply <- b + }() + + wrapped := &peekedConn{Conn: cl, r: bufio.NewReader(cl)} + done := make(chan struct{}) + go func() { + tunnel(wrapped, up) + close(done) + }() + + select { + case got := <-reply: + if string(got) != "REPLY" { + t.Fatalf("reply = %q", got) + } + case <-time.After(2 * time.Second): + t.Fatal("wrapped client never got the reply") + } + <-done + <-upDone +} + +// If the client half-closes but the upstream never does, tunnel still has to +// return so the close row can settle. Pre-fix the downstream io.Copy blocked +// forever; post-fix halfCloseDrainVar bounds the wait. +func TestTunnel_ReturnsWhenUpstreamNeverHalfCloses(t *testing.T) { + cl, clPeer := tcpPair(t) + up, _ := tcpPair(t) // upstream peer kept alive by tcpPair's cleanup; never writes EOF + + go func() { + _, _ = clPeer.Write([]byte("hello")) + _ = clPeer.(*net.TCPConn).CloseWrite() + }() + + start := time.Now() + done := make(chan struct{}) + go func() { + tunnel(cl, up) + close(done) + }() + select { + case <-done: + case <-time.After(halfCloseDrainVar + 2*time.Second): + t.Fatalf("tunnel did not return after %v", time.Since(start)) + } +} + +// A reply that keeps streaming past halfCloseDrainVar is not truncated: the +// drain extends while bytes continue flowing and only full-closes once the +// peer actually goes idle. +func TestTunnel_ReplyStreamingPastDrainCompletes(t *testing.T) { + cl, clPeer := tcpPair(t) + up, upPeer := tcpPair(t) + + // Upstream: read to EOF, then drip-feed a reply longer than halfCloseDrainVar. + chunks := 6 + gap := halfCloseDrainVar / 2 // < halfCloseDrainVar, so each gap extends the drain + wantTotal := chunks * 5 // "chunk" × 5 bytes + upDone := make(chan struct{}) + go func() { + defer close(upDone) + _, _ = io.ReadAll(upPeer) + for range chunks { + if _, err := upPeer.Write([]byte("chunk")); err != nil { + return + } + time.Sleep(gap) + } + _ = upPeer.Close() + }() + + reply := make(chan []byte, 1) + go func() { + _, _ = clPeer.Write([]byte("hello")) + _ = clPeer.(*net.TCPConn).CloseWrite() + b, _ := io.ReadAll(clPeer) + reply <- b + }() + + done := make(chan struct{}) + go func() { + tunnel(cl, up) + close(done) + }() + + deadline := time.Duration(chunks)*gap + halfCloseDrainVar + 2*time.Second + select { + case got := <-reply: + if len(got) != wantTotal { + t.Fatalf("reply = %d bytes, want %d; drain truncated a slow stream", len(got), wantTotal) + } + case <-time.After(deadline): + t.Fatal("tunnel never returned a complete reply") + } + <-done + <-upDone +} + +// tcpPair returns two ends of a localhost TCP connection as *net.TCPConn, +// so CloseWrite is available on both. +func tcpPair(t *testing.T) (a, b net.Conn) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + type accepted struct { + c net.Conn + err error + } + ch := make(chan accepted, 1) + go func() { + c, err := ln.Accept() + ch <- accepted{c, err} + }() + dialer, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + got := <-ch + if got.err != nil { + t.Fatal(got.err) + } + t.Cleanup(func() { _ = dialer.Close(); _ = got.c.Close() }) + return dialer, got.c +} diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 59bce0e2e..f19920d9e 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -2789,7 +2789,14 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { if s.TLSBridge != nil { pc := &peekedConn{Conn: clientConn, r: bufio.NewReaderSize(clientConn, sniffBufSize)} clientConn = pc // replay peeked bytes into whichever path runs + // A server-first protocol (SSH, SMTP) sends nothing until it reads the + // banner, so bound the peek: a short or empty first classifies as + // not-TLS and the tunnel takes the opaque path. Shorter than + // sniffHost's 5s because this delay is user-visible — the banner only + // flows once this returns. + _ = clientConn.SetReadDeadline(time.Now().Add(connectSniffTimeout)) first, _ := pc.Peek(5) + _ = clientConn.SetReadDeadline(time.Time{}) authority := r.Host // CONNECT target is already host:port key := hostOnly(r.Host) var bridge bool diff --git a/core/listener/forwardproxy/sniff.go b/core/listener/forwardproxy/sniff.go index 4c21d5bd5..e53f25952 100644 --- a/core/listener/forwardproxy/sniff.go +++ b/core/listener/forwardproxy/sniff.go @@ -27,6 +27,10 @@ const ( // (or a server-first protocol) can't pin a goroutine. Only relevant on the // sniffed ports, where the client speaks first, so it is rarely hit. sniffTimeout = 5 * time.Second + // connectSniffTimeout bounds the TLS-bridge classify peek on a CONNECT + // tunnel. The banner of a server-first protocol only flows after this + // returns, so it is shorter than sniffTimeout to keep the delay unnoticeable. + connectSniffTimeout = 500 * time.Millisecond ) // errSniffDone aborts the throwaway TLS handshake once we have the SNI. @@ -152,6 +156,15 @@ type peekedConn struct { func (c *peekedConn) Read(p []byte) (int, error) { return c.r.Read(p) } +// CloseWrite forwards to the embedded conn when it supports half-close, so +// tunnel's half-close relay works when the client was wrapped for sniffing. +func (c *peekedConn) CloseWrite() error { + if cw, ok := c.Conn.(interface{ CloseWrite() error }); ok { + return cw.CloseWrite() + } + return c.Conn.Close() +} + // Peek returns the next n buffered bytes without consuming them. Used by the // TLS-bridge classify step on the CONNECT and transparent paths. func (c *peekedConn) Peek(n int) ([]byte, error) { return c.r.Peek(n) } diff --git a/core/listener/forwardproxy/transparent.go b/core/listener/forwardproxy/transparent.go index 151a242a3..7b73568bd 100644 --- a/core/listener/forwardproxy/transparent.go +++ b/core/listener/forwardproxy/transparent.go @@ -7,6 +7,7 @@ import ( "log/slog" "net" "net/http" + "sync/atomic" "time" "github.com/rossoctl/cortex/core/listener/egressguard" @@ -394,25 +395,95 @@ func (s *Server) recordTunnelClosed(pctx *pipeline.Context, b *session.Bucket, r }) } -// tunnel bidirectionally copies between two connections until either side -// closes, then propagates the close to the other so both io.Copy goroutines -// exit. Close-on-each-side is idempotent on net.Conn. Shared by handleConnect -// and HandleTransparentConn. +// tunnel bidirectionally copies between two connections, relaying each +// direction's EOF as a CloseWrite on the receiving side so a half-close +// survives. Shared by handleConnect and HandleTransparentConn. // -// It returns how many bytes crossed each way — up is client to upstream, down is -// upstream to client — for the tunnel's close row. It waits for the other copy to -// finish before returning, which it did not need to while nothing read its count: -// closing both ends first is what makes that wait short. +// It returns how many bytes crossed each way — up is client to upstream, down +// is upstream to client — for the tunnel's close row. After the first copy +// returns, the second gets halfCloseDrainVar of idle time before both ends are +// full-closed, so a reply that keeps streaming extends the wait and finishes +// without truncation. func tunnel(client, upstream net.Conn) (up, down int64) { + var upBytes, downBytes atomic.Int64 upc := make(chan int64, 1) + downc := make(chan int64, 1) go func() { - n, _ := io.Copy(upstream, client) - _ = upstream.Close() - _ = client.Close() + n, _ := io.Copy(&countingWriter{w: upstream, n: &upBytes}, client) + closeWrite(upstream) // relay the client's half-close; let the reply flow back upc <- n }() - down, _ = io.Copy(client, upstream) + go func() { + n, _ := io.Copy(&countingWriter{w: client, n: &downBytes}, upstream) + closeWrite(client) + downc <- n + }() + select { + case up = <-upc: + down = drainTunnel(&downBytes, downc, client, upstream) + case down = <-downc: + up = drainTunnel(&upBytes, upc, client, upstream) + } _ = client.Close() _ = upstream.Close() - return <-upc, down + return up, down +} + +// drainTunnel waits for the second copy, extending the wait while bytes keep +// flowing (seen via n). Idle for halfCloseDrainVar → full-close both ends to +// unblock the copy, read its total, return. +func drainTunnel(n *atomic.Int64, done <-chan int64, client, upstream net.Conn) int64 { + last := n.Load() + tick := time.NewTicker(halfCloseDrainVar / 4) + defer tick.Stop() + idle := time.Duration(0) + for { + select { + case total := <-done: + return total + case <-tick.C: + now := n.Load() + if now != last { + last, idle = now, 0 + continue + } + idle += halfCloseDrainVar / 4 + if idle >= halfCloseDrainVar { + _ = client.Close() + _ = upstream.Close() + return <-done + } + } + } +} + +// countingWriter forwards writes to w and records the running byte count in n. +// Used by tunnel so drainTunnel can see progress mid-copy. +type countingWriter struct { + w io.Writer + n *atomic.Int64 +} + +func (c *countingWriter) Write(p []byte) (int, error) { + m, err := c.w.Write(p) + c.n.Add(int64(m)) + return m, err +} + +// halfCloseDrainVar is how long tunnel waits for the second copy after the +// first has returned. Short enough to settle the tunnel's close row promptly, +// long enough to carry a reply the peer is actively writing. A var so a test +// can shrink it; a test that mutates it relies on package tests running +// sequentially, since tunnel reads it on every teardown. +var halfCloseDrainVar = 2 * time.Second + +// closeWrite shuts down c's write side when the underlying conn supports it +// (TCPConn, UnixConn), so the peer sees EOF on its reads without losing the +// other direction. Falls back to a full close for anything that does not. +func closeWrite(c net.Conn) { + if cw, ok := c.(interface{ CloseWrite() error }); ok { + _ = cw.CloseWrite() + return + } + _ = c.Close() }