From e3ac31414f145695bd37eb1c2d3e371a9eb682d9 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Wed, 9 Sep 2020 14:52:41 +0200
Subject: [PATCH] Add support for the PROXY protocol

---
 README.md |  3 +++
 go.mod    |  1 +
 go.sum    |  2 ++
 main.go   |  5 ++++-
 server.go | 37 +++++++++++++++++++++++++++++++++++++
 5 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index bcd0add..3c7ff30 100644
--- a/README.md
+++ b/README.md
@@ -4,6 +4,7 @@ A TLS reverse proxy.
 
 - Automatic TLS with Let's Encrypt
 - Route incoming connections to backends using Server Name Indication
+- Support for the [PROXY protocol]
 
 Example configuration:
 
@@ -14,3 +15,5 @@ Example configuration:
 ## License
 
 MIT
+
+[PROXY protocol]: https://www.haproxy.org/download/2.3/doc/proxy-protocol.txt
diff --git a/go.mod b/go.mod
index 6524e2f..8dfbcea 100644
--- a/go.mod
+++ b/go.mod
@@ -5,4 +5,5 @@ go 1.15
 require (
 	github.com/caddyserver/certmagic v0.11.2
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
+	github.com/pires/go-proxyproto v0.1.3
 )
diff --git a/go.sum b/go.sum
index 88500b6..28452c3 100644
--- a/go.sum
+++ b/go.sum
@@ -195,6 +195,8 @@ github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJ
 github.com/oracle/oci-go-sdk v7.0.0+incompatible/go.mod h1:VQb79nF8Z2cwLkLS35ukwStZIg5F66tcBccjip/j888=
 github.com/ovh/go-ovh v0.0.0-20181109152953-ba5adb4cf014/go.mod h1:joRatxRJaZBsY3JAOEMcoOp05CnZzsx4scTxi95DHyQ=
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
+github.com/pires/go-proxyproto v0.1.3 h1:2XEuhsQluSNA5QIQkiUv8PfgZ51sNYIQkq/yFquiSQM=
+github.com/pires/go-proxyproto v0.1.3/go.mod h1:Odh9VFOZJCf9G8cLW5o435Xf1J95Jw9Gw5rnCjcwzAY=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/main.go b/main.go
index f789063..9b5755d 100644
--- a/main.go
+++ b/main.go
@@ -98,7 +98,10 @@ func parseBackend(backend *Backend, d *Directive) error {
 		return fmt.Errorf("failed to parse backend URI %q: %v", backendURI, err)
 	}
 
-	// TODO: +proxy to use the PROXY protocol
+	if strings.HasSuffix(u.Scheme, "+proxy") {
+		u.Scheme = strings.TrimSuffix(u.Scheme, "+proxy")
+		backend.Proxy = true
+	}
 
 	switch u.Scheme {
 	case "", "tcp":
diff --git a/server.go b/server.go
index 91af70e..276ca88 100644
--- a/server.go
+++ b/server.go
@@ -9,6 +9,7 @@ import (
 	"net"
 
 	"github.com/caddyserver/certmagic"
+	"github.com/pires/go-proxyproto"
 )
 
 type Server struct {
@@ -152,12 +153,20 @@ func (fe *Frontend) handle(downstream net.Conn) error {
 	}
 	defer upstream.Close()
 
+	if be.Proxy {
+		h := proxyHeader(downstream.RemoteAddr(), downstream.LocalAddr())
+		if _, err := h.WriteTo(upstream); err != nil {
+			return fmt.Errorf("failed to write PROXY protocol header: %v", err)
+		}
+	}
+
 	return duplexCopy(upstream, downstream)
 }
 
 type Backend struct {
 	Network string
 	Address string
+	Proxy   bool
 }
 
 func duplexCopy(a, b io.ReadWriter) error {
@@ -172,3 +181,31 @@ func duplexCopy(a, b io.ReadWriter) error {
 	}()
 	return <-done
 }
+
+func proxyHeader(sourceAddr, destAddr net.Addr) *proxyproto.Header {
+	h := proxyproto.Header{
+		Version: 2,
+		Command: proxyproto.PROXY,
+	}
+
+	switch sourceAddr := sourceAddr.(type) {
+	case *net.TCPAddr:
+		destAddr, ok := destAddr.(*net.TCPAddr)
+		if !ok {
+			break
+		}
+		if localIP4 := sourceAddr.IP.To4(); len(localIP4) == net.IPv4len {
+			h.TransportProtocol = proxyproto.TCPv4
+		} else if len(sourceAddr.IP) == net.IPv6len {
+			h.TransportProtocol = proxyproto.TCPv6
+		} else {
+			break
+		}
+		h.SourceAddress = sourceAddr.IP
+		h.DestinationAddress = destAddr.IP
+		h.SourcePort = uint16(sourceAddr.Port)
+		h.DestinationPort = uint16(destAddr.Port)
+	}
+
+	return &h
+}