Add support for UDP IPv6

Also ironed out all the bugs found on the IPv6 code path that affects socket
bind, send and receive.

PiperOrigin-RevId: 321202653

10 files changed, 464 insertions, 73 deletions

diff --git a/test/packetimpact/dut/posix_server.cc b/test/packetimpact/dut/posix_server.cc
index a1a5c3612..29d4cc6fe 100644
--- a/test/packetimpact/dut/posix_server.cc
+++ b/test/packetimpact/dut/posix_server.cc
@@ -53,7 +53,10 @@
       response_in6->set_flowinfo(ntohl(addr_in6->sin6_flowinfo));
       response_in6->mutable_addr()->assign(
           reinterpret_cast<const char *>(&addr_in6->sin6_addr.s6_addr), 16);
-      response_in6->set_scope_id(ntohl(addr_in6->sin6_scope_id));
+      // sin6_scope_id is stored in host byte order.
+      //
+      // https://www.gnu.org/software/libc/manual/html_node/Internet-Address-Formats.html
+      response_in6->set_scope_id(addr_in6->sin6_scope_id);
       return ::grpc::Status::OK;
     }
   }
@@ -89,7 +92,10 @@
       addr_in6->sin6_flowinfo = htonl(proto_in6.flowinfo());
       proto_in6.addr().copy(
           reinterpret_cast<char *>(&addr_in6->sin6_addr.s6_addr), 16);
-      addr_in6->sin6_scope_id = htonl(proto_in6.scope_id());
+      // sin6_scope_id is stored in host byte order.
+      //
+      // https://www.gnu.org/software/libc/manual/html_node/Internet-Address-Formats.html
+      addr_in6->sin6_scope_id = proto_in6.scope_id();
       *addr_len = sizeof(*addr_in6);
       break;
     }
diff --git a/test/packetimpact/netdevs/BUILD b/test/packetimpact/netdevs/BUILD
index 422bb9b0c..8d1193fed 100644
--- a/test/packetimpact/netdevs/BUILD
+++ b/test/packetimpact/netdevs/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library")
+load("//tools:defs.bzl", "go_library", "go_test")
 
 package(
     licenses = ["notice"],
@@ -13,3 +13,11 @@ go_library(
         "//pkg/tcpip/header",
     ],
 )
+
+go_test(
+    name = "netdevs_test",
+    size = "small",
+    srcs = ["netdevs_test.go"],
+    library = ":netdevs",
+    deps = ["@com_github_google_go_cmp//cmp:go_default_library"],
+)
diff --git a/test/packetimpact/netdevs/netdevs.go b/test/packetimpact/netdevs/netdevs.go
index d2c9cfeaf..eecfe0730 100644
--- a/test/packetimpact/netdevs/netdevs.go
+++ b/test/packetimpact/netdevs/netdevs.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"net"
 	"regexp"
+	"strconv"
 	"strings"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -27,6 +28,7 @@ import (
 
 // A DeviceInfo represents a network device.
 type DeviceInfo struct {
+	ID       uint32
 	MAC      net.HardwareAddr
 	IPv4Addr net.IP
 	IPv4Net  *net.IPNet
@@ -35,7 +37,7 @@ type DeviceInfo struct {
 }
 
 var (
-	deviceLine = regexp.MustCompile(`^\s*\d+: (\w+)`)
+	deviceLine = regexp.MustCompile(`^\s*(\d+): (\w+)`)
 	linkLine   = regexp.MustCompile(`^\s*link/\w+ ([0-9a-fA-F:]+)`)
 	inetLine   = regexp.MustCompile(`^\s*inet ([0-9./]+)`)
 	inet6Line  = regexp.MustCompile(`^\s*inet6 ([0-9a-fA-Z:/]+)`)
@@ -43,6 +45,11 @@ var (
 
 // ParseDevices parses the output from `ip addr show` into a map from device
 // name to information about the device.
+//
+// Note: if multiple IPv6 addresses are assigned to a device, the last address
+// displayed by `ip addr show` will be used. This is fine for packetimpact
+// because we will always only have at most one IPv6 address assigned to each
+// device.
 func ParseDevices(cmdOutput string) (map[string]DeviceInfo, error) {
 	var currentDevice string
 	var currentInfo DeviceInfo
@@ -52,8 +59,12 @@ func ParseDevices(cmdOutput string) (map[string]DeviceInfo, error) {
 			if currentDevice != "" {
 				deviceInfos[currentDevice] = currentInfo
 			}
-			currentInfo = DeviceInfo{}
-			currentDevice = m[1]
+			id, err := strconv.ParseUint(m[1], 10, 32)
+			if err != nil {
+				return nil, fmt.Errorf("parsing device ID %s: %w", m[1], err)
+			}
+			currentInfo = DeviceInfo{ID: uint32(id)}
+			currentDevice = m[2]
 		} else if m := linkLine.FindStringSubmatch(line); m != nil {
 			mac, err := net.ParseMAC(m[1])
 			if err != nil {
diff --git a/test/packetimpact/netdevs/netdevs_test.go b/test/packetimpact/netdevs/netdevs_test.go
new file mode 100644
index 000000000..24ad12198
--- /dev/null
+++ b/test/packetimpact/netdevs/netdevs_test.go
@@ -0,0 +1,227 @@
+// 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 netdevs
+
+import (
+	"fmt"
+	"net"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func mustParseMAC(s string) net.HardwareAddr {
+	mac, err := net.ParseMAC(s)
+	if err != nil {
+		panic(fmt.Sprintf("failed to parse test MAC %q: %s", s, err))
+	}
+	return mac
+}
+
+func TestParseDevices(t *testing.T) {
+	for _, v := range []struct {
+		desc      string
+		cmdOutput string
+		want      map[string]DeviceInfo
+	}{
+		{
+			desc: "v4 and v6",
+			cmdOutput: `
+1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue state UNKNOWN group default qlen 1000
+  link/loopback 00:00:00:00:00:00 brd 00:00:00:00:00:00
+  inet 127.0.0.1/8 scope host lo
+    valid_lft forever preferred_lft forever
+  inet6 ::1/128 scope host
+    valid_lft forever preferred_lft forever
+2613: eth0@if2614: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:c0:a8:09:02 brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 192.168.9.2/24 brd 192.168.9.255 scope global eth0
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:c0ff:fea8:902/64 scope link tentative
+    valid_lft forever preferred_lft forever
+2615: eth2@if2616: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:df:f5:e1:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 223.245.225.10/24 brd 223.245.225.255 scope global eth2
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:dfff:fef5:e10a/64 scope link tentative
+    valid_lft forever preferred_lft forever
+2617: eth1@if2618: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:da:33:13:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 218.51.19.10/24 brd 218.51.19.255 scope global eth1
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:daff:fe33:130a/64 scope link tentative
+    valid_lft forever preferred_lft forever`,
+			want: map[string]DeviceInfo{
+				"lo": DeviceInfo{
+					ID:       1,
+					MAC:      mustParseMAC("00:00:00:00:00:00"),
+					IPv4Addr: net.IPv4(127, 0, 0, 1),
+					IPv4Net: &net.IPNet{
+						IP:   net.IPv4(127, 0, 0, 0),
+						Mask: net.CIDRMask(8, 32),
+					},
+					IPv6Addr: net.ParseIP("::1"),
+					IPv6Net: &net.IPNet{
+						IP:   net.ParseIP("::1"),
+						Mask: net.CIDRMask(128, 128),
+					},
+				},
+				"eth0": DeviceInfo{
+					ID:       2613,
+					MAC:      mustParseMAC("02:42:c0:a8:09:02"),
+					IPv4Addr: net.IPv4(192, 168, 9, 2),
+					IPv4Net: &net.IPNet{
+						IP:   net.IPv4(192, 168, 9, 0),
+						Mask: net.CIDRMask(24, 32),
+					},
+					IPv6Addr: net.ParseIP("fe80::42:c0ff:fea8:902"),
+					IPv6Net: &net.IPNet{
+						IP:   net.ParseIP("fe80::"),
+						Mask: net.CIDRMask(64, 128),
+					},
+				},
+				"eth1": DeviceInfo{
+					ID:       2617,
+					MAC:      mustParseMAC("02:42:da:33:13:0a"),
+					IPv4Addr: net.IPv4(218, 51, 19, 10),
+					IPv4Net: &net.IPNet{
+						IP:   net.IPv4(218, 51, 19, 0),
+						Mask: net.CIDRMask(24, 32),
+					},
+					IPv6Addr: net.ParseIP("fe80::42:daff:fe33:130a"),
+					IPv6Net: &net.IPNet{
+						IP:   net.ParseIP("fe80::"),
+						Mask: net.CIDRMask(64, 128),
+					},
+				},
+				"eth2": DeviceInfo{
+					ID:       2615,
+					MAC:      mustParseMAC("02:42:df:f5:e1:0a"),
+					IPv4Addr: net.IPv4(223, 245, 225, 10),
+					IPv4Net: &net.IPNet{
+						IP:   net.IPv4(223, 245, 225, 0),
+						Mask: net.CIDRMask(24, 32),
+					},
+					IPv6Addr: net.ParseIP("fe80::42:dfff:fef5:e10a"),
+					IPv6Net: &net.IPNet{
+						IP:   net.ParseIP("fe80::"),
+						Mask: net.CIDRMask(64, 128),
+					},
+				},
+			},
+		},
+		{
+			desc: "v4 only",
+			cmdOutput: `
+2613: eth0@if2614: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:c0:a8:09:02 brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 192.168.9.2/24 brd 192.168.9.255 scope global eth0
+    valid_lft forever preferred_lft forever`,
+			want: map[string]DeviceInfo{
+				"eth0": DeviceInfo{
+					ID:       2613,
+					MAC:      mustParseMAC("02:42:c0:a8:09:02"),
+					IPv4Addr: net.IPv4(192, 168, 9, 2),
+					IPv4Net: &net.IPNet{
+						IP:   net.IPv4(192, 168, 9, 0),
+						Mask: net.CIDRMask(24, 32),
+					},
+				},
+			},
+		},
+		{
+			desc: "v6 only",
+			cmdOutput: `
+2615: eth2@if2616: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:df:f5:e1:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet6 fe80::42:dfff:fef5:e10a/64 scope link tentative
+    valid_lft forever preferred_lft forever`,
+			want: map[string]DeviceInfo{
+				"eth2": DeviceInfo{
+					ID:       2615,
+					MAC:      mustParseMAC("02:42:df:f5:e1:0a"),
+					IPv6Addr: net.ParseIP("fe80::42:dfff:fef5:e10a"),
+					IPv6Net: &net.IPNet{
+						IP:   net.ParseIP("fe80::"),
+						Mask: net.CIDRMask(64, 128),
+					},
+				},
+			},
+		},
+	} {
+		t.Run(v.desc, func(t *testing.T) {
+			got, err := ParseDevices(v.cmdOutput)
+			if err != nil {
+				t.Errorf("ParseDevices(\n%s\n) got unexpected error: %s", v.cmdOutput, err)
+			}
+			if diff := cmp.Diff(v.want, got); diff != "" {
+				t.Errorf("ParseDevices(\n%s\n) got output diff (-want, +got):\n%s", v.cmdOutput, diff)
+			}
+		})
+	}
+}
+
+func TestParseDevicesErrors(t *testing.T) {
+	for _, v := range []struct {
+		desc      string
+		cmdOutput string
+	}{
+		{
+			desc: "invalid MAC addr",
+			cmdOutput: `
+2617: eth1@if2618: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+	link/ether 02:42:da:33:13:0a:ffffffff brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 218.51.19.10/24 brd 218.51.19.255 scope global eth1
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:daff:fe33:130a/64 scope link tentative
+    valid_lft forever preferred_lft forever`,
+		},
+		{
+			desc: "invalid v4 addr",
+			cmdOutput: `
+2617: eth1@if2618: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:da:33:13:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 1234.4321.424242.0/24 brd 218.51.19.255 scope global eth1
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:daff:fe33:130a/64 scope link tentative
+    valid_lft forever preferred_lft forever`,
+		},
+		{
+			desc: "invalid v6 addr",
+			cmdOutput: `
+2617: eth1@if2618: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:da:33:13:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 218.51.19.10/24 brd 218.51.19.255 scope global eth1
+    valid_lft forever preferred_lft forever
+		inet6 fe80:ffffffff::42:daff:fe33:130a/64 scope link tentative
+    valid_lft forever preferred_lft forever`,
+		},
+		{
+			desc: "invalid CIDR missing prefixlen",
+			cmdOutput: `
+2617: eth1@if2618: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc noqueue state UP group default
+  link/ether 02:42:da:33:13:0a brd ff:ff:ff:ff:ff:ff link-netnsid 0
+  inet 218.51.19.10 brd 218.51.19.255 scope global eth1
+    valid_lft forever preferred_lft forever
+  inet6 fe80::42:daff:fe33:130a scope link tentative
+    valid_lft forever preferred_lft forever`,
+		},
+	} {
+		t.Run(v.desc, func(t *testing.T) {
+			if _, err := ParseDevices(v.cmdOutput); err == nil {
+				t.Errorf("ParseDevices(\n%s\n) succeeded unexpectedly, want error", v.cmdOutput)
+			}
+		})
+	}
+}
diff --git a/test/packetimpact/runner/packetimpact_test.go b/test/packetimpact/runner/packetimpact_test.go
index 397ca3ba5..1a99719c2 100644
--- a/test/packetimpact/runner/packetimpact_test.go
+++ b/test/packetimpact/runner/packetimpact_test.go
@@ -304,6 +304,7 @@ func TestOne(t *testing.T) {
 		"--local_ipv4", addressInSubnet(testbenchAddr, *testNet.Subnet).String(),
 		"--remote_ipv6", remoteIPv6.String(),
 		"--remote_mac", remoteMAC.String(),
+		"--remote_interface_id", fmt.Sprintf("%d", dutDeviceInfo.ID),
 		"--device", testNetDev,
 		"--dut_type", *dutPlatform,
 	)
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 8b4a4d905..5d9cec73e 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -43,14 +43,16 @@ func portFromSockaddr(sa unix.Sockaddr) (uint16, error) {
 
 // pickPort makes a new socket and returns the socket FD and port. The domain should be AF_INET or AF_INET6. The caller must close the FD when done with
 // the port if there is no error.
-func pickPort(domain, typ int) (int, uint16, error) {
-	fd, err := unix.Socket(domain, typ, 0)
+func pickPort(domain, typ int) (fd int, port uint16, err error) {
+	fd, err = unix.Socket(domain, typ, 0)
 	if err != nil {
-		return -1, 0, err
+		return -1, 0, fmt.Errorf("creating socket: %w", err)
 	}
 	defer func() {
 		if err != nil {
-			err = multierr.Append(err, unix.Close(fd))
+			if cerr := unix.Close(fd); cerr != nil {
+				err = multierr.Append(err, fmt.Errorf("failed to close socket %d: %w", fd, cerr))
+			}
 		}
 	}()
 	var sa unix.Sockaddr
@@ -60,22 +62,22 @@ func pickPort(domain, typ int) (int, uint16, error) {
 		copy(sa4.Addr[:], net.ParseIP(LocalIPv4).To4())
 		sa = &sa4
 	case unix.AF_INET6:
-		var sa6 unix.SockaddrInet6
+		sa6 := unix.SockaddrInet6{ZoneId: uint32(LocalInterfaceID)}
 		copy(sa6.Addr[:], net.ParseIP(LocalIPv6).To16())
 		sa = &sa6
 	default:
 		return -1, 0, fmt.Errorf("invalid domain %d, it should be one of unix.AF_INET or unix.AF_INET6", domain)
 	}
 	if err = unix.Bind(fd, sa); err != nil {
-		return -1, 0, err
+		return -1, 0, fmt.Errorf("binding to %+v: %w", sa, err)
 	}
 	sa, err = unix.Getsockname(fd)
 	if err != nil {
-		return -1, 0, err
+		return -1, 0, fmt.Errorf("Getsocketname(%d): %w", fd, err)
 	}
-	port, err := portFromSockaddr(sa)
+	port, err = portFromSockaddr(sa)
 	if err != nil {
-		return -1, 0, err
+		return -1, 0, fmt.Errorf("extracting port from socket address %+v: %w", sa, err)
 	}
 	return fd, port, nil
 }
@@ -378,7 +380,7 @@ var _ layerState = (*udpState)(nil)
 func newUDPState(domain int, out, in UDP) (*udpState, error) {
 	portPickerFD, localPort, err := pickPort(domain, unix.SOCK_DGRAM)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("picking port: %w", err)
 	}
 	s := udpState{
 		out:          UDP{SrcPort: &localPort},
@@ -916,14 +918,14 @@ func (conn *UDPIPv4) SendIP(ip IPv4, udp UDP, additionalLayers ...Layer) {
 func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
 	conn.t.Helper()
 	layer, err := (*Connection)(conn).Expect(&udp, timeout)
-	if layer == nil {
+	if err != nil {
 		return nil, err
 	}
 	gotUDP, ok := layer.(*UDP)
 	if !ok {
 		conn.t.Fatalf("expected %s to be UDP", layer)
 	}
-	return gotUDP, err
+	return gotUDP, nil
 }
 
 // ExpectData is a convenient method that expects a Layer and the Layer after
@@ -948,3 +950,114 @@ func (conn *UDPIPv4) Close() {
 func (conn *UDPIPv4) Drain() {
 	conn.sniffer.Drain()
 }
+
+// UDPIPv6 maintains the state for all the layers in a UDP/IPv6 connection.
+type UDPIPv6 Connection
+
+// NewUDPIPv6 creates a new UDPIPv6 connection with reasonable defaults.
+func NewUDPIPv6(t *testing.T, outgoingUDP, incomingUDP UDP) UDPIPv6 {
+	etherState, err := newEtherState(Ether{}, Ether{})
+	if err != nil {
+		t.Fatalf("can't make etherState: %s", err)
+	}
+	ipv6State, err := newIPv6State(IPv6{}, IPv6{})
+	if err != nil {
+		t.Fatalf("can't make IPv6State: %s", err)
+	}
+	udpState, err := newUDPState(unix.AF_INET6, outgoingUDP, incomingUDP)
+	if err != nil {
+		t.Fatalf("can't make udpState: %s", err)
+	}
+	injector, err := NewInjector(t)
+	if err != nil {
+		t.Fatalf("can't make injector: %s", err)
+	}
+	sniffer, err := NewSniffer(t)
+	if err != nil {
+		t.Fatalf("can't make sniffer: %s", err)
+	}
+	return UDPIPv6{
+		layerStates: []layerState{etherState, ipv6State, udpState},
+		injector:    injector,
+		sniffer:     sniffer,
+		t:           t,
+	}
+}
+
+func (conn *UDPIPv6) udpState() *udpState {
+	state, ok := conn.layerStates[2].(*udpState)
+	if !ok {
+		conn.t.Fatalf("got transport-layer state type=%T, expected udpState", conn.layerStates[2])
+	}
+	return state
+}
+
+func (conn *UDPIPv6) ipv6State() *ipv6State {
+	state, ok := conn.layerStates[1].(*ipv6State)
+	if !ok {
+		conn.t.Fatalf("got network-layer state type=%T, expected ipv6State", conn.layerStates[1])
+	}
+	return state
+}
+
+// LocalAddr gets the local socket address of this connection.
+func (conn *UDPIPv6) LocalAddr() *unix.SockaddrInet6 {
+	sa := &unix.SockaddrInet6{
+		Port: int(*conn.udpState().out.SrcPort),
+		// Local address is in perspective to the remote host, so it's scoped to the
+		// ID of the remote interface.
+		ZoneId: uint32(RemoteInterfaceID),
+	}
+	copy(sa.Addr[:], *conn.ipv6State().out.SrcAddr)
+	return sa
+}
+
+// Send sends a packet with reasonable defaults, potentially overriding the UDP
+// layer and adding additionLayers.
+func (conn *UDPIPv6) Send(udp UDP, additionalLayers ...Layer) {
+	(*Connection)(conn).send(Layers{&udp}, additionalLayers...)
+}
+
+// SendIPv6 sends a packet with reasonable defaults, potentially overriding the
+// UDP and IPv6 headers and adding additionLayers.
+func (conn *UDPIPv6) SendIPv6(ip IPv6, udp UDP, additionalLayers ...Layer) {
+	(*Connection)(conn).send(Layers{&ip, &udp}, additionalLayers...)
+}
+
+// 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 *UDPIPv6) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
+	conn.t.Helper()
+	layer, err := (*Connection)(conn).Expect(&udp, timeout)
+	if err != nil {
+		return nil, err
+	}
+	gotUDP, ok := layer.(*UDP)
+	if !ok {
+		conn.t.Fatalf("expected %s to be UDP", layer)
+	}
+	return gotUDP, nil
+}
+
+// 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 *UDPIPv6) 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 UDPIPv6 connection.
+func (conn *UDPIPv6) Close() {
+	(*Connection)(conn).Close()
+}
+
+// Drain drains the sniffer's receive buffer by receiving packets until there's
+// nothing else to receive.
+func (conn *UDPIPv6) Drain() {
+	conn.sniffer.Drain()
+}
diff --git a/test/packetimpact/testbench/dut.go b/test/packetimpact/testbench/dut.go
index 2a2afecb5..51be13759 100644
--- a/test/packetimpact/testbench/dut.go
+++ b/test/packetimpact/testbench/dut.go
@@ -87,7 +87,7 @@ func (dut *DUT) sockaddrToProto(sa unix.Sockaddr) *pb.Sockaddr {
 			},
 		}
 	}
-	dut.t.Fatalf("can't parse Sockaddr: %+v", sa)
+	dut.t.Fatalf("can't parse Sockaddr struct: %+v", sa)
 	return nil
 }
 
@@ -106,8 +106,9 @@ func (dut *DUT) protoToSockaddr(sa *pb.Sockaddr) unix.Sockaddr {
 			ZoneId: s.In6.GetScopeId(),
 		}
 		copy(ret.Addr[:], s.In6.GetAddr())
+		return &ret
 	}
-	dut.t.Fatalf("can't parse Sockaddr: %+v", sa)
+	dut.t.Fatalf("can't parse Sockaddr proto: %+v", sa)
 	return nil
 }
 
@@ -126,6 +127,7 @@ func (dut *DUT) CreateBoundSocket(typ, proto int32, addr net.IP) (int32, uint16)
 		fd = dut.Socket(unix.AF_INET6, typ, proto)
 		sa := unix.SockaddrInet6{}
 		copy(sa.Addr[:], addr.To16())
+		sa.ZoneId = uint32(RemoteInterfaceID)
 		dut.Bind(fd, &sa)
 	} else {
 		dut.t.Fatalf("unknown ip addr type for remoteIP")
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index a8121b0da..95bafa876 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -904,12 +904,14 @@ func payload(l Layer) (buffer.VectorisedView, error) {
 func layerChecksum(l Layer, protoNumber tcpip.TransportProtocolNumber) (uint16, error) {
 	totalLength := uint16(totalLength(l))
 	var xsum uint16
-	switch s := l.Prev().(type) {
+	switch p := l.Prev().(type) {
 	case *IPv4:
-		xsum = header.PseudoHeaderChecksum(protoNumber, *s.SrcAddr, *s.DstAddr, totalLength)
+		xsum = header.PseudoHeaderChecksum(protoNumber, *p.SrcAddr, *p.DstAddr, totalLength)
+	case *IPv6:
+		xsum = header.PseudoHeaderChecksum(protoNumber, *p.SrcAddr, *p.DstAddr, totalLength)
 	default:
-		// TODO(b/150301488): Support more protocols, like IPv6.
-		return 0, fmt.Errorf("can't get src and dst addr from previous layer: %#v", s)
+		// TODO(b/161246171): Support more protocols.
+		return 0, fmt.Errorf("checksum for protocol %d is not supported when previous layer is %T", protoNumber, p)
 	}
 	payloadBytes, err := payload(l)
 	if err != nil {
diff --git a/test/packetimpact/testbench/testbench.go b/test/packetimpact/testbench/testbench.go
index 6530036a8..242464e3a 100644
--- a/test/packetimpact/testbench/testbench.go
+++ b/test/packetimpact/testbench/testbench.go
@@ -44,6 +44,13 @@ var (
 	// RemoteIPv6 is the DUT's IPv6 address on the test network.
 	RemoteIPv6 = ""
 
+	// LocalInterfaceID is the ID of the local interface on the test network.
+	LocalInterfaceID uint32
+	// RemoteInterfaceID is the ID of the remote interface on the test network.
+	//
+	// Not using uint32 because package flag does not support uint32.
+	RemoteInterfaceID uint64
+
 	// LocalMAC is the local MAC address on the test network.
 	LocalMAC = ""
 	// RemoteMAC is the DUT's MAC address on the test network.
@@ -75,6 +82,7 @@ func RegisterFlags(fs *flag.FlagSet) {
 	fs.StringVar(&RemoteMAC, "remote_mac", RemoteMAC, "remote mac address for test packets")
 	fs.StringVar(&Device, "device", Device, "local device for test packets")
 	fs.StringVar(&DUTType, "dut_type", DUTType, "type of device under test")
+	fs.Uint64Var(&RemoteInterfaceID, "remote_interface_id", RemoteInterfaceID, "remote interface ID for test packets")
 }
 
 // genPseudoFlags populates flag-like global config based on real flags.
@@ -97,6 +105,7 @@ func genPseudoFlags() error {
 
 	LocalMAC = deviceInfo.MAC.String()
 	LocalIPv6 = deviceInfo.IPv6Addr.String()
+	LocalInterfaceID = deviceInfo.ID
 
 	if deviceInfo.IPv4Net != nil {
 		IPv4PrefixLength, _ = deviceInfo.IPv4Net.Mask.Size()
diff --git a/test/packetimpact/tests/udp_send_recv_dgram_test.go b/test/packetimpact/tests/udp_send_recv_dgram_test.go
index 224feef85..bd53ad90b 100644
--- a/test/packetimpact/tests/udp_send_recv_dgram_test.go
+++ b/test/packetimpact/tests/udp_send_recv_dgram_test.go
@@ -28,62 +28,74 @@ func init() {
 	testbench.RegisterFlags(flag.CommandLine)
 }
 
-func TestUDPRecv(t *testing.T) {
+type udpConn interface {
+	Send(testbench.UDP, ...testbench.Layer)
+	ExpectData(testbench.UDP, testbench.Payload, time.Duration) (testbench.Layers, error)
+	Drain()
+	Close()
+}
+
+func TestUDP(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"))
-	defer dut.Close(boundFD)
-	conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
-	defer conn.Close()
 
-	testCases := []struct {
-		name    string
-		payload []byte
-	}{
-		{"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: 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)
+	for _, isIPv4 := range []bool{true, false} {
+		ipVersionName := "IPv6"
+		if isIPv4 {
+			ipVersionName = "IPv4"
+		}
+		t.Run(ipVersionName, func(t *testing.T) {
+			var addr string
+			if isIPv4 {
+				addr = testbench.RemoteIPv4
+			} else {
+				addr = testbench.RemoteIPv6
 			}
-		})
-	}
-}
+			boundFD, remotePort := dut.CreateBoundSocket(unix.SOCK_DGRAM, unix.IPPROTO_UDP, net.ParseIP(addr))
+			defer dut.Close(boundFD)
 
-func TestUDPSend(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"))
-	defer dut.Close(boundFD)
-	conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
-	defer conn.Close()
+			var conn udpConn
+			var localAddr unix.Sockaddr
+			if isIPv4 {
+				v4Conn := testbench.NewUDPIPv4(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
+				localAddr = v4Conn.LocalAddr()
+				conn = &v4Conn
+			} else {
+				v6Conn := testbench.NewUDPIPv6(t, testbench.UDP{DstPort: &remotePort}, testbench.UDP{SrcPort: &remotePort})
+				localAddr = v6Conn.LocalAddr()
+				conn = &v6Conn
+			}
+			defer conn.Close()
 
-	testCases := []struct {
-		name    string
-		payload []byte
-	}{
-		{"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.Drain()
-			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)
+			testCases := []struct {
+				name    string
+				payload []byte
+			}{
+				{"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.
 			}
-			if _, err := conn.ExpectData(testbench.UDP{SrcPort: &remotePort}, testbench.Payload{Bytes: tc.payload}, 1*time.Second); err != nil {
-				t.Fatal(err)
+			for _, tc := range testCases {
+				t.Run(tc.name, func(t *testing.T) {
+					t.Run("Send", func(t *testing.T) {
+						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)
+						}
+					})
+					t.Run("Recv", func(t *testing.T) {
+						conn.Drain()
+						if got, want := int(dut.SendTo(boundFD, tc.payload, 0, 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: tc.payload}, time.Second); err != nil {
+							t.Fatal(err)
+						}
+					})
+				})
 			}
 		})
 	}