Make GenerateRandomPayload available to all tests

Moved the function for generating a payload of random byets of a specified
length into the testbench package so that it's availbale for all tests to use.

Added a test case to the IPv4 ID uniqueness test which uses a payload length
of 512 bytes. This test case passes for gVisor currently, whereas the test case
with a small payload of 11 bytes fails because gVisor only assigns the ID field
if the IP payload is sufficiently large.

PiperOrigin-RevId: 316185097

3 files changed, 79 insertions, 65 deletions

diff --git a/test/packetimpact/testbench/testbench.go b/test/packetimpact/testbench/testbench.go
index 82eda6565..d64f32a5b 100644
--- a/test/packetimpact/testbench/testbench.go
+++ b/test/packetimpact/testbench/testbench.go
@@ -17,8 +17,10 @@ package testbench
 import (
 	"flag"
 	"fmt"
+	"math/rand"
 	"net"
 	"os/exec"
+	"testing"
 	"time"
 
 	"gvisor.dev/gvisor/test/packetimpact/netdevs"
@@ -91,3 +93,14 @@ func genPseudoFlags() error {
 
 	return nil
 }
+
+// GenerateRandomPayload generates a random byte slice of the specified length,
+// causing a fatal test failure if it is unable to do so.
+func GenerateRandomPayload(t *testing.T, n int) []byte {
+	t.Helper()
+	buf := make([]byte, n)
+	if _, err := rand.Read(buf); err != nil {
+		t.Fatalf("rand.Read(buf) failed: %s", err)
+	}
+	return buf
+}
diff --git a/test/packetimpact/tests/ipv4_id_uniqueness_test.go b/test/packetimpact/tests/ipv4_id_uniqueness_test.go
index 7db7ae02d..70f6df5e0 100644
--- a/test/packetimpact/tests/ipv4_id_uniqueness_test.go
+++ b/test/packetimpact/tests/ipv4_id_uniqueness_test.go
@@ -37,7 +37,7 @@ func recvTCPSegment(conn *testbench.TCPIPv4, expect *testbench.TCP, expectPayloa
 		return 0, fmt.Errorf("failed to receive TCP segment: %s", err)
 	}
 	if len(layers) < 2 {
-		return 0, fmt.Errorf("got packet with layers: %s, expected to have at least 2 layers (link and network)", layers)
+		return 0, fmt.Errorf("got packet with layers: %v, expected to have at least 2 layers (link and network)", layers)
 	}
 	ipv4, ok := layers[1].(*testbench.IPv4)
 	if !ok {
@@ -56,56 +56,67 @@ func recvTCPSegment(conn *testbench.TCPIPv4, expect *testbench.TCP, expectPayloa
 // to force the DF bit to be 0, and checks that a retransmitted segment has a
 // different IPv4 Identification value than the original segment.
 func TestIPv4RetransmitIdentificationUniqueness(t *testing.T) {
-	dut := testbench.NewDUT(t)
-	defer dut.TearDown()
+	for _, tc := range []struct {
+		name    string
+		payload []byte
+	}{
+		{"SmallPayload", []byte("sample data")},
+		// 512 bytes is chosen because sending more than this in a single segment
+		// causes the retransmission to send less than the original amount.
+		{"512BytePayload", testbench.GenerateRandomPayload(t, 512)},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			dut := testbench.NewDUT(t)
+			defer dut.TearDown()
 
-	listenFD, remotePort := dut.CreateListener(unix.SOCK_STREAM, unix.IPPROTO_TCP, 1)
-	defer dut.Close(listenFD)
+			listenFD, remotePort := dut.CreateListener(unix.SOCK_STREAM, unix.IPPROTO_TCP, 1)
+			defer dut.Close(listenFD)
 
-	conn := testbench.NewTCPIPv4(t, testbench.TCP{DstPort: &remotePort}, testbench.TCP{SrcPort: &remotePort})
-	defer conn.Close()
+			conn := testbench.NewTCPIPv4(t, testbench.TCP{DstPort: &remotePort}, testbench.TCP{SrcPort: &remotePort})
+			defer conn.Close()
 
-	conn.Connect()
-	remoteFD, _ := dut.Accept(listenFD)
-	defer dut.Close(remoteFD)
+			conn.Connect()
+			remoteFD, _ := dut.Accept(listenFD)
+			defer dut.Close(remoteFD)
 
-	dut.SetSockOptInt(remoteFD, unix.IPPROTO_TCP, unix.TCP_NODELAY, 1)
+			dut.SetSockOptInt(remoteFD, unix.IPPROTO_TCP, unix.TCP_NODELAY, 1)
 
-	// TODO(b/129291778) The following socket option clears the DF bit on
-	// IP packets sent over the socket, and is currently not supported by
-	// gVisor. gVisor by default sends packets with DF=0 anyway, so the
-	// socket option being not supported does not affect the operation of
-	// this test. Once the socket option is supported, the following call
-	// can be changed to simply assert success.
-	ret, errno := dut.SetSockOptIntWithErrno(context.Background(), remoteFD, unix.IPPROTO_IP, linux.IP_MTU_DISCOVER, linux.IP_PMTUDISC_DONT)
-	if ret == -1 && errno != unix.ENOTSUP {
-		t.Fatalf("failed to set IP_MTU_DISCOVER socket option to IP_PMTUDISC_DONT: %s", errno)
-	}
+			// TODO(b/129291778) The following socket option clears the DF bit on
+			// IP packets sent over the socket, and is currently not supported by
+			// gVisor. gVisor by default sends packets with DF=0 anyway, so the
+			// socket option being not supported does not affect the operation of
+			// this test. Once the socket option is supported, the following call
+			// can be changed to simply assert success.
+			ret, errno := dut.SetSockOptIntWithErrno(context.Background(), remoteFD, unix.IPPROTO_IP, linux.IP_MTU_DISCOVER, linux.IP_PMTUDISC_DONT)
+			if ret == -1 && errno != unix.ENOTSUP {
+				t.Fatalf("failed to set IP_MTU_DISCOVER socket option to IP_PMTUDISC_DONT: %s", errno)
+			}
 
-	sampleData := []byte("Sample Data")
-	samplePayload := &testbench.Payload{Bytes: sampleData}
+			samplePayload := &testbench.Payload{Bytes: tc.payload}
 
-	dut.Send(remoteFD, sampleData, 0)
-	if _, err := conn.ExpectData(&testbench.TCP{}, samplePayload, time.Second); err != nil {
-		t.Fatalf("failed to receive TCP segment sent for RTT calculation: %s", err)
-	}
-	// Let the DUT estimate RTO with RTT from the DATA-ACK.
-	// TODO(gvisor.dev/issue/2685) Estimate RTO during handshake, after which
-	// we can skip sending this ACK.
-	conn.Send(testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
+			dut.Send(remoteFD, tc.payload, 0)
+			if _, err := conn.ExpectData(&testbench.TCP{}, samplePayload, time.Second); err != nil {
+				t.Fatalf("failed to receive TCP segment sent for RTT calculation: %s", err)
+			}
+			// Let the DUT estimate RTO with RTT from the DATA-ACK.
+			// TODO(gvisor.dev/issue/2685) Estimate RTO during handshake, after which
+			// we can skip sending this ACK.
+			conn.Send(testbench.TCP{Flags: testbench.Uint8(header.TCPFlagAck)})
 
-	expectTCP := &testbench.TCP{SeqNum: testbench.Uint32(uint32(*conn.RemoteSeqNum()))}
-	dut.Send(remoteFD, sampleData, 0)
-	originalID, err := recvTCPSegment(&conn, expectTCP, samplePayload)
-	if err != nil {
-		t.Fatalf("failed to receive TCP segment: %s", err)
-	}
+			dut.Send(remoteFD, tc.payload, 0)
+			expectTCP := &testbench.TCP{SeqNum: testbench.Uint32(uint32(*conn.RemoteSeqNum()))}
+			originalID, err := recvTCPSegment(&conn, expectTCP, samplePayload)
+			if err != nil {
+				t.Fatalf("failed to receive TCP segment: %s", err)
+			}
 
-	retransmitID, err := recvTCPSegment(&conn, expectTCP, samplePayload)
-	if err != nil {
-		t.Fatalf("failed to receive retransmitted TCP segment: %s", err)
-	}
-	if originalID == retransmitID {
-		t.Fatalf("unexpectedly got retransmitted TCP segment with same IPv4 ID field=%d", originalID)
+			retransmitID, err := recvTCPSegment(&conn, expectTCP, samplePayload)
+			if err != nil {
+				t.Fatalf("failed to receive retransmitted TCP segment: %s", err)
+			}
+			if originalID == retransmitID {
+				t.Fatalf("unexpectedly got retransmitted TCP segment with same IPv4 ID field=%d", originalID)
+			}
+		})
 	}
 }
diff --git a/test/packetimpact/tests/udp_send_recv_dgram_test.go b/test/packetimpact/tests/udp_send_recv_dgram_test.go
index a7db384ad..224feef85 100644
--- a/test/packetimpact/tests/udp_send_recv_dgram_test.go
+++ b/test/packetimpact/tests/udp_send_recv_dgram_test.go
@@ -16,7 +16,6 @@ package udp_send_recv_dgram_test
 
 import (
 	"flag"
-	"math/rand"
 	"net"
 	"testing"
 	"time"
@@ -29,15 +28,6 @@ func init() {
 	testbench.RegisterFlags(flag.CommandLine)
 }
 
-func generateRandomPayload(t *testing.T, n int) string {
-	t.Helper()
-	buf := make([]byte, n)
-	if _, err := rand.Read(buf); err != nil {
-		t.Fatalf("rand.Read(buf) failed: %s", err)
-	}
-	return string(buf)
-}
-
 func TestUDPRecv(t *testing.T) {
 	dut := testbench.NewDUT(t)
 	defer dut.TearDown()
@@ -48,19 +38,19 @@ func TestUDPRecv(t *testing.T) {
 
 	testCases := []struct {
 		name    string
-		payload string
+		payload []byte
 	}{
-		{"emptypayload", ""},
-		{"small payload", "hello world"},
-		{"1kPayload", generateRandomPayload(t, 1<<10)},
+		{"emptypayload", nil},
+		{"small payload", []byte("hello world")},
+		{"1kPayload", testbench.GenerateRandomPayload(t, 1<<10)},
 		// Even though UDP allows larger dgrams we don't test it here as
 		// they need to be fragmented and written out as individual
 		// frames.
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			conn.Send(testbench.UDP{}, &testbench.Payload{Bytes: []byte(tc.payload)})
-			if got, want := string(dut.Recv(boundFD, int32(len(tc.payload)), 0)), tc.payload; got != want {
+			conn.Send(testbench.UDP{}, &testbench.Payload{Bytes: tc.payload})
+			if got, want := string(dut.Recv(boundFD, int32(len(tc.payload)), 0)), string(tc.payload); got != want {
 				t.Fatalf("received payload does not match sent payload got: %s, want: %s", got, want)
 			}
 		})
@@ -77,11 +67,11 @@ func TestUDPSend(t *testing.T) {
 
 	testCases := []struct {
 		name    string
-		payload string
+		payload []byte
 	}{
-		{"emptypayload", ""},
-		{"small payload", "hello world"},
-		{"1kPayload", generateRandomPayload(t, 1<<10)},
+		{"emptypayload", nil},
+		{"small payload", []byte("hello world")},
+		{"1kPayload", testbench.GenerateRandomPayload(t, 1<<10)},
 		// Even though UDP allows larger dgrams we don't test it here as
 		// they need to be fragmented and written out as individual
 		// frames.
@@ -89,10 +79,10 @@ func TestUDPSend(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			conn.Drain()
-			if got, want := int(dut.SendTo(boundFD, []byte(tc.payload), 0, conn.LocalAddr())), len(tc.payload); got != want {
+			if got, want := int(dut.SendTo(boundFD, tc.payload, 0, conn.LocalAddr())), len(tc.payload); got != want {
 				t.Fatalf("short write got: %d, want: %d", got, want)
 			}
-			if _, err := conn.ExpectData(testbench.UDP{SrcPort: &remotePort}, testbench.Payload{Bytes: []byte(tc.payload)}, 1*time.Second); err != nil {
+			if _, err := conn.ExpectData(testbench.UDP{SrcPort: &remotePort}, testbench.Payload{Bytes: tc.payload}, 1*time.Second); err != nil {
 				t.Fatal(err)
 			}
 		})