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