From b2153f89cb19bf7cb995a5270bd983717800c066 Mon Sep 17 00:00:00 2001
From: Geoff Bourne <itzgeoff@gmail.com>
Date: Sat, 3 Feb 2024 19:37:44 -0600
Subject: [PATCH] Corrected addresses used in PROXY header (#262)

---
 server/connector.go | 31 +++++++++++++++++++------------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/server/connector.go b/server/connector.go
index a8437fa..4089770 100644
--- a/server/connector.go
+++ b/server/connector.go
@@ -7,7 +7,6 @@ import (
 	"golang.ngrok.com/ngrok/config"
 	"io"
 	"net"
-	"strconv"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -248,21 +247,29 @@ func (c *Connector) findAndConnectBackend(ctx context.Context, frontendConn net.
 
 	// PROXY protocol implementation
 	if c.sendProxyProto {
-		remoteHostStr, _, _ := net.SplitHostPort(backendHostPort)
-		sourceAddrStr, sourcePortStr, _ := net.SplitHostPort(clientAddr.String())
-		sourcePort, _ := strconv.Atoi(sourcePortStr)
+
+		// Determine transport protocol for the PROXY header by "analyzing" the frontend connection's address
+		transportProtocol := proxyproto.TCPv4
+		ourHostIpPart, _, err := net.SplitHostPort(frontendConn.LocalAddr().String())
+		if err != nil {
+			logrus.
+				WithError(err).
+				WithField("localAddr", frontendConn.LocalAddr()).
+				Error("Failed to extract host part of our address")
+			_ = backendConn.Close()
+			return
+		}
+		ourFrontendIp := net.ParseIP(ourHostIpPart)
+		if ourFrontendIp.To4() == nil {
+			transportProtocol = proxyproto.TCPv6
+		}
 
 		header := &proxyproto.Header{
 			Version:           2,
 			Command:           proxyproto.PROXY,
-			TransportProtocol: proxyproto.TCPv4,
-			SourceAddr: &net.TCPAddr{
-				IP:   net.ParseIP(sourceAddrStr),
-				Port: sourcePort,
-			},
-			DestinationAddr: &net.TCPAddr{
-				IP: net.ParseIP(remoteHostStr),
-			},
+			TransportProtocol: transportProtocol,
+			SourceAddr:        clientAddr,
+			DestinationAddr:   frontendConn.LocalAddr(), // our end of the client's connection
 		}
 
 		_, err = header.WriteTo(backendConn)