protocol: avoid double buffering

protocol.go

@@ -43,10 +43,11 @@ type Conn struct {
 	once              sync.Once
 	readErr           error
 	conn              net.Conn
-	Validate          Validator
 	bufReader         *bufio.Reader
+	reader            io.Reader
 	header            *Header
 	ProxyHeaderPolicy Policy
+	Validate          Validator
 	readHeaderTimeout time.Duration
 }
 
@@ -134,8 +135,15 @@ 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
+	br := bufio.NewReaderSize(conn, bufSize)
+
 	pConn := &Conn{
-		bufReader: bufio.NewReader(conn),
+		bufReader: br,
+		reader:    io.MultiReader(br, conn),
 		conn:      conn,
 	}
 
@@ -157,7 +165,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
@@ -339,5 +347,27 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
 	if p.readErr != nil {
 		return 0, p.readErr
 	}
-	return p.bufReader.WriteTo(w)
+
+	b := make([]byte, p.bufReader.Buffered())
+	if _, err := p.bufReader.Read(b); err != nil {
+		return 0, err // this should never as we read buffered data
+	}
+
+	var n int64
+	{
+		nn, err := w.Write(b)
+		n += int64(nn)
+		if err != nil {
+			return n, err
+		}
+	}
+	{
+		nn, err := io.Copy(w, p.conn)
+		n += nn
+		if err != nil {
+			return n, err
+		}
+	}
+
+	return n, nil
 }

protocol_test.go

@@ -11,7 +11,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net"
 	"testing"
 	"time"
@@ -1480,7 +1479,7 @@ type testConn struct {
 
 func (c *testConn) ReadFrom(r io.Reader) (int64, error) {
 	c.readFromCalledWith = r
-	b, err := ioutil.ReadAll(r)
+	b, err := io.ReadAll(r)
 	return int64(len(b)), err
 }
 func (c *testConn) Write(p []byte) (int, error) {
@@ -1625,7 +1624,7 @@ func benchmarkTCPProxy(size int, b *testing.B) {
 
 		}()
 		//receive data
-		n, err := io.Copy(ioutil.Discard, conn)
+		n, err := io.Copy(io.Discard, conn)
 		if n != int64(len(data)) {
 			b.Fatalf("Expected to receive %d bytes, got %d", len(data), n)
 		}