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) + } +}