4 files changed, 144 insertions, 18 deletions
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()