From bbe67adf1ea929110e3909cc3702bdb5c08c0629 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 7 Sep 2020 21:28:24 +0200
Subject: [PATCH] Add support for the PROXY protocol

IPs whitelisted in accept-proxy-ip can now use the PROXY protocol to
indicate the original source/destination addresses.

Closes: https://todo.sr.ht/~emersion/soju/81
---
 cmd/soju/main.go | 21 +++++++++++++++++++++
 go.mod           |  1 +
 go.sum           |  2 ++
 3 files changed, 24 insertions(+)

diff --git a/cmd/soju/main.go b/cmd/soju/main.go
index 86bfd7d..9d74e0f 100644
--- a/cmd/soju/main.go
+++ b/cmd/soju/main.go
@@ -9,6 +9,8 @@ import (
 	"net/url"
 	"strings"
 
+	"github.com/pires/go-proxyproto"
+
 	"git.sr.ht/~emersion/soju"
 	"git.sr.ht/~emersion/soju/config"
 )
@@ -85,6 +87,7 @@ func main() {
 			if err != nil {
 				log.Fatalf("failed to start TLS listener on %q: %v", listen, err)
 			}
+			ln = proxyProtoListener(ln, srv)
 			go func() {
 				log.Fatal(srv.Serve(ln))
 			}()
@@ -97,6 +100,7 @@ func main() {
 			if err != nil {
 				log.Fatalf("failed to start listener on %q: %v", listen, err)
 			}
+			ln = proxyProtoListener(ln, srv)
 			go func() {
 				log.Fatal(srv.Serve(ln))
 			}()
@@ -138,6 +142,7 @@ func main() {
 			if err != nil {
 				log.Fatalf("failed to start listener on %q: %v", listen, err)
 			}
+			ln = proxyProtoListener(ln, srv)
 			go func() {
 				log.Fatal(srv.Identd.Serve(ln))
 			}()
@@ -149,3 +154,19 @@ func main() {
 	}
 	log.Fatal(srv.Run())
 }
+
+func proxyProtoListener(ln net.Listener, srv *soju.Server) net.Listener {
+	return &proxyproto.Listener{
+		Listener: ln,
+		Policy: func(upstream net.Addr) (proxyproto.Policy, error) {
+			tcpAddr, ok := upstream.(*net.TCPAddr)
+			if !ok {
+				return proxyproto.IGNORE, nil
+			}
+			if srv.AcceptProxyIPs.Contains(tcpAddr.IP) {
+				return proxyproto.USE, nil
+			}
+			return proxyproto.IGNORE, nil
+		},
+	}
+}
diff --git a/go.mod b/go.mod
index 478206e..6d25354 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/klauspost/compress v1.10.11 // indirect
 	github.com/mattn/go-sqlite3 v1.14.1
+	github.com/pires/go-proxyproto v0.1.3
 	golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a
 	golang.org/x/sys v0.0.0-20200824131525-c12d262b63d8 // indirect
 	gopkg.in/irc.v3 v3.1.3
diff --git a/go.sum b/go.sum
index cd46071..8e1787b 100644
--- a/go.sum
+++ b/go.sum
@@ -50,6 +50,8 @@ github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 h1:ZqeYNhU3OH
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742 h1:Esafd1046DLDQ0W1YjYsBW+p8U2u7vzgW2SQVmlNazg=
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+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/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=