seperate protolistener

go.mod

@@ -4,6 +4,7 @@ go 1.13
 
 require (
 	github.com/fangdingjun/go-log v4.0.0+incompatible
+	github.com/fangdingjun/protolistener v0.0.0-20190413090244-355464816fc0
 	github.com/go-yaml/yaml v2.1.0+incompatible
 	github.com/gorilla/handlers v1.4.0
 	github.com/gorilla/mux v1.7.1

go.sum

@@ -1,5 +1,7 @@
 github.com/fangdingjun/go-log v4.0.0+incompatible h1:h/3S0FRNiFR1A5xkK/ECHtmc2O/DVESKCzjhpgzszfA=
 github.com/fangdingjun/go-log v4.0.0+incompatible/go.mod h1:Y6ko/KySCpIcAX3mrfiUJtuPY2saPwofv9WYcUatPJY=
+github.com/fangdingjun/protolistener v0.0.0-20190413090244-355464816fc0 h1:dafx8KP2yvsX8oAFwJ6Bx54h3YE2P7UXLB453z7LbPo=
+github.com/fangdingjun/protolistener v0.0.0-20190413090244-355464816fc0/go.mod h1:bp6oYi3nsUwqGyzikoHJ672i9vX39Qp1h1C/r/vkIIg=
 github.com/go-yaml/yaml v2.1.0+incompatible h1:RYi2hDdss1u4YE7GwixGzWwVo47T8UQwnTLB6vQiq+o=
 github.com/go-yaml/yaml v2.1.0+incompatible/go.mod h1:w2MrLa16VYP0jy6N7M5kHaCkaLENm+P+Tv+MfurjSw0=
 github.com/gorilla/handlers v1.4.0 h1:XulKRWSQK5uChr4pEgSE4Tc/OcmnU9GJuSwdog/tZsA=

proxy_protocol.go

@@ -1,50 +0,0 @@
-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)}
-}

server.go

@@ -11,6 +11,7 @@ import (
 	"syscall"
 
 	"github.com/fangdingjun/go-log"
+	"github.com/fangdingjun/protolistener"
 	"github.com/gorilla/handlers"
 	"github.com/gorilla/mux"
 	"golang.org/x/net/http2"
@@ -69,7 +70,7 @@ func initServer(c *conf) error {
 			return err
 		}
 
-		l = &protoListener{Listener: l}
+		l = protolistener.New(l)
 
 		if len(certs) > 0 {
 			tlsconfig.Certificates = certs