Initial commit

README.md

@@ -0,0 +1,4 @@
+protolistener
+=====
+
+wrap github.com/pires/go-proxyproto to net.Listener

go.mod

@@ -0,0 +1,5 @@
+module github.com/fangdingjun/proto_listener
+
+go 1.13
+
+require github.com/pires/go-proxyproto v0.0.0-20190111085350-4d51b51e3bfc

go.sum

@@ -0,0 +1,2 @@
+github.com/pires/go-proxyproto v0.0.0-20190111085350-4d51b51e3bfc h1:lNOt1SMsgHXTdpuGw+RpnJtzUcCb/oRKZP65pBy9pr8=
+github.com/pires/go-proxyproto v0.0.0-20190111085350-4d51b51e3bfc/go.mod h1:6/gX3+E/IYGa0wMORlSMla999awQFdbaeQCHjSMKIzY=

proxy_protocol.go

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