From 8bbe643f4dd5dce396691355bba0c97d61087660 Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Sat, 13 Apr 2019 17:16:22 +0800
Subject: [PATCH] seperate protolistener

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

diff --git a/go.mod b/go.mod
index f037b9e..dc58746 100644
--- a/go.mod
+++ b/go.mod
@@ -4,10 +4,10 @@ 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/kr/pretty v0.1.0 // indirect
 	github.com/miekg/dns v1.1.8
-	github.com/pires/go-proxyproto v0.0.0-20190111085350-4d51b51e3bfc
 	golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3
 	golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6 // indirect
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
diff --git a/go.sum b/go.sum
index 13b15c7..5bcc31f 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/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
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 c3a76c5..ad976e4 100644
--- a/server.go
+++ b/server.go
@@ -8,6 +8,7 @@ import (
 	"strconv"
 
 	"github.com/fangdingjun/go-log"
+	"github.com/fangdingjun/protolistener"
 )
 
 type server struct {
@@ -88,7 +89,7 @@ func (srv *server) serveTLS() {
 	defer l.Close()
 
 	log.Debugf("listen tls://%s", l.Addr().String())
-	tl := tls.NewListener(&protoListener{l}, &tls.Config{
+	tl := tls.NewListener(protolistener.New(l), &tls.Config{
 		Certificates: []tls.Certificate{cert},
 		//NextProtos:   []string{"h2"},
 	})
@@ -115,7 +116,7 @@ func (srv *server) serveHTTPS() {
 	httpsrv := &http.Server{
 		Handler: LogHandler(srv),
 	}
-	if err := httpsrv.ServeTLS(&protoListener{l}, srv.cert, srv.key); err != nil {
+	if err := httpsrv.ServeTLS(protolistener.New(l), srv.cert, srv.key); err != nil {
 		log.Fatal(err)
 	}
 }