From ed44759412da102f90262a51b8d9fdf9060931b8 Mon Sep 17 00:00:00 2001
From: Andrei Belov <defanator@gmail.com>
Date: Sun, 16 Feb 2025 21:56:54 +0400
Subject: [PATCH] Add proxy protocol support

---
 README.md   |  2 ++
 exporter.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 go.mod      |  1 +
 go.sum      |  2 ++
 4 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 0e99782e..11efc9dc 100644
--- a/README.md
+++ b/README.md
@@ -181,6 +181,8 @@ Flags:
       --nginx.ssl-client-cert=""
                                  Path to the PEM encoded client certificate file to use when connecting to the server. ($SSL_CLIENT_CERT)
       --nginx.ssl-client-key=""  Path to the PEM encoded client certificate key file to use when connecting to the server. ($SSL_CLIENT_KEY)
+      --[no-]nginx.proxy-protocol
+                                 Pass proxy protocol payload to nginx listeners. ($PROXY_PROTOCOL)
       --nginx.timeout=5s         A timeout for scraping metrics from NGINX or NGINX Plus. ($TIMEOUT)
       --prometheus.const-label=PROMETHEUS.CONST-LABEL ...
                                  Label that will be used in every metric. Format is label=value. It can be repeated multiple times. ($CONST_LABELS)
diff --git a/exporter.go b/exporter.go
index b70e39b6..9beddfa4 100644
--- a/exporter.go
+++ b/exporter.go
@@ -30,6 +30,8 @@ import (
 
 	"github.com/prometheus/exporter-toolkit/web"
 	"github.com/prometheus/exporter-toolkit/web/kingpinflag"
+
+	proxyproto "github.com/pires/go-proxyproto"
 )
 
 // positiveDuration is a wrapper of time.Duration to ensure only positive values are accepted.
@@ -90,6 +92,7 @@ var (
 	sslCaCert     = kingpin.Flag("nginx.ssl-ca-cert", "Path to the PEM encoded CA certificate file used to validate the servers SSL certificate.").Default("").Envar("SSL_CA_CERT").String()
 	sslClientCert = kingpin.Flag("nginx.ssl-client-cert", "Path to the PEM encoded client certificate file to use when connecting to the server.").Default("").Envar("SSL_CLIENT_CERT").String()
 	sslClientKey  = kingpin.Flag("nginx.ssl-client-key", "Path to the PEM encoded client certificate key file to use when connecting to the server.").Default("").Envar("SSL_CLIENT_KEY").String()
+	useProxyProto = kingpin.Flag("nginx.proxy-protocol", "Pass proxy protocol payload to nginx listeners.").Default("false").Envar("PROXY_PROTOCOL").Bool()
 
 	// Custom command-line flags.
 	timeout = createPositiveDurationFlag(kingpin.Flag("nginx.timeout", "A timeout for scraping metrics from NGINX or NGINX Plus.").Default("5s").Envar("TIMEOUT").HintOptions("5s", "10s", "30s", "1m", "5m"))
@@ -223,17 +226,65 @@ func main() {
 func registerCollector(logger *slog.Logger, transport *http.Transport,
 	addr string, labels map[string]string,
 ) {
+	var socketPath string
+
 	if strings.HasPrefix(addr, "unix:") {
-		socketPath, requestPath, err := parseUnixSocketAddress(addr)
+		var err error
+		var requestPath string
+		socketPath, requestPath, err = parseUnixSocketAddress(addr)
 		if err != nil {
 			logger.Error("parsing unix domain socket scrape address failed", "uri", addr, "error", err.Error())
 			os.Exit(1)
 		}
+		addr = "http://unix" + requestPath
+	}
 
+	if !*useProxyProto && socketPath != "" {
 		transport.DialContext = func(_ context.Context, _, _ string) (net.Conn, error) {
 			return net.Dial("unix", socketPath)
 		}
-		addr = "http://unix" + requestPath
+	}
+
+	if *useProxyProto {
+		transport.DialContext = func(_ context.Context, network, addr string) (net.Conn, error) {
+			if socketPath != "" {
+				network = "unix"
+				addr = socketPath
+			}
+
+			conn, err := (&net.Dialer{}).Dial(network, addr)
+			if err != nil {
+				return nil, fmt.Errorf("dialing %s %s: %w", network, addr, err)
+			}
+
+			localAddr := conn.LocalAddr()
+			remoteAddr := conn.RemoteAddr()
+			transportProtocol := proxyproto.TCPv4
+
+			switch addr := remoteAddr.(type) {
+			case *net.TCPAddr:
+				if addr.IP.To4() == nil {
+					transportProtocol = proxyproto.TCPv6
+				}
+			case *net.UnixAddr:
+				transportProtocol = proxyproto.UnixStream
+			}
+
+			header := &proxyproto.Header{
+				Version:           2,
+				Command:           proxyproto.PROXY,
+				TransportProtocol: transportProtocol,
+				SourceAddr:        localAddr,
+				DestinationAddr:   remoteAddr,
+			}
+
+			_, err = header.WriteTo(conn)
+			if err != nil {
+				return nil, fmt.Errorf("writing proxyproto header: %w", err)
+			}
+
+			return conn, nil
+		}
 	}
 
 	userAgent := fmt.Sprintf("NGINX-Prometheus-Exporter/v%v", common_version.Version)
diff --git a/go.mod b/go.mod
index ad77106a..34cb7238 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/prometheus/client_golang v1.20.5
 	github.com/prometheus/common v0.62.0
 	github.com/prometheus/exporter-toolkit v0.14.0
+	github.com/pires/go-proxyproto v0.8.0
 )
 
 require (
diff --git a/go.sum b/go.sum
index d0489b3f..b9421a29 100644
--- a/go.sum
+++ b/go.sum
@@ -34,6 +34,8 @@ github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/nginx/nginx-plus-go-client/v2 v2.3.0 h1:ciKh1lwadNzUaOGjLcKWu/BGigASxU6p7v/6US71fhA=
 github.com/nginx/nginx-plus-go-client/v2 v2.3.0/go.mod h1:U7G5pqucUS1V4Uecs1xCsJ9knSsfwqhwu8ZEjoCYnmk=
+github.com/pires/go-proxyproto v0.8.0 h1:5unRmEAPbHXHuLjDg01CxJWf91cw3lKHc/0xzKpXEe0=
+github.com/pires/go-proxyproto v0.8.0/go.mod h1:iknsfgnH8EkjrMeMyvfKByp9TiBZCKZM0jx2xmKqnVY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y=