Allow DUT binaries to define their own flags

Updates #6835

PiperOrigin-RevId: 413772074

test/packetimpact/dut/BUILD

@@ -33,6 +33,10 @@ go_library(
     name = "dut",
     testonly = True,
     srcs = ["dut.go"],
+    visibility = [
+        "//test/packetimpact:__subpackages__",
+        "//turquoise/connectivity/netstack/gvisor_tests/packetimpact:__subpackages__",
+    ],
     deps = [
         "//test/packetimpact/testbench",
         "@org_golang_x_sync//errgroup:go_default_library",

test/packetimpact/dut/dut.go

@@ -39,6 +39,10 @@ const (
 	completeFd = 3
 	// PosixServerPort is the port the posix server should listen on.
 	PosixServerPort = 54321
+	// CtrlIface is the command switch name for passing name of the control interface.
+	CtrlIface = "ctrl_iface"
+	// TestIface is the command switch name for passing name of the test interface.
+	TestIface = "test_iface"
 )
 
 // Ifaces describe the names of the interfaces on DUT.
@@ -51,18 +55,15 @@ type Ifaces struct {
 
 // Init puts the current process into the target network namespace, the user of
 // this library should call this function in the beginning.
-func Init() (Ifaces, error) {
+func Init(fs *flag.FlagSet) (Ifaces, error) {
 	// The DUT might create child processes, we don't want this fd to leak into
 	// those processes as it keeps the pipe open and the testbench will hang
 	// waiting for an EOF on the pipe.
 	unix.CloseOnExec(completeFd)
 	var ifaces Ifaces
-	// Parse command line flags. It is effectively the same as using top-level
-	// functions in flag package, but more explicit that we exit if the parsing
-	// failed.
-	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
-	fs.StringVar(&ifaces.Ctrl, "ctrl_iface", "", "the name of the control interface")
-	fs.StringVar(&ifaces.Test, "test_iface", "", "the name of the test interface")
+	// Parse command line flags that is defined by the caller and us.
+	fs.StringVar(&ifaces.Ctrl, CtrlIface, "", "the name of the control interface")
+	fs.StringVar(&ifaces.Test, TestIface, "", "the name of the test interface")
 	if err := fs.Parse(os.Args[1:]); err != nil {
 		return Ifaces{}, err
 	}

test/packetimpact/dut/native/main.go

@@ -20,6 +20,7 @@ package main
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -40,7 +41,7 @@ type native struct {
 }
 
 func main() {
-	ifaces, err := dut.Init()
+	ifaces, err := dut.Init(flag.CommandLine)
 	if err != nil {
 		log.Fatal(err)
 	}

test/packetimpact/dut/runsc/main.go

@@ -20,6 +20,7 @@ package main
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -49,7 +50,7 @@ type runsc struct {
 var _ dut.DUT = (*runsc)(nil)
 
 func main() {
-	ifaces, err := dut.Init()
+	ifaces, err := dut.Init(flag.CommandLine)
 	if err != nil {
 		log.Fatal(err)
 	}

test/packetimpact/internal/testing/BUILD

@@ -1,7 +1,6 @@
 load("//tools:defs.bzl", "go_library")
 
 package(
-    default_visibility = ["//test/packetimpact:__subpackages__"],
     licenses = ["notice"],
 )
 
@@ -9,4 +8,8 @@ go_library(
     name = "testing",
     testonly = True,
     srcs = ["testing.go"],
+    visibility = [
+        "//test/packetimpact:__subpackages__",
+        "//turquoise/connectivity/netstack/gvisor_tests/packetimpact:__subpackages__",
+    ],
 )

test/packetimpact/runner/BUILD

@@ -41,7 +41,12 @@ go_binary(
     name = "main",
     testonly = True,
     srcs = ["main.go"],
+    visibility = [
+        "//test/packetimpact:__subpackages__",
+        "//turquoise/connectivity/netstack/gvisor_tests/packetimpact:__subpackages__",
+    ],
     deps = [
+        "//test/packetimpact/dut",
         "//test/packetimpact/internal/testing",
         "//test/packetimpact/netdevs/netlink",
         "//test/packetimpact/testbench",

test/packetimpact/runner/main.go

@@ -30,6 +30,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"syscall"
 
 	"github.com/google/gopacket"
@@ -38,11 +39,25 @@ import (
 	"github.com/vishvananda/netlink"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/test/packetimpact/dut"
 	"gvisor.dev/gvisor/test/packetimpact/internal/testing"
 	netdevs "gvisor.dev/gvisor/test/packetimpact/netdevs/netlink"
 	"gvisor.dev/gvisor/test/packetimpact/testbench"
 )
 
+type dutArgList []string
+
+// String implements flag.Value.
+func (l *dutArgList) String() string {
+	return strings.Join(*l, " ")
+}
+
+// Set implements flag.Value.
+func (l *dutArgList) Set(value string) error {
+	*l = append(*l, value)
+	return nil
+}
+
 func main() {
 	const procSelfExe = "/proc/self/exe"
 	if os.Args[0] != procSelfExe {
@@ -88,6 +103,7 @@ func main() {
 		runtime         string
 		partition       int
 		totalPartitions int
+		dutArgs         dutArgList
 	)
 	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
 	fs.StringVar(&dutBinary, "dut_binary", "", "path to the DUT binary")
@@ -95,6 +111,7 @@ func main() {
 	fs.BoolVar(&expectFailure, "expect_failure", false, "whether the test is expected to fail")
 	fs.IntVar(&numDUTs, "num_duts", 1, "number of DUTs to create")
 	fs.StringVar(&variant, "variant", "", "test variant could be native, gvisor or fuchsia")
+	fs.Var(&dutArgs, "dut_arg", "argument to the DUT binary")
 	// The following args are passed by CI environment which are not used by us.
 	fs.StringVar(&runtime, "runtime", "", "docker runtime to use (unused)")
 	fs.IntVar(&partition, "partition", 1, "1-indexed partition (unused)")
@@ -107,9 +124,9 @@ func main() {
 
 	// Create all the DUTs.
 	infoCh := make(chan testbench.DUTInfo, numDUTs)
-	var duts []*dut
+	var duts []*dutProcess
 	for i := 0; i < numDUTs; i++ {
-		d, err := newDUT(ctx, i, dutBinary)
+		d, err := newDUT(ctx, i, dutBinary, dutArgs)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -194,15 +211,18 @@ func main() {
 	}
 }
 
-type dut struct {
+type dutProcess struct {
 	cmd       *exec.Cmd
 	id        int
 	completeR *os.File
 	dutNetNS  netNS
 }
 
-func newDUT(ctx context.Context, id int, dutBinary string) (*dut, error) {
-	cmd := exec.CommandContext(ctx, dutBinary, "--ctrl_iface", dutSide.ifaceName(ctrlLink, id), "--test_iface", dutSide.ifaceName(testLink, id))
+func newDUT(ctx context.Context, id int, dutBinary string, dutArgs dutArgList) (*dutProcess, error) {
+	cmd := exec.CommandContext(ctx, dutBinary, append([]string{
+		"--" + dut.CtrlIface, dutSide.ifaceName(ctrlLink, id),
+		"--" + dut.TestIface, dutSide.ifaceName(testLink, id),
+	}, dutArgs...)...)
 
 	// Create the pipe for completion signal
 	completeR, completeW, err := os.Pipe()
@@ -297,10 +317,10 @@ func newDUT(ctx context.Context, id int, dutBinary string) (*dut, error) {
 		}
 	}
 
-	return &dut{cmd: cmd, id: id, completeR: completeR, dutNetNS: dutNetNS}, nil
+	return &dutProcess{cmd: cmd, id: id, completeR: completeR, dutNetNS: dutNetNS}, nil
 }
 
-func (d *dut) bootstrap(ctx context.Context) (testbench.DUTInfo, func() error, error) {
+func (d *dutProcess) bootstrap(ctx context.Context) (testbench.DUTInfo, func() error, error) {
 	if err := d.dutNetNS.Do(func() error {
 		return d.cmd.Start()
 	}); err != nil {
@@ -335,16 +355,16 @@ func (d *dut) bootstrap(ctx context.Context) (testbench.DUTInfo, func() error, e
 	return dutInfo, d.cmd.Wait, nil
 }
 
-func (d *dut) name() string {
+func (d *dutProcess) name() string {
 	return fmt.Sprintf("dut-%d", d.id)
 }
 
-func (d *dut) peerIface() string {
+func (d *dutProcess) peerIface() string {
 	return tbSide.ifaceName(testLink, d.id)
 }
 
 // writePcap creates the packet capture while the test is running.
-func (d *dut) writePcap(ctx context.Context, testName string) error {
+func (d *dutProcess) writePcap(ctx context.Context, testName string) error {
 	iface := d.peerIface()
 	// Create the pcap file.
 	fileName, err := testing.UndeclaredOutput(fmt.Sprintf("%s_%s.pcap", testName, iface))