diff --git a/protocol_test.go b/protocol_test.go index c4f5932..e4c9d3d 100644 --- a/protocol_test.go +++ b/protocol_test.go @@ -798,6 +798,131 @@ func TestCopyFromWrappedConnectionToWrappedConnection(t *testing.T) { } } +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.