From 341802505f77903e7417e9e5503ba0665e55970e Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Sat, 13 Apr 2019 17:26:54 +0800
Subject: [PATCH] seperate protolistener

---
 go.mod            |  1 +
 go.sum            |  2 ++
 proxy_protocol.go | 50 -----------------------------------------------
 server.go         |  3 ++-
 4 files changed, 5 insertions(+), 51 deletions(-)
 delete mode 100644 proxy_protocol.go

diff --git a/go.mod b/go.mod
index 6f015c3..9000a63 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index d70325d..113ee9d 100644
--- a/go.sum
+++ b/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=
diff --git a/proxy_protocol.go b/proxy_protocol.go
deleted file mode 100644
index 6536fed..0000000
--- a/proxy_protocol.go
+++ /dev/null
@@ -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)}
-}
diff --git a/server.go b/server.go
index 816c452..97a874d 100644
--- a/server.go
+++ b/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