3 files changed, 88 insertions, 9 deletions
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index a90046f69..8fa585804 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -839,6 +839,61 @@ func (conn *TCPIPv4) Drain(t *testing.T) {
 	conn.sniffer.Drain(t)
 }
 
+// IPv4Conn maintains the state for all the layers in a IPv4 connection.
+type IPv4Conn Connection
+
+// NewIPv4Conn creates a new IPv4Conn connection with reasonable defaults.
+func NewIPv4Conn(t *testing.T, outgoingIPv4, incomingIPv4 IPv4) IPv4Conn {
+	t.Helper()
+
+	etherState, err := newEtherState(Ether{}, Ether{})
+	if err != nil {
+		t.Fatalf("can't make EtherState: %s", err)
+	}
+	ipv4State, err := newIPv4State(outgoingIPv4, incomingIPv4)
+	if err != nil {
+		t.Fatalf("can't make IPv4State: %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 IPv4Conn{
+		layerStates: []layerState{etherState, ipv4State},
+		injector:    injector,
+		sniffer:     sniffer,
+	}
+}
+
+// Send sends a frame with ipv4 overriding the IPv4 layer defaults and
+// additionalLayers added after it.
+func (c *IPv4Conn) Send(t *testing.T, ipv4 IPv4, additionalLayers ...Layer) {
+	t.Helper()
+
+	(*Connection)(c).send(t, Layers{&ipv4}, additionalLayers...)
+}
+
+// Close cleans up any resources held.
+func (c *IPv4Conn) Close(t *testing.T) {
+	t.Helper()
+
+	(*Connection)(c).Close(t)
+}
+
+// ExpectFrame expects a frame that matches the provided Layers within the
+// timeout specified. If it doesn't arrive in time, an error is returned.
+func (c *IPv4Conn) ExpectFrame(t *testing.T, frame Layers, timeout time.Duration) (Layers, error) {
+	t.Helper()
+
+	return (*Connection)(c).ExpectFrame(t, frame, timeout)
+}
+
 // IPv6Conn maintains the state for all the layers in a IPv6 connection.
 type IPv6Conn Connection
 
diff --git a/test/packetimpact/testbench/dut_client.go b/test/packetimpact/testbench/dut_client.go
index d0e68c5da..0fc3d97b4 100644
--- a/test/packetimpact/testbench/dut_client.go
+++ b/test/packetimpact/testbench/dut_client.go
@@ -19,7 +19,7 @@ import (
 	pb "gvisor.dev/gvisor/test/packetimpact/proto/posix_server_go_proto"
 )
 
-// PosixClient is a gRPC client for the Posix service.
+// POSIXClient is a gRPC client for the Posix service.
 type POSIXClient pb.PosixClient
 
 // NewPOSIXClient makes a new gRPC client for the POSIX service.
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index a35562ca8..af7a2ba4e 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -879,6 +879,9 @@ type ICMPv4 struct {
 	Type     *header.ICMPv4Type
 	Code     *header.ICMPv4Code
 	Checksum *uint16
+	Ident    *uint16
+	Sequence *uint16
+	Payload  []byte
 }
 
 func (l *ICMPv4) String() string {
@@ -887,7 +890,7 @@ func (l *ICMPv4) String() string {
 
 // ToBytes implements Layer.ToBytes.
 func (l *ICMPv4) ToBytes() ([]byte, error) {
-	b := make([]byte, header.ICMPv4MinimumSize)
+	b := make([]byte, header.ICMPv4MinimumSize+len(l.Payload))
 	h := header.ICMPv4(b)
 	if l.Type != nil {
 		h.SetType(*l.Type)
@@ -895,15 +898,33 @@ func (l *ICMPv4) ToBytes() ([]byte, error) {
 	if l.Code != nil {
 		h.SetCode(*l.Code)
 	}
+	if copied := copy(h.Payload(), l.Payload); copied != len(l.Payload) {
+		panic(fmt.Sprintf("wrong number of bytes copied into h.Payload(): got = %d, want = %d", len(h.Payload()), len(l.Payload)))
+	}
+	if l.Ident != nil {
+		h.SetIdent(*l.Ident)
+	}
+	if l.Sequence != nil {
+		h.SetSequence(*l.Sequence)
+	}
+
+	// The checksum must be handled last because the ICMPv4 header fields are
+	// included in the computation.
 	if l.Checksum != nil {
 		h.SetChecksum(*l.Checksum)
-		return h, nil
-	}
-	payload, err := payload(l)
-	if err != nil {
-		return nil, err
+	} else {
+		// Compute the checksum based on the ICMPv4.Payload and also the subsequent
+		// layers.
+		payload, err := payload(l)
+		if err != nil {
+			return nil, err
+		}
+		var vv buffer.VectorisedView
+		vv.AppendView(buffer.View(l.Payload))
+		vv.Append(payload)
+		h.SetChecksum(header.ICMPv4Checksum(h, vv))
 	}
-	h.SetChecksum(header.ICMPv4Checksum(h, payload))
+
 	return h, nil
 }
 
@@ -915,8 +936,11 @@ func parseICMPv4(b []byte) (Layer, layerParser) {
 		Type:     ICMPv4Type(h.Type()),
 		Code:     ICMPv4Code(h.Code()),
 		Checksum: Uint16(h.Checksum()),
+		Ident:    Uint16(h.Ident()),
+		Sequence: Uint16(h.Sequence()),
+		Payload:  h.Payload(),
 	}
-	return &icmpv4, parsePayload
+	return &icmpv4, nil
 }
 
 func (l *ICMPv4) match(other Layer) bool {