From d1788369dd2af8176f5a72a9682cef904b255d09 Mon Sep 17 00:00:00 2001
From: Cameron Moore <moorereason@gmail.com>
Date: Fri, 10 Nov 2017 22:09:05 -0600
Subject: [PATCH] Add PROXY protocol support

This commit adds a custom net.Listener to the webhook HTTP server to
enable PROXY protocol support.  I've copied in the keep-alive listener
from the Go net/http package, so the non-PROXY server should behave just
like the stdlib.
---
 http.go    | 30 ++++++++++++++++++++++++++++++
 webhook.go | 29 +++++++++++++++++++++++++----
 2 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 http.go

diff --git a/http.go b/http.go
new file mode 100644
index 0000000..8eb0831
--- /dev/null
+++ b/http.go
@@ -0,0 +1,30 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// tcpKeepAliveListener copied from Go net/http package.
+
+package main
+
+import (
+	"net"
+	"time"
+)
+
+// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
+// connections. It's used by ListenAndServe and ListenAndServeTLS so
+// dead TCP connections (e.g. closing laptop mid-download) eventually
+// go away.
+type tcpKeepAliveListener struct {
+	*net.TCPListener
+}
+
+func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
+	tc, err := ln.AcceptTCP()
+	if err != nil {
+		return
+	}
+	tc.SetKeepAlive(true)
+	tc.SetKeepAlivePeriod(3 * time.Minute)
+	return tc, nil
+}
diff --git a/webhook.go b/webhook.go
index 4df45f8..da12343 100644
--- a/webhook.go
+++ b/webhook.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -15,6 +16,7 @@ import (
 
 	"github.com/adnanh/webhook/hook"
 
+	"github.com/armon/go-proxyproto"
 	"github.com/codegangsta/negroni"
 	"github.com/gorilla/mux"
 	"github.com/satori/go.uuid"
@@ -38,6 +40,7 @@ var (
 	cert               = flag.String("cert", "cert.pem", "path to the HTTPS certificate pem file")
 	key                = flag.String("key", "key.pem", "path to the HTTPS certificate private key pem file")
 	justDisplayVersion = flag.Bool("version", false, "display webhook version and quit")
+	proxyProtocol      = flag.Bool("proxyProtocol", false, "enable PROXY protocol support")
 
 	responseHeaders hook.ResponseHeaders
 	hooksFiles      hook.HooksFiles
@@ -186,17 +189,35 @@ func main() {
 	}
 
 	router.HandleFunc(hooksURL, hookHandler)
-
 	n.UseHandler(router)
 
+	addr := fmt.Sprintf("%s:%d", *ip, *port)
+
+	svr := &http.Server{
+		Addr:    addr,
+		Handler: n,
+	}
+
+	// In order to support the PROXY protocol, we need a custom listener.
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Add keep-alive settings
+	listener = tcpKeepAliveListener{listener.(*net.TCPListener)}
+
+	if *proxyProtocol {
+		listener = &proxyproto.Listener{Listener: listener}
+	}
+
 	if *secure {
 		log.Printf("serving hooks on https://%s:%d%s", *ip, *port, hooksURL)
-		log.Fatal(http.ListenAndServeTLS(fmt.Sprintf("%s:%d", *ip, *port), *cert, *key, n))
+		log.Fatal(svr.ServeTLS(listener, *cert, *key))
 	} else {
 		log.Printf("serving hooks on http://%s:%d%s", *ip, *port, hooksURL)
-		log.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", *ip, *port), n))
+		log.Fatal(svr.Serve(listener))
 	}
-
 }
 
 func hookHandler(w http.ResponseWriter, r *http.Request) {