WIP: Add support to haproxy's proxy-protocol

main.go

@@ -126,6 +126,7 @@ func runServer(c *cli.Context) {
 		ReadHeaderTimeout: c.Duration("client-read-header-timeout"),
 		WriteTimeout:      c.Duration("client-write-timeout"),
 		IdleTimeout:       c.Duration("client-idle-timeout"),
+		ProxyProtocol:     c.Bool("proxy-protocol"),
 	})
 	if err != nil {
 		log.Fatal(err)
@@ -331,6 +332,10 @@ The value 'none' can be used to disable access logs.`,
 			Value: "native",
 			Usage: "Reverse proxy engine, options are 'native', 'sni' and 'fasthttp'. Using 'sni' and 'fasthttp' is highly experimental and not recommended for production environments.",
 		},
+		cli.BoolFlag{
+			Name:  "proxy-protocol",
+			Usage: "Enable parsing HAProxy's PROXY protocol v2 header in tcp connections.",
+		},
 		cli.BoolFlag{
 			Name:  "backend-cache",
 			Usage: "Enable caching backend results for 2 seconds. This may cause temporary inconsistencies.",

reverseproxy/native.go

@@ -72,9 +72,11 @@ func init() {
 type NativeReverseProxy struct {
 	http.Transport
 	ReverseProxyConfig
-	servers []*http.Server
-	rp      *httputil.ReverseProxy
-	dialer  *net.Dialer
+	servers   []*http.Server
+	rp        *httputil.ReverseProxy
+	dialer    *net.Dialer
+	mu        sync.RWMutex
+	remoteMap map[string]string
 }
 
 type fixedReadCloser struct {
@@ -108,6 +110,9 @@ func (p *bufferPool) Put(b []byte) {
 func (rp *NativeReverseProxy) Initialize(rpConfig ReverseProxyConfig) error {
 	rp.ReverseProxyConfig = rpConfig
 	rp.servers = make([]*http.Server, 0)
+	if rpConfig.ProxyProtocol {
+		rp.remoteMap = make(map[string]string)
+	}
 
 	rp.dialer = &net.Dialer{
 		Timeout:   rp.DialTimeout,
@@ -139,11 +144,24 @@ func (rp *NativeReverseProxy) Listen(listener net.Listener, tlsConfig *tls.Confi
 		ConnState: func(c net.Conn, s http.ConnState) {
 			switch s {
 			case http.StateNew:
+				if rp.remoteMap != nil {
+					remoteAddr, err := readProxyProtoV2Header(c)
+					if err == nil {
+						rp.mu.Lock()
+						rp.remoteMap[c.RemoteAddr().String()] = remoteAddr
+						rp.mu.Unlock()
+					}
+				}
 				openConnections.Inc()
 			case http.StateHijacked:
-				openConnections.Dec()
+				fallthrough
 			case http.StateClosed:
 				openConnections.Dec()
+				if rp.remoteMap != nil {
+					rp.mu.Lock()
+					delete(rp.remoteMap, c.RemoteAddr().String())
+					rp.mu.Unlock()
+				}
 			}
 		},
 	}
@@ -263,12 +281,20 @@ func (rp *NativeReverseProxy) RoundTrip(req *http.Request) (*http.Response, erro
 func (rp *NativeReverseProxy) doResponse(req *http.Request, reqData *RequestData, rsp *http.Response, isDebug bool, isDead bool, backendDuration time.Duration, originalForwardedFor string) *http.Response {
 	totalDuration := time.Since(reqData.StartTime)
 	logEntry := func() *log.LogEntry {
+		remoteAddr := req.RemoteAddr
+		if rp.remoteMap != nil {
+			rp.mu.RLock()
+			if remoteFromMap, ok := rp.remoteMap[req.RemoteAddr]; ok {
+				remoteAddr = remoteFromMap
+			}
+			rp.mu.RUnlock()
+		}
 		return &log.LogEntry{
 			Now:             time.Now(),
 			BackendDuration: backendDuration,
 			TotalDuration:   totalDuration,
 			BackendKey:      reqData.BackendKey,
-			RemoteAddr:      req.RemoteAddr,
+			RemoteAddr:      remoteAddr,
 			Method:          req.Method,
 			Path:            req.URL.Path,
 			Proto:           req.Proto,

reverseproxy/proxyprotocol.go

@@ -0,0 +1,79 @@
+// Copyright 2017 tsuru authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package reverseproxy
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+)
+
+type proxyHdrV2 struct {
+	Sig    [12]byte /* hex 0D 0A 0D 0A 00 0D 0A 51 55 49 54 0A */
+	VerCmd byte     /* protocol version and command */
+	Fam    byte     /* protocol family and address */
+	Len    uint16   /* number of following bytes part of the header */
+}
+
+type proxyHdrV2Ipv4 struct { /* for TCP/UDP over IPv4, len = 12 */
+	SrcAddr [4]byte
+	DstAddr [4]byte
+	SrcPort uint16
+	DstPort uint16
+}
+
+type proxyHdrV2Ipv6 struct { /* for TCP/UDP over IPv6, len = 36 */
+	SrcAddr [16]byte
+	DstAddr [16]byte
+	SrcPort uint16
+	DstPort uint16
+}
+
+type proxyHdrV2Sock struct { /* for AF_UNIX sockets, len = 216 */
+	SrcAddr [108]byte
+	DstAddr [108]byte
+}
+
+var proxyHdrSig = []byte("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A")
+
+func readProxyProtoV2Header(conn net.Conn) (string, error) {
+	var proxyHdr proxyHdrV2
+	err := binary.Read(conn, binary.BigEndian, &proxyHdr)
+	if err != nil {
+		return "", err
+	}
+	if !bytes.Equal(proxyHdr.Sig[:], proxyHdrSig) {
+		return "", errors.New("invalid proxy protocol header")
+	}
+	switch proxyHdr.Len {
+	case 12:
+		var data proxyHdrV2Ipv4
+		err := binary.Read(conn, binary.BigEndian, &data)
+		if err != nil {
+			return "", err
+		}
+		ip := net.IP(data.SrcAddr[:])
+		return fmt.Sprintf("%s:%d", ip.String(), data.SrcPort), nil
+	case 36:
+		var data proxyHdrV2Ipv6
+		err := binary.Read(conn, binary.BigEndian, &data)
+		if err != nil {
+			return "", err
+		}
+		ip := net.IP(data.SrcAddr[:])
+		return fmt.Sprintf("%s:%d", ip.String(), data.SrcPort), nil
+	case 216:
+		var data proxyHdrV2Sock
+		err := binary.Read(conn, binary.BigEndian, &data)
+		if err != nil {
+			return "", err
+		}
+		return string(data.SrcAddr[:]), nil
+	default:
+		return "", errors.New("invalid proxy protocol header")
+	}
+}

reverseproxy/reverseproxy.go

@@ -45,6 +45,7 @@ type ReverseProxyConfig struct {
 	WriteTimeout      time.Duration
 	IdleTimeout       time.Duration
 	RequestIDHeader   string
+	ProxyProtocol     bool
 }
 
 type RequestData struct {