From ebc8d4df1ffdbc3b14480b3e7b734206ee148979 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Wed, 24 May 2023 11:08:13 -0700
Subject: [PATCH] tcp_benchmark: allow packet sniffing/logging

PiperOrigin-RevId: 534910831
---
 test/benchmarks/tcp/BUILD            | 1 +
 test/benchmarks/tcp/tcp_benchmark.sh | 4 ++++
 test/benchmarks/tcp/tcp_proxy.go     | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/test/benchmarks/tcp/BUILD b/test/benchmarks/tcp/BUILD
index 4442753a6..adb1efca0 100644
--- a/test/benchmarks/tcp/BUILD
+++ b/test/benchmarks/tcp/BUILD
@@ -14,6 +14,7 @@ go_binary(
         "//pkg/tcpip/adapters/gonet",
         "//pkg/tcpip/link/fdbased",
         "//pkg/tcpip/link/qdisc/fifo",
+        "//pkg/tcpip/link/sniffer",
         "//pkg/tcpip/network/arp",
         "//pkg/tcpip/network/ipv4",
         "//pkg/tcpip/network/ipv6",
diff --git a/test/benchmarks/tcp/tcp_benchmark.sh b/test/benchmarks/tcp/tcp_benchmark.sh
index 910154246..8394b688f 100755
--- a/test/benchmarks/tcp/tcp_benchmark.sh
+++ b/test/benchmarks/tcp/tcp_benchmark.sh
@@ -50,6 +50,7 @@ disable_linux_gso=
 disable_linux_gro=
 gro=0
 num_client_threads=1
+sniff=false
 
 # Check for netem support.
 lsmod_output=$(lsmod | grep sch_netem)
@@ -199,6 +200,9 @@ while [[ $# -gt 0 ]]; do
       [[ "$#" -le 0 ]] && echo "no iperf name provided" && exit 1
       iperf_binary_name=$1
       ;;
+    --sniff)
+      netstack_opts="${netstack_opts} -sniff"
+      ;;
     *)
       echo "unknown option: $1"
       echo ""
diff --git a/test/benchmarks/tcp/tcp_proxy.go b/test/benchmarks/tcp/tcp_proxy.go
index 28eacd9a7..75585ad8d 100644
--- a/test/benchmarks/tcp/tcp_proxy.go
+++ b/test/benchmarks/tcp/tcp_proxy.go
@@ -37,6 +37,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"gvisor.dev/gvisor/pkg/tcpip/link/fdbased"
 	"gvisor.dev/gvisor/pkg/tcpip/link/qdisc/fifo"
+	"gvisor.dev/gvisor/pkg/tcpip/link/sniffer"
 	"gvisor.dev/gvisor/pkg/tcpip/network/arp"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
@@ -72,6 +73,7 @@ var (
 	mutexprofile       = flag.String("mutexprofile", "", "write a mutex profile to the specified file.")
 	traceprofile       = flag.String("traceprofile", "", "write a 5s trace of the benchmark to the specified file.")
 	useIpv6            = flag.Bool("ipv6", false, "use ipv6 instead of ipv4.")
+	sniff              = flag.Bool("sniff", false, "log sniffed packets")
 )
 
 type impl interface {
@@ -229,6 +231,10 @@ func newNetstackImpl(mode string) (impl, error) {
 		return nil, fmt.Errorf("failed to create FD endpoint: %v", err)
 	}
 
+	if *sniff {
+		ep = sniffer.New(ep)
+	}
+
 	qDisc := fifo.New(ep, runtime.GOMAXPROCS(0), 1000)
 	opts := stack.NICOptions{QDisc: qDisc, GROTimeout: *gro}
 	if err := s.CreateNICWithOptions(nicID, ep, opts); err != nil {