3 files changed, 49 insertions, 13 deletions
diff --git a/test/packetimpact/testbench/testbench.go b/test/packetimpact/testbench/testbench.go
index d64f32a5b..6530036a8 100644
--- a/test/packetimpact/testbench/testbench.go
+++ b/test/packetimpact/testbench/testbench.go
@@ -31,23 +31,30 @@ var (
 	DUTType = ""
 	// Device is the local device on the test network.
 	Device = ""
+
 	// LocalIPv4 is the local IPv4 address on the test network.
 	LocalIPv4 = ""
+	// RemoteIPv4 is the DUT's IPv4 address on the test network.
+	RemoteIPv4 = ""
+	// IPv4PrefixLength is the network prefix length of the IPv4 test network.
+	IPv4PrefixLength = 0
+
 	// LocalIPv6 is the local IPv6 address on the test network.
 	LocalIPv6 = ""
+	// RemoteIPv6 is the DUT's IPv6 address on the test network.
+	RemoteIPv6 = ""
+
 	// LocalMAC is the local MAC address on the test network.
 	LocalMAC = ""
+	// RemoteMAC is the DUT's MAC address on the test network.
+	RemoteMAC = ""
+
 	// POSIXServerIP is the POSIX server's IP address on the control network.
 	POSIXServerIP = ""
 	// POSIXServerPort is the UDP port the POSIX server is bound to on the
 	// control network.
 	POSIXServerPort = 40000
-	// RemoteIPv4 is the DUT's IPv4 address on the test network.
-	RemoteIPv4 = ""
-	// RemoteIPv6 is the DUT's IPv6 address on the test network.
-	RemoteIPv6 = ""
-	// RemoteMAC is the DUT's MAC address on the test network.
-	RemoteMAC = ""
+
 	// RPCKeepalive is the gRPC keepalive.
 	RPCKeepalive = 10 * time.Second
 	// RPCTimeout is the gRPC timeout.
@@ -91,6 +98,12 @@ func genPseudoFlags() error {
 	LocalMAC = deviceInfo.MAC.String()
 	LocalIPv6 = deviceInfo.IPv6Addr.String()
 
+	if deviceInfo.IPv4Net != nil {
+		IPv4PrefixLength, _ = deviceInfo.IPv4Net.Mask.Size()
+	} else {
+		IPv4PrefixLength, _ = net.ParseIP(LocalIPv4).DefaultMask().Size()
+	}
+
 	return nil
 }
 
diff --git a/test/packetimpact/tests/BUILD b/test/packetimpact/tests/BUILD
index 3ecbe83eb..3f538b5c6 100644
--- a/test/packetimpact/tests/BUILD
+++ b/test/packetimpact/tests/BUILD
@@ -27,8 +27,8 @@ packetimpact_go_test(
 )
 
 packetimpact_go_test(
-    name = "udp_recv_multicast",
-    srcs = ["udp_recv_multicast_test.go"],
+    name = "udp_recv_mcast_bcast",
+    srcs = ["udp_recv_mcast_bcast_test.go"],
     # TODO(b/152813495): Fix netstack then remove the line below.
     expect_netstack_failure = True,
     deps = [
diff --git a/test/packetimpact/tests/udp_recv_multicast_test.go b/test/packetimpact/tests/udp_recv_mcast_bcast_test.go
index 77a9bfa1d..263a54291 100644
--- a/test/packetimpact/tests/udp_recv_multicast_test.go
+++ b/test/packetimpact/tests/udp_recv_mcast_bcast_test.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package udp_recv_multicast_test
+package udp_recv_mcast_bcast_test
 
 import (
 	"flag"
@@ -28,13 +28,36 @@ func init() {
 	testbench.RegisterFlags(flag.CommandLine)
 }
 
-func TestUDPRecvMulticast(t *testing.T) {
+func TestUDPRecvMulticastBroadcast(t *testing.T) {
 	dut := testbench.NewDUT(t)
 	defer dut.TearDown()
-	boundFD, remotePort := dut.CreateBoundSocket(unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.ParseIP("0.0.0.0"))
+	boundFD, remotePort := dut.CreateBoundSocket(unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.IPv4(0, 0, 0, 0))
 	defer dut.Close(boundFD)
 	conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
 	defer conn.Close()
-	conn.SendIP(testbench.IPv4{DstAddr: testbench.Address(tcpip.Address(net.ParseIP("224.0.0.1").To4()))}, testbench.UDP{})
-	dut.Recv(boundFD, 100, 0)
+
+	for _, bcastAddr := range []net.IP{
+		broadcastAddr(net.ParseIP(testbench.RemoteIPv4), net.CIDRMask(testbench.IPv4PrefixLength, 32)),
+		net.IPv4(255, 255, 255, 255),
+		net.IPv4(224, 0, 0, 1),
+	} {
+		payload := testbench.GenerateRandomPayload(t, 1<<10)
+		conn.SendIP(
+			testbench.IPv4{DstAddr: testbench.Address(tcpip.Address(bcastAddr.To4()))},
+			testbench.UDP{},
+			&testbench.Payload{Bytes: payload},
+		)
+		t.Logf("Receiving packet sent to address: %s", bcastAddr)
+		if got, want := string(dut.Recv(boundFD, int32(len(payload)), 0)), string(payload); got != want {
+			t.Errorf("received payload does not match sent payload got: %s, want: %s", got, want)
+		}
+	}
+}
+
+func broadcastAddr(ip net.IP, mask net.IPMask) net.IP {
+	ip4 := ip.To4()
+	for i := range ip4 {
+		ip4[i] |= ^mask[i]
+	}
+	return ip4
 }