Repository: pires/go-proxyproto
Pull request #68 (2021-02-19 23:33:40)

Implement io.ReaderFrom/WriterTo for Conn
This change increase performance when proxying wrapped connections using io.Copy.
Since go 1.11 copying between tcp connections uses the splice system call on linux yielding considerable performance improvments.
See: https://golang.org/doc/go1.11#net

Signed-off-by: Fabian Ruff <fabian.ruff@sap.com>

coveralls: 
[![Coverage Status](https://coveralls.io/builds/37289817/badge)](https://coveralls.io/builds/37289817)

Coverage decreased (-1.4%) to 92.794% when pulling **d38c8e36a06a5fcea0e594fcf9b26df8174cf18d on databus23:readerfrom-writerto** into **fff0abf66729bfbb17e1e0f9131c43644d576225 on pires:main**.


Patch:
diff --git a/protocol.go b/protocol.go
index 878ca1c..d044335 100644
--- a/protocol.go
+++ b/protocol.go
@@ -2,6 +2,7 @@ package proxyproto
 
 import (
 	"bufio"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -237,3 +238,20 @@ func (p *Conn) readHeader() error {
 
 	return err
 }
+
+// ReadFrom implements the io.ReaderFrom ReadFrom method
+func (p *Conn) ReadFrom(r io.Reader) (int64, error) {
+	if rf, ok := p.conn.(io.ReaderFrom); ok {
+		return rf.ReadFrom(r)
+	}
+	return io.Copy(p.conn, r)
+}
+
+// WriteTo implements io.WriterTo
+func (p *Conn) WriteTo(w io.Writer) (int64, error) {
+	p.once.Do(func() { p.readErr = p.readHeader() })
+	if p.readErr != nil {
+		return 0, p.readErr
+	}
+	return p.bufReader.WriteTo(w)
+}


Test patch:
diff --git a/protocol_test.go b/protocol_test.go
index f5c08be..e4c9d3d 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -9,6 +9,8 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
 	"testing"
 )
@@ -741,6 +743,186 @@ func Test_MisconfiguredTLSServerRespondsWithUnderlyingError(t *testing.T) {
 	}
 }
 
+type testConn struct {
+	readFromCalledWith io.Reader
+	reads              int
+	net.Conn           // nil; crash on any unexpected use
+}
+
+func (c *testConn) ReadFrom(r io.Reader) (int64, error) {
+	c.readFromCalledWith = r
+	b, err := ioutil.ReadAll(r)
+	return int64(len(b)), err
+}
+func (c *testConn) Write(p []byte) (int, error) {
+	return len(p), nil
+}
+func (c *testConn) Read(p []byte) (int, error) {
+	if c.reads == 0 {
+		return 0, io.EOF
+	}
+	c.reads--
+	return 1, nil
+}
+
+func TestCopyToWrappedConnection(t *testing.T) {
+	innerConn := &testConn{}
+	wrappedConn := NewConn(innerConn)
+	dummySrc := &testConn{reads: 1}
+
+	io.Copy(wrappedConn, dummySrc)
+	if innerConn.readFromCalledWith != dummySrc {
+		t.Error("Expected io.Copy to delegate to ReadFrom function of inner destination connection")
+	}
+}
+
+func TestCopyFromWrappedConnection(t *testing.T) {
+	wrappedConn := NewConn(&testConn{reads: 1})
+	dummyDst := &testConn{}
+
+	io.Copy(dummyDst, wrappedConn)
+	if dummyDst.readFromCalledWith != wrappedConn.conn {
+		t.Errorf("Expected io.Copy to pass inner source connection to ReadFrom method of destination")
+	}
+}
+
+func TestCopyFromWrappedConnectionToWrappedConnection(t *testing.T) {
+	innerConn1 := &testConn{reads: 1}
+	wrappedConn1 := NewConn(innerConn1)
+	innerConn2 := &testConn{}
+	wrappedConn2 := NewConn(innerConn2)
+
+	io.Copy(wrappedConn1, wrappedConn2)
+	if innerConn1.readFromCalledWith != innerConn2 {
+		t.Errorf("Expected io.Copy to pass inner source connection to ReadFrom of inner destination connection")
+	}
+}
+
+func benchmarkTCPProxy(size int, b *testing.B) {
+	//create and start the echo backend
+	backend, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatalf("err: %v", err)
+	}
+	defer backend.Close()
+	go func() {
+		for {
+			conn, err := backend.Accept()
+			if err != nil {
+				break
+			}
+			_, err = io.Copy(conn, conn)
+			conn.Close()
+			if err != nil {
+				b.Fatalf("Failed to read entire payload: %v", err)
+			}
+		}
+	}()
+
+	//start the proxyprotocol enabled tcp proxy
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatalf("err: %v", err)
+	}
+	defer l.Close()
+	pl := &Listener{Listener: l}
+	go func() {
+		for {
+			conn, err := pl.Accept()
+			if err != nil {
+				break
+			}
+			bConn, err := net.Dial("tcp", backend.Addr().String())
+			if err != nil {
+				b.Fatalf("failed to dial backend: %v", err)
+			}
+			go func() {
+				_, err = io.Copy(bConn, conn)
+				if err != nil {
+					b.Fatalf("Failed to proxy incoming data to backend: %v", err)
+				}
+				bConn.(*net.TCPConn).CloseWrite()
+			}()
+			_, err = io.Copy(conn, bConn)
+			if err != nil {
+				b.Fatalf("Failed to proxy data from backend: %v", err)
+			}
+			conn.Close()
+			bConn.Close()
+		}
+	}()
+
+	data := make([]byte, size)
+
+	header := &Header{
+		Version:           2,
+		Command:           PROXY,
+		TransportProtocol: TCPv4,
+		SourceAddr: &net.TCPAddr{
+			IP:   net.ParseIP("10.1.1.1"),
+			Port: 1000,
+		},
+		DestinationAddr: &net.TCPAddr{
+			IP:   net.ParseIP("20.2.2.2"),
+			Port: 2000,
+		},
+	}
+
+	//now for the actual benchmark
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			b.Fatalf("err: %v", err)
+		}
+		// Write out the header!
+		header.WriteTo(conn)
+		//send data
+		go func() {
+			_, err = conn.Write(data)
+			if err != nil {
+				b.Fatalf("Failed to write data: %v", err)
+			}
+			conn.(*net.TCPConn).CloseWrite()
+
+		}()
+		//receive data
+		n, err := io.Copy(ioutil.Discard, conn)
+		if n != int64(len(data)) {
+			b.Fatalf("Expected to receive %d bytes, got %d", len(data), n)
+		}
+		if err != nil {
+			b.Fatalf("Failed to read data: %v", err)
+		}
+		conn.Close()
+	}
+}
+
+func BenchmarkTCPProxy16KB(b *testing.B) {
+	benchmarkTCPProxy(16*1024, b)
+}
+func BenchmarkTCPProxy32KB(b *testing.B) {
+	benchmarkTCPProxy(32*1024, b)
+}
+func BenchmarkTCPProxy64KB(b *testing.B) {
+	benchmarkTCPProxy(64*1024, b)
+}
+func BenchmarkTCPProxy128KB(b *testing.B) {
+	benchmarkTCPProxy(128*1024, b)
+}
+func BenchmarkTCPProxy256KB(b *testing.B) {
+	benchmarkTCPProxy(256*1024, b)
+}
+func BenchmarkTCPProxy512KB(b *testing.B) {
+	benchmarkTCPProxy(512*1024, b)
+}
+func BenchmarkTCPProxy1024KB(b *testing.B) {
+	benchmarkTCPProxy(1024*1024, b)
+}
+func BenchmarkTCPProxy2048KB(b *testing.B) {
+	benchmarkTCPProxy(2048*1024, b)
+}
+
 // copied from src/net/http/internal/testcert.go
 
 // Copyright 2015 The Go Authors. All rights reserved.