From e5010163d10456739159ad738dc5b20305a6e31c Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Sat, 13 Apr 2019 15:50:24 +0800
Subject: [PATCH] add proxy proto v2 support

---
 middleware.go     |  8 ++++----
 proxy_protocol.go | 50 +++++++++++++++++++++++++++++++++++++++++++++++
 server.go         | 26 +++++++++++++++++-------
 3 files changed, 73 insertions(+), 11 deletions(-)
 create mode 100644 proxy_protocol.go

diff --git a/middleware.go b/middleware.go
index ee8177b..189886d 100644
--- a/middleware.go
+++ b/middleware.go
@@ -43,14 +43,14 @@ func LogHandler(handler http.Handler) http.Handler {
 			if err := recover(); err != nil {
 				log.Error(err)
 				w.WriteHeader(http.StatusInternalServerError)
-				log.Infof("\"%s %s %s\" - %d %d \"%s\"",
-					r.Method, r.RequestURI, r.Proto, 500, 0, r.UserAgent())
+				log.Infof("\"%s - %s %s %s\" - %d %d \"%s\"",
+					r.RemoteAddr, r.Method, r.RequestURI, r.Proto, 500, 0, r.UserAgent())
 			}
 		}()
 
 		lh := &logHandler{w: w}
 		handler.ServeHTTP(lh, r)
-		log.Infof("\"%s %s %s\" - %d %d \"%s\"",
-			r.Method, r.RequestURI, r.Proto, lh.Status(), lh.size, r.UserAgent())
+		log.Infof("\"%s - %s %s %s\" - %d %d \"%s\"",
+			r.RemoteAddr, r.Method, r.RequestURI, r.Proto, lh.Status(), lh.size, r.UserAgent())
 	})
 }
diff --git a/proxy_protocol.go b/proxy_protocol.go
new file mode 100644
index 0000000..6536fed
--- /dev/null
+++ b/proxy_protocol.go
@@ -0,0 +1,50 @@
+package main
+
+import (
+	"bufio"
+	"net"
+
+	proxyproto "github.com/pires/go-proxyproto"
+)
+
+type protoListener struct {
+	net.Listener
+}
+
+type protoConn struct {
+	net.Conn
+	headerDone bool
+	r          *bufio.Reader
+	proxy      *proxyproto.Header
+}
+
+func (l *protoListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &protoConn{Conn: c}, err
+}
+
+func (c *protoConn) Read(buf []byte) (int, error) {
+	var err error
+	if !c.headerDone {
+		c.r = bufio.NewReader(c.Conn)
+		c.proxy, err = proxyproto.Read(c.r)
+		if err != nil && err != proxyproto.ErrNoProxyProtocol {
+			return 0, err
+		}
+		c.headerDone = true
+		return c.r.Read(buf)
+	}
+	return c.r.Read(buf)
+}
+
+func (c *protoConn) RemoteAddr() net.Addr {
+	if c.proxy == nil {
+		return c.Conn.RemoteAddr()
+	}
+	return &net.TCPAddr{
+		IP:   c.proxy.SourceAddress,
+		Port: int(c.proxy.SourcePort)}
+}
diff --git a/server.go b/server.go
index 5029ad3..c3a76c5 100644
--- a/server.go
+++ b/server.go
@@ -80,18 +80,21 @@ func (srv *server) serveTLS() {
 	if err != nil {
 		log.Fatalln("load certificate failed", err)
 	}
-	l, err := tls.Listen("tcp", srv.addr.Host,
-		&tls.Config{
-			Certificates: []tls.Certificate{cert},
-			//NextProtos:   []string{"h2"},
-		})
+
+	l, err := net.Listen("tcp", srv.addr.Host)
 	if err != nil {
 		log.Fatalln("listen tls", err)
 	}
 	defer l.Close()
+
 	log.Debugf("listen tls://%s", l.Addr().String())
+	tl := tls.NewListener(&protoListener{l}, &tls.Config{
+		Certificates: []tls.Certificate{cert},
+		//NextProtos:   []string{"h2"},
+	})
+
 	for {
-		conn, err := l.Accept()
+		conn, err := tl.Accept()
 		if err != nil {
 			log.Debugln("tls accept", err)
 			break
@@ -102,8 +105,17 @@ func (srv *server) serveTLS() {
 
 func (srv *server) serveHTTPS() {
 	log.Debugf("listen https://%s", srv.addr.Host)
-	err := http.ListenAndServeTLS(srv.addr.Host, srv.cert, srv.key, LogHandler(srv))
+
+	l, err := net.Listen("tcp", srv.addr.Host)
 	if err != nil {
+		log.Fatalln("listen https", err)
+	}
+	defer l.Close()
+
+	httpsrv := &http.Server{
+		Handler: LogHandler(srv),
+	}
+	if err := httpsrv.ServeTLS(&protoListener{l}, srv.cert, srv.key); err != nil {
 		log.Fatal(err)
 	}
 }