From ba5c55f1a9ea0c138bedfe84d154ae843ac7d97f Mon Sep 17 00:00:00 2001
From: breken <hello@breken.ai>
Date: Sun, 13 Sep 2026 19:28:18 -0700
Subject: [PATCH] fix(proxy): keep bytes after PROXY UNKNOWN header

When a trusted upstream sends "PROXY UNKNOWN\r\n", parseProxyProtocolHeader
returned the raw connection and discarded whatever followed the header in
the same read. The TLS ClientHello usually arrives in that same segment, so
it was lost, the SNI extraction failed on the truncated stream, and the
connection was dropped. The 5s parsing read deadline was also left set on
this path.

Wrap the connection so the remaining buffered bytes are replayed ahead of
the socket, and clear the read deadline, matching the other header
branches.
---
 proxy/proxy.go      | 15 ++++++++--
 proxy/proxy_test.go | 68 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 94d6888..3df4caa 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -213,8 +213,19 @@ func (p *SNIProxy) parseProxyProtocolHeader(conn net.Conn) (*ProxyProtocolInfo,
 	if len(parts) != 6 || parts[0] != "PROXY" {
 		// Check for PROXY UNKNOWN
 		if len(parts) == 2 && parts[0] == "PROXY" && parts[1] == "UNKNOWN" {
-			// PROXY UNKNOWN - use original connection info
-			return nil, conn, nil
+			// PROXY UNKNOWN - use original connection info, but keep any
+			// bytes that arrived after the header (the TLS ClientHello).
+			if err := conn.SetReadDeadline(time.Time{}); err != nil {
+				return nil, conn, fmt.Errorf("failed to clear read deadline: %w", err)
+			}
+			if len(remainingData) == 0 {
+				return nil, conn, nil
+			}
+			wrappedConn := &proxyProtocolConn{
+				Conn:   conn,
+				reader: io.MultiReader(bytes.NewReader(remainingData), conn),
+			}
+			return nil, wrappedConn, nil
 		}
 		// Invalid PROXY protocol, but might be regular TLS - treat as such
 		logger.Debug("Invalid PROXY protocol from trusted upstream %s, treating as regular TLS connection: %s", remoteHost, headerLine)
diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go
index 747c81d..7e3e6e6 100644
--- a/proxy/proxy_test.go
+++ b/proxy/proxy_test.go
@@ -1,8 +1,11 @@
 package proxy
 
 import (
+	"bytes"
+	"io"
 	"net"
 	"testing"
+	"time"
 )
 
 func TestBuildProxyProtocolHeader(t *testing.T) {
@@ -117,3 +120,68 @@ func TestBuildProxyProtocolHeaderFromInfo(t *testing.T) {
 		t.Errorf("Expected header '%s', got '%s'", expected, header)
 	}
 }
+
+// TestParseProxyProtocolHeaderUnknownPreservesPayload checks that when a
+// trusted upstream sends "PROXY UNKNOWN\r\n" the bytes that follow the header
+// (the TLS ClientHello) are still readable from the returned connection.
+func TestParseProxyProtocolHeaderUnknownPreservesPayload(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer listener.Close()
+
+	proxy, err := NewSNIProxy(8443, "", "", "127.0.0.1", 443, nil, false, []string{"127.0.0.1"})
+	if err != nil {
+		t.Fatalf("Failed to create SNI proxy: %v", err)
+	}
+
+	payload := []byte("\x16\x03\x01client-hello-bytes")
+	accepted := make(chan net.Conn, 1)
+	go func() {
+		conn, err := listener.Accept()
+		if err != nil {
+			t.Errorf("Accept failed: %v", err)
+			accepted <- nil
+			return
+		}
+		accepted <- conn
+	}()
+
+	client, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatalf("Failed to dial: %v", err)
+	}
+	defer client.Close()
+
+	// Header and payload arrive in a single segment, as they do when the
+	// upstream writes both before the first flush.
+	if _, err := client.Write(append([]byte("PROXY UNKNOWN\r\n"), payload...)); err != nil {
+		t.Fatalf("Failed to write: %v", err)
+	}
+
+	serverConn := <-accepted
+	if serverConn == nil {
+		t.FailNow()
+	}
+	defer serverConn.Close()
+
+	proxyInfo, wrapped, err := proxy.parseProxyProtocolHeader(serverConn)
+	if err != nil {
+		t.Fatalf("parseProxyProtocolHeader returned error: %v", err)
+	}
+	if proxyInfo != nil {
+		t.Fatalf("Expected nil proxyInfo for PROXY UNKNOWN, got %+v", proxyInfo)
+	}
+
+	if err := wrapped.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
+		t.Fatalf("Failed to set read deadline: %v", err)
+	}
+	got := make([]byte, len(payload))
+	if _, err := io.ReadFull(wrapped, got); err != nil {
+		t.Fatalf("Payload after PROXY UNKNOWN header was lost: %v", err)
+	}
+	if !bytes.Equal(got, payload) {
+		t.Fatalf("Expected payload %q, got %q", payload, got)
+	}
+}