added the proxy protocol to the server

1 files changed, 14 insertions, 1 deletions

diff --git a/icanhazallips.go b/icanhazallips.go
index 1564ef4..48222bf 100644
--- a/icanhazallips.go
+++ b/icanhazallips.go
@@ -10,6 +10,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/pires/go-proxyproto"
 )
 
 const (
@@ -166,5 +168,16 @@ func main() {
 		Handler:           nil,
 	}
 
-	log.Fatal(server.ListenAndServe())
+	ln, err := net.Listen("tcp", config.Addr)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	proxyListener := &proxyproto.Listener{
+		Listener:          ln,
+		ReadHeaderTimeout: time.Duration(config.ReadHeaderTimeout) * time.Second,
+	}
+	defer proxyListener.Close()
+
+	log.Fatal(server.Serve(proxyListener))
 }