From 1f160533349e92f8a90962da03909b2457576b50 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Fri, 9 Oct 2020 12:21:19 +0200
Subject: [PATCH] Set PROXY protocol PP2_TYPE_AUTHORITY TLV

---
 server.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/server.go b/server.go
index cc86f85..0767845 100644
--- a/server.go
+++ b/server.go
@@ -145,7 +145,7 @@ func (ln *Listener) handle(conn net.Conn) error {
 		return fmt.Errorf("can't find frontend for server name %q", tlsState.ServerName)
 	}
 
-	return fe.handle(tlsConn)
+	return fe.handle(tlsConn, &tlsState)
 }
 
 type Frontend struct {
@@ -153,7 +153,7 @@ type Frontend struct {
 	Backend Backend
 }
 
-func (fe *Frontend) handle(downstream net.Conn) error {
+func (fe *Frontend) handle(downstream net.Conn, tlsState *tls.ConnectionState) error {
 	defer downstream.Close()
 
 	be := &fe.Backend
@@ -165,6 +165,15 @@ func (fe *Frontend) handle(downstream net.Conn) error {
 
 	if be.Proxy {
 		h := proxyproto.HeaderProxyFromAddrs(2, downstream.RemoteAddr(), downstream.LocalAddr())
+
+		var tlvs []proxyproto.TLV
+		if tlsState.ServerName != "" {
+			tlvs = append(tlvs, authorityTLV(tlsState.ServerName))
+		}
+		if err := h.SetTLVs(tlvs); err != nil {
+			return fmt.Errorf("failed to set PROXY protocol header TLVs: %v", err)
+		}
+
 		if _, err := h.WriteTo(upstream); err != nil {
 			return fmt.Errorf("failed to write PROXY protocol header: %v", err)
 		}
@@ -191,3 +200,11 @@ func duplexCopy(a, b io.ReadWriter) error {
 	}()
 	return <-done
 }
+
+func authorityTLV(name string) proxyproto.TLV {
+	return proxyproto.TLV{
+		Type:   proxyproto.PP2_TYPE_AUTHORITY,
+		Length: len(name),
+		Value:  []byte(name),
+	}
+}