Add UDP send/recv packetimpact tests.

Fixes #2654

PiperOrigin-RevId: 310642216

5 files changed, 120 insertions, 5 deletions

diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 56ac3fa54..28c841612 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -841,6 +841,7 @@ func (conn *UDPIPv4) SendIP(additionalLayers ...Layer) {
 // Expect expects a frame with the UDP layer matching the provided UDP within
 // the timeout specified. If it doesn't arrive in time, an error is returned.
 func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
+	conn.t.Helper()
 	layer, err := (*Connection)(conn).Expect(&udp, timeout)
 	if layer == nil {
 		return nil, err
@@ -852,6 +853,18 @@ func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
 	return gotUDP, err
 }
 
+// ExpectData is a convenient method that expects a Layer and the Layer after
+// it. If it doens't arrive in time, it returns nil.
+func (conn *UDPIPv4) ExpectData(udp UDP, payload Payload, timeout time.Duration) (Layers, error) {
+	conn.t.Helper()
+	expected := make([]Layer, len(conn.layerStates))
+	expected[len(expected)-1] = &udp
+	if payload.length() != 0 {
+		expected = append(expected, &payload)
+	}
+	return (*Connection)(conn).ExpectFrame(expected, timeout)
+}
+
 // Close frees associated resources held by the UDPIPv4 connection.
 func (conn *UDPIPv4) Close() {
 	(*Connection)(conn).Close()
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index 165f62d3b..49370377d 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -898,10 +898,7 @@ func (l *UDP) match(other Layer) bool {
 }
 
 func (l *UDP) length() int {
-	if l.Length == nil {
-		return header.UDPMinimumSize
-	}
-	return int(*l.Length)
+	return header.UDPMinimumSize
 }
 
 // merge implements Layer.merge.
diff --git a/test/packetimpact/testbench/rawsockets.go b/test/packetimpact/testbench/rawsockets.go
index ff722d4a6..a9ad72b63 100644
--- a/test/packetimpact/testbench/rawsockets.go
+++ b/test/packetimpact/testbench/rawsockets.go
@@ -169,7 +169,7 @@ func NewInjector(t *testing.T) (Injector, error) {
 // Send a raw frame.
 func (i *Injector) Send(b []byte) {
 	if _, err := unix.Write(i.fd, b); err != nil {
-		i.t.Fatalf("can't write: %s", err)
+		i.t.Fatalf("can't write: %s of len %d", err, len(b))
 	}
 }
 
diff --git a/test/packetimpact/tests/BUILD b/test/packetimpact/tests/BUILD
index 563823d04..c25b3b8c1 100644
--- a/test/packetimpact/tests/BUILD
+++ b/test/packetimpact/tests/BUILD
@@ -148,6 +148,15 @@ packetimpact_go_test(
     ],
 )
 
+packetimpact_go_test(
+    name = "udp_send_recv_dgram",
+    srcs = ["udp_send_recv_dgram_test.go"],
+    deps = [
+        "//test/packetimpact/testbench",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
 sh_binary(
     name = "test_runner",
     srcs = ["test_runner.sh"],
diff --git a/test/packetimpact/tests/udp_send_recv_dgram_test.go b/test/packetimpact/tests/udp_send_recv_dgram_test.go
new file mode 100644
index 000000000..1c682831c
--- /dev/null
+++ b/test/packetimpact/tests/udp_send_recv_dgram_test.go
@@ -0,0 +1,96 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package udp_send_recv_dgram_test
+
+import (
+	"math/rand"
+	"net"
+	"testing"
+	"time"
+
+	"golang.org/x/sys/unix"
+	tb "gvisor.dev/gvisor/test/packetimpact/testbench"
+)
+
+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 := tb.NewDUT(t)
+	defer dut.TearDown()
+	boundFD, remotePort := dut.CreateBoundSocket(unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.ParseIP("0.0.0.0"))
+	defer dut.Close(boundFD)
+	conn := tb.NewUDPIPv4(t, tb.UDP{DstPort: &remotePort}, tb.UDP{SrcPort: &remotePort})
+	defer conn.Close()
+
+	testCases := []struct {
+		name    string
+		payload string
+	}{
+		{"emptypayload", ""},
+		{"small payload", "hello world"},
+		{"1kPayload", 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) {
+			frame := conn.CreateFrame(&tb.UDP{}, &tb.Payload{Bytes: []byte(tc.payload)})
+			conn.SendFrame(frame)
+			if got, want := string(dut.Recv(boundFD, int32(len(tc.payload)), 0)), tc.payload; got != want {
+				t.Fatalf("received payload does not match sent payload got: %s, want: %s", got, want)
+			}
+		})
+	}
+}
+
+func TestUDPSend(t *testing.T) {
+	dut := tb.NewDUT(t)
+	defer dut.TearDown()
+	boundFD, remotePort := dut.CreateBoundSocket(unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.ParseIP("0.0.0.0"))
+	defer dut.Close(boundFD)
+	conn := tb.NewUDPIPv4(t, tb.UDP{DstPort: &remotePort}, tb.UDP{SrcPort: &remotePort})
+	defer conn.Close()
+
+	testCases := []struct {
+		name    string
+		payload string
+	}{
+		{"emptypayload", ""},
+		{"small payload", "hello world"},
+		{"1kPayload", 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.Drain()
+			if got, want := int(dut.SendTo(boundFD, []byte(tc.payload), 0, conn.LocalAddr())), len(tc.payload); got != want {
+				t.Fatalf("short write got: %d, want: %d", got, want)
+			}
+			if _, err := conn.ExpectData(tb.UDP{SrcPort: &remotePort}, tb.Payload{Bytes: []byte(tc.payload)}, 1*time.Second); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}