Buffer only proxy header data

Reverts and re-implements incorrect #116
which passed all connection data through the buffered reader.

Author: AlexanderYastrebov

header.go

@@ -42,6 +42,7 @@ type Header struct {
 	SourceAddr        net.Addr
 	DestinationAddr   net.Addr
 	rawTLVs           []byte
+	length            int
 }
 
 // HeaderProxyFromAddrs creates a new PROXY header from a source and a

protocol.go

@@ -2,6 +2,7 @@ package proxyproto
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -51,7 +52,8 @@ type Conn struct {
 	once              sync.Once
 	readErr           error
 	conn              net.Conn
-	bufReader         *bufio.Reader
+	reader            io.Reader
+	headerBuf         *bytes.Buffer
 	header            *Header
 	ProxyHeaderPolicy Policy
 	Validate          Validator
@@ -150,14 +152,8 @@ func (p *Listener) Addr() net.Addr {
 // NewConn is used to wrap a net.Conn that may be speaking
 // the proxy protocol into a proxyproto.Conn
 func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
-	// For v1 the header length is at most 108 bytes.
-	// For v2 the header length is at most 52 bytes plus the length of the TLVs.
-	// We use 256 bytes to be safe.
-	const bufSize = 256
-
 	pConn := &Conn{
-		bufReader: bufio.NewReaderSize(conn, bufSize),
-		conn:      conn,
+		conn: conn,
 	}
 
 	for _, opt := range opts {
@@ -178,7 +174,7 @@ func (p *Conn) Read(b []byte) (int, error) {
 		return 0, p.readErr
 	}
 
-	return p.bufReader.Read(b)
+	return p.reader.Read(b)
 }
 
 // Write wraps original conn.Write
@@ -294,7 +290,22 @@ func (p *Conn) readHeader() error {
 		}
 	}
 
-	header, err := Read(p.bufReader)
+	// For v1 the header length is at most 108 bytes.
+	// For v2 the header length is at most 52 bytes plus the length of the TLVs.
+	// We use 256 bytes to be safe.
+	const bufSize = 256
+
+	bb := bytes.NewBuffer(make([]byte, 0, bufSize))
+	tr := io.TeeReader(p.conn, bb)
+	br := bufio.NewReaderSize(tr, bufSize)
+
+	header, err := Read(br)
+
+	if err == nil {
+		io.CopyN(io.Discard, bb, int64(header.length))
+	}
+	p.headerBuf = bb
+	p.reader = io.MultiReader(bb, p.conn)
 
 	// If the connection's readHeaderTimeout is more than 0, undo the change to the
 	// deadline that we made above. Because we retain the readDeadline as part of our
@@ -360,5 +371,9 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
 	if p.readErr != nil {
 		return 0, p.readErr
 	}
-	return p.bufReader.WriteTo(w)
+
+	if n, err := p.headerBuf.WriteTo(w); err != nil {
+		return n, err
+	}
+	return io.Copy(w, p.conn)
 }

v1.go

@@ -125,6 +125,7 @@ func parseVersion1(reader *bufio.Reader) (*Header, error) {
 	// Command doesn't exist in v1 but set it for other parts of this library
 	// to rely on it for determining connection details.
 	header := initVersion1()
+	header.length = len(buf)
 
 	// Transport protocol has been processed already.
 	header.TransportProtocol = transportProtocol

v2.go

@@ -100,6 +100,7 @@ func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
 	if !header.validateLength(length) {
 		return nil, ErrInvalidLength
 	}
+	header.length = 16 + int(length)
 
 	// Return early if the length is zero, which means that
 	// there's no address information and TLVs present for UNSPEC.