feat: implement PROXY protocol to propagate source IPs (#22)

Fixes #13 

Co-authored-by: Geoff Bourne <itzgeoff@gmail.com>

server/connector.go

@@ -6,9 +6,11 @@ import (
 	"github.com/go-kit/kit/metrics"
 	"github.com/itzg/mc-router/mcproto"
 	"github.com/juju/ratelimit"
+	"github.com/pires/go-proxyproto"
 	"github.com/sirupsen/logrus"
 	"io"
 	"net"
+	"strconv"
 	"time"
 )
 
@@ -29,16 +31,18 @@ type ConnectorMetrics struct {
 	ActiveConnections metrics.Gauge
 }
 
-func NewConnector(metrics *ConnectorMetrics) Connector {
+func NewConnector(metrics *ConnectorMetrics, sendProxyProto bool) Connector {
 
 	return &connectorImpl{
-		metrics: metrics,
+		metrics:        metrics,
+		sendProxyProto: sendProxyProto,
 	}
 }
 
 type connectorImpl struct {
-	state   mcproto.State
-	metrics *ConnectorMetrics
+	state          mcproto.State
+	metrics        *ConnectorMetrics
+	sendProxyProto bool
 }
 
 func (c *connectorImpl) StartAcceptingConnections(ctx context.Context, listenAddress string, connRateLimit int) error {
@@ -189,12 +193,45 @@ func (c *connectorImpl) findAndConnectBackend(ctx context.Context, frontendConn
 	c.metrics.ActiveConnections.Add(1)
 	defer c.metrics.ActiveConnections.Add(-1)
 
+	// PROXY protocol implementation
+	if c.sendProxyProto {
+		remoteHostStr, _, _ := net.SplitHostPort(backendHostPort)
+		sourceAddrStr, sourcePortStr, _ := net.SplitHostPort(clientAddr.String())
+		sourcePort, _ := strconv.Atoi(sourcePortStr)
+
+		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),
+			},
+		}
+
+		_, err = header.WriteTo(backendConn)
+		if err != nil {
+			logrus.
+				WithError(err).
+				WithField("clientAddr", header.SourceAddr).
+				WithField("destAddr", header.DestinationAddr).
+				Error("Failed to write PROXY header")
+			c.metrics.Errors.With("type", "proxy_write").Add(1)
+			_ = backendConn.Close()
+			return
+		}
+	}
+
 	amount, err := io.Copy(backendConn, preReadContent)
 	if err != nil {
 		logrus.WithError(err).Error("Failed to write handshake to backend connection")
 		c.metrics.Errors.With("type", "backend_failed").Add(1)
 		return
 	}
+
 	logrus.WithField("amount", amount).Debug("Relayed handshake to backend")
 	if err = frontendConn.SetReadDeadline(noDeadline); err != nil {
 		logrus.
@@ -204,6 +241,7 @@ func (c *connectorImpl) findAndConnectBackend(ctx context.Context, frontendConn
 		c.metrics.Errors.With("type", "read_deadline").Add(1)
 		return
 	}
+
 	c.pumpConnections(ctx, frontendConn, backendConn)
 	return
 }