Add IPv4 reassembly packetimpact test

The IPv6 reassembly test was also refactored to be easily extended with
more cases.

PiperOrigin-RevId: 339768605

3 files changed, 265 insertions, 127 deletions

diff --git a/test/packetimpact/tests/BUILD b/test/packetimpact/tests/BUILD
index 8c2de5a9f..c30c77a17 100644
--- a/test/packetimpact/tests/BUILD
+++ b/test/packetimpact/tests/BUILD
@@ -298,6 +298,18 @@ packetimpact_testbench(
 )
 
 packetimpact_testbench(
+    name = "ipv4_fragment_reassembly",
+    srcs = ["ipv4_fragment_reassembly_test.go"],
+    deps = [
+        "//pkg/tcpip/buffer",
+        "//pkg/tcpip/header",
+        "//test/packetimpact/testbench",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+packetimpact_testbench(
     name = "ipv6_fragment_reassembly",
     srcs = ["ipv6_fragment_reassembly_test.go"],
     deps = [
@@ -305,6 +317,7 @@ packetimpact_testbench(
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
         "//test/packetimpact/testbench",
+        "@com_github_google_go_cmp//cmp:go_default_library",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
diff --git a/test/packetimpact/tests/ipv4_fragment_reassembly_test.go b/test/packetimpact/tests/ipv4_fragment_reassembly_test.go
new file mode 100644
index 000000000..65c0df140
--- /dev/null
+++ b/test/packetimpact/tests/ipv4_fragment_reassembly_test.go
@@ -0,0 +1,142 @@
+// 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 ipv4_fragment_reassembly_test
+
+import (
+	"flag"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"gvisor.dev/gvisor/pkg/tcpip/buffer"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/test/packetimpact/testbench"
+)
+
+func init() {
+	testbench.RegisterFlags(flag.CommandLine)
+}
+
+type fragmentInfo struct {
+	offset uint16
+	size   uint16
+	more   uint8
+}
+
+func TestIPv4FragmentReassembly(t *testing.T) {
+	const fragmentID = 42
+	icmpv4ProtoNum := uint8(header.ICMPv4ProtocolNumber)
+
+	tests := []struct {
+		description  string
+		ipPayloadLen int
+		fragments    []fragmentInfo
+		expectReply  bool
+	}{
+		{
+			description:  "basic reassembly",
+			ipPayloadLen: 2000,
+			fragments: []fragmentInfo{
+				{offset: 0, size: 1000, more: header.IPv4FlagMoreFragments},
+				{offset: 1000, size: 1000, more: 0},
+			},
+			expectReply: true,
+		},
+		{
+			description:  "out of order fragments",
+			ipPayloadLen: 2000,
+			fragments: []fragmentInfo{
+				{offset: 1000, size: 1000, more: 0},
+				{offset: 0, size: 1000, more: header.IPv4FlagMoreFragments},
+			},
+			expectReply: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			dut := testbench.NewDUT(t)
+			defer dut.TearDown()
+			conn := testbench.NewIPv4Conn(t, testbench.IPv4{}, testbench.IPv4{})
+			defer conn.Close(t)
+
+			data := make([]byte, test.ipPayloadLen)
+			icmp := header.ICMPv4(data[:header.ICMPv4MinimumSize])
+			icmp.SetType(header.ICMPv4Echo)
+			icmp.SetCode(header.ICMPv4UnusedCode)
+			icmp.SetChecksum(0)
+			icmp.SetSequence(0)
+			icmp.SetIdent(0)
+			originalPayload := data[header.ICMPv4MinimumSize:]
+			if _, err := rand.Read(originalPayload); err != nil {
+				t.Fatalf("rand.Read: %s", err)
+			}
+			cksum := header.ICMPv4Checksum(
+				icmp,
+				buffer.NewVectorisedView(len(originalPayload), []buffer.View{originalPayload}),
+			)
+			icmp.SetChecksum(cksum)
+
+			for _, fragment := range test.fragments {
+				conn.Send(t,
+					testbench.IPv4{
+						Protocol:       &icmpv4ProtoNum,
+						FragmentOffset: testbench.Uint16(fragment.offset),
+						Flags:          testbench.Uint8(fragment.more),
+						ID:             testbench.Uint16(fragmentID),
+					},
+					&testbench.Payload{
+						Bytes: data[fragment.offset:][:fragment.size],
+					})
+			}
+
+			var bytesReceived int
+			reassembledPayload := make([]byte, test.ipPayloadLen)
+			for {
+				incomingFrame, err := conn.ExpectFrame(t, testbench.Layers{
+					&testbench.Ether{},
+					&testbench.IPv4{},
+					&testbench.ICMPv4{},
+				}, time.Second)
+				if err != nil {
+					// Either an unexpected frame was received, or none at all.
+					if bytesReceived < test.ipPayloadLen {
+						t.Fatalf("received %d bytes out of %d, then conn.ExpectFrame(_, _, time.Second) failed with %s", bytesReceived, test.ipPayloadLen, err)
+					}
+					break
+				}
+				if !test.expectReply {
+					t.Fatalf("unexpected reply received:\n%s", incomingFrame)
+				}
+				ipPayload, err := incomingFrame[2 /* ICMPv4 */].ToBytes()
+				if err != nil {
+					t.Fatalf("failed to parse ICMPv4 header: incomingPacket[2].ToBytes() = (_, %s)", err)
+				}
+				offset := *incomingFrame[1 /* IPv4 */].(*testbench.IPv4).FragmentOffset
+				if copied := copy(reassembledPayload[offset:], ipPayload); copied != len(ipPayload) {
+					t.Fatalf("wrong number of bytes copied into reassembledPayload: got = %d, want = %d", copied, len(ipPayload))
+				}
+				bytesReceived += len(ipPayload)
+			}
+
+			if test.expectReply {
+				if diff := cmp.Diff(originalPayload, reassembledPayload[header.ICMPv4MinimumSize:]); diff != "" {
+					t.Fatalf("reassembledPayload mismatch (-want +got):\n%s", diff)
+				}
+			}
+		})
+	}
+}
diff --git a/test/packetimpact/tests/ipv6_fragment_reassembly_test.go b/test/packetimpact/tests/ipv6_fragment_reassembly_test.go
index a24c85566..4a29de688 100644
--- a/test/packetimpact/tests/ipv6_fragment_reassembly_test.go
+++ b/test/packetimpact/tests/ipv6_fragment_reassembly_test.go
@@ -15,154 +15,137 @@
 package ipv6_fragment_reassembly_test
 
 import (
-	"bytes"
-	"encoding/binary"
-	"encoding/hex"
 	"flag"
+	"math/rand"
 	"net"
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/test/packetimpact/testbench"
 )
 
-const (
-	// The payload length for the first fragment we send. This number
-	// is a multiple of 8 near 750 (half of 1500).
-	firstPayloadLength = 752
-	// The ID field for our outgoing fragments.
-	fragmentID = 1
-	// A node must be able to accept a fragmented packet that,
-	// after reassembly, is as large as 1500 octets.
-	reassemblyCap = 1500
-)
-
 func init() {
 	testbench.RegisterFlags(flag.CommandLine)
 }
 
-func TestIPv6FragmentReassembly(t *testing.T) {
-	dut := testbench.NewDUT(t)
-	defer dut.TearDown()
-	conn := testbench.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
-	defer conn.Close(t)
-
-	firstPayloadToSend := make([]byte, firstPayloadLength)
-	for i := range firstPayloadToSend {
-		firstPayloadToSend[i] = 'A'
-	}
-
-	secondPayloadLength := reassemblyCap - firstPayloadLength - header.ICMPv6EchoMinimumSize
-	secondPayloadToSend := firstPayloadToSend[:secondPayloadLength]
-
-	icmpv6EchoPayload := make([]byte, 4)
-	binary.BigEndian.PutUint16(icmpv6EchoPayload[0:], 0)
-	binary.BigEndian.PutUint16(icmpv6EchoPayload[2:], 0)
-	icmpv6EchoPayload = append(icmpv6EchoPayload, firstPayloadToSend...)
-
-	lIP := tcpip.Address(net.ParseIP(testbench.LocalIPv6).To16())
-	rIP := tcpip.Address(net.ParseIP(testbench.RemoteIPv6).To16())
-	icmpv6 := testbench.ICMPv6{
-		Type:    testbench.ICMPv6Type(header.ICMPv6EchoRequest),
-		Code:    testbench.ICMPv6Code(header.ICMPv6UnusedCode),
-		Payload: icmpv6EchoPayload,
-	}
-	icmpv6Bytes, err := icmpv6.ToBytes()
-	if err != nil {
-		t.Fatalf("failed to serialize ICMPv6: %s", err)
-	}
-	cksum := header.ICMPv6Checksum(
-		header.ICMPv6(icmpv6Bytes),
-		lIP,
-		rIP,
-		buffer.NewVectorisedView(len(secondPayloadToSend), []buffer.View{secondPayloadToSend}),
-	)
-
-	conn.Send(t, testbench.IPv6{},
-		&testbench.IPv6FragmentExtHdr{
-			FragmentOffset: testbench.Uint16(0),
-			MoreFragments:  testbench.Bool(true),
-			Identification: testbench.Uint32(fragmentID),
-		},
-		&testbench.ICMPv6{
-			Type:     testbench.ICMPv6Type(header.ICMPv6EchoRequest),
-			Code:     testbench.ICMPv6Code(header.ICMPv6UnusedCode),
-			Payload:  icmpv6EchoPayload,
-			Checksum: &cksum,
-		})
+type fragmentInfo struct {
+	offset uint16
+	size   uint16
+	more   bool
+}
 
+func TestIPv6FragmentReassembly(t *testing.T) {
+	const fragmentID = 42
 	icmpv6ProtoNum := header.IPv6ExtensionHeaderIdentifier(header.ICMPv6ProtocolNumber)
 
-	conn.Send(t, testbench.IPv6{},
-		&testbench.IPv6FragmentExtHdr{
-			NextHeader:     &icmpv6ProtoNum,
-			FragmentOffset: testbench.Uint16((firstPayloadLength + header.ICMPv6EchoMinimumSize) / 8),
-			MoreFragments:  testbench.Bool(false),
-			Identification: testbench.Uint32(fragmentID),
+	tests := []struct {
+		description  string
+		ipPayloadLen int
+		fragments    []fragmentInfo
+		expectReply  bool
+	}{
+		{
+			description:  "basic reassembly",
+			ipPayloadLen: 1500,
+			fragments: []fragmentInfo{
+				{offset: 0, size: 760, more: true},
+				{offset: 760, size: 740, more: false},
+			},
+			expectReply: true,
 		},
-		&testbench.Payload{
-			Bytes: secondPayloadToSend,
-		})
-
-	gotEchoReplyFirstPart, err := conn.ExpectFrame(t, testbench.Layers{
-		&testbench.Ether{},
-		&testbench.IPv6{},
-		&testbench.IPv6FragmentExtHdr{
-			FragmentOffset: testbench.Uint16(0),
-			MoreFragments:  testbench.Bool(true),
+		{
+			description:  "out of order fragments",
+			ipPayloadLen: 3000,
+			fragments: []fragmentInfo{
+				{offset: 0, size: 1024, more: true},
+				{offset: 2048, size: 952, more: false},
+				{offset: 1024, size: 1024, more: true},
+			},
+			expectReply: true,
 		},
-		&testbench.ICMPv6{
-			Type: testbench.ICMPv6Type(header.ICMPv6EchoReply),
-			Code: testbench.ICMPv6Code(header.ICMPv6UnusedCode),
-		},
-	}, time.Second)
-	if err != nil {
-		t.Fatalf("expected a fragmented ICMPv6 Echo Reply, but got none: %s", err)
 	}
 
-	id := *gotEchoReplyFirstPart[2].(*testbench.IPv6FragmentExtHdr).Identification
-	gotFirstPayload, err := gotEchoReplyFirstPart[len(gotEchoReplyFirstPart)-1].ToBytes()
-	if err != nil {
-		t.Fatalf("failed to serialize ICMPv6: %s", err)
-	}
-	icmpPayload := gotFirstPayload[header.ICMPv6EchoMinimumSize:]
-	receivedLen := len(icmpPayload)
-	wantSecondPayloadLen := reassemblyCap - header.ICMPv6EchoMinimumSize - receivedLen
-	wantFirstPayload := make([]byte, receivedLen)
-	for i := range wantFirstPayload {
-		wantFirstPayload[i] = 'A'
-	}
-	wantSecondPayload := wantFirstPayload[:wantSecondPayloadLen]
-	if !bytes.Equal(icmpPayload, wantFirstPayload) {
-		t.Fatalf("received unexpected payload, got: %s, want: %s",
-			hex.Dump(icmpPayload),
-			hex.Dump(wantFirstPayload))
-	}
-
-	gotEchoReplySecondPart, err := conn.ExpectFrame(t, testbench.Layers{
-		&testbench.Ether{},
-		&testbench.IPv6{},
-		&testbench.IPv6FragmentExtHdr{
-			NextHeader:     &icmpv6ProtoNum,
-			FragmentOffset: testbench.Uint16(uint16((receivedLen + header.ICMPv6EchoMinimumSize) / 8)),
-			MoreFragments:  testbench.Bool(false),
-			Identification: &id,
-		},
-		&testbench.ICMPv6{},
-	}, time.Second)
-	if err != nil {
-		t.Fatalf("expected the rest of ICMPv6 Echo Reply, but got none: %s", err)
-	}
-	secondPayload, err := gotEchoReplySecondPart[len(gotEchoReplySecondPart)-1].ToBytes()
-	if err != nil {
-		t.Fatalf("failed to serialize ICMPv6 Echo Reply: %s", err)
-	}
-	if !bytes.Equal(secondPayload, wantSecondPayload) {
-		t.Fatalf("received unexpected payload, got: %s, want: %s",
-			hex.Dump(secondPayload),
-			hex.Dump(wantSecondPayload))
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			dut := testbench.NewDUT(t)
+			defer dut.TearDown()
+			conn := testbench.NewIPv6Conn(t, testbench.IPv6{}, testbench.IPv6{})
+			defer conn.Close(t)
+
+			lIP := tcpip.Address(net.ParseIP(testbench.LocalIPv6).To16())
+			rIP := tcpip.Address(net.ParseIP(testbench.RemoteIPv6).To16())
+
+			data := make([]byte, test.ipPayloadLen)
+			icmp := header.ICMPv6(data[:header.ICMPv6HeaderSize])
+			icmp.SetType(header.ICMPv6EchoRequest)
+			icmp.SetCode(header.ICMPv6UnusedCode)
+			icmp.SetChecksum(0)
+			originalPayload := data[header.ICMPv6HeaderSize:]
+			if _, err := rand.Read(originalPayload); err != nil {
+				t.Fatalf("rand.Read: %s", err)
+			}
+
+			cksum := header.ICMPv6Checksum(
+				icmp,
+				lIP,
+				rIP,
+				buffer.NewVectorisedView(len(originalPayload), []buffer.View{originalPayload}),
+			)
+			icmp.SetChecksum(cksum)
+
+			for _, fragment := range test.fragments {
+				conn.Send(t, testbench.IPv6{},
+					&testbench.IPv6FragmentExtHdr{
+						NextHeader:     &icmpv6ProtoNum,
+						FragmentOffset: testbench.Uint16(fragment.offset / header.IPv6FragmentExtHdrFragmentOffsetBytesPerUnit),
+						MoreFragments:  testbench.Bool(fragment.more),
+						Identification: testbench.Uint32(fragmentID),
+					},
+					&testbench.Payload{
+						Bytes: data[fragment.offset:][:fragment.size],
+					})
+			}
+
+			var bytesReceived int
+			reassembledPayload := make([]byte, test.ipPayloadLen)
+			for {
+				incomingFrame, err := conn.ExpectFrame(t, testbench.Layers{
+					&testbench.Ether{},
+					&testbench.IPv6{},
+					&testbench.IPv6FragmentExtHdr{},
+					&testbench.ICMPv6{},
+				}, time.Second)
+				if err != nil {
+					// Either an unexpected frame was received, or none at all.
+					if bytesReceived < test.ipPayloadLen {
+						t.Fatalf("received %d bytes out of %d, then conn.ExpectFrame(_, _, time.Second) failed with %s", bytesReceived, test.ipPayloadLen, err)
+					}
+					break
+				}
+				if !test.expectReply {
+					t.Fatalf("unexpected reply received:\n%s", incomingFrame)
+				}
+				ipPayload, err := incomingFrame[3 /* ICMPv6 */].ToBytes()
+				if err != nil {
+					t.Fatalf("failed to parse ICMPv6 header: incomingPacket[3].ToBytes() = (_, %s)", err)
+				}
+				offset := *incomingFrame[2 /* IPv6FragmentExtHdr */].(*testbench.IPv6FragmentExtHdr).FragmentOffset
+				offset *= header.IPv6FragmentExtHdrFragmentOffsetBytesPerUnit
+				if copied := copy(reassembledPayload[offset:], ipPayload); copied != len(ipPayload) {
+					t.Fatalf("wrong number of bytes copied into reassembledPayload: got = %d, want = %d", copied, len(ipPayload))
+				}
+				bytesReceived += len(ipPayload)
+			}
+
+			if test.expectReply {
+				if diff := cmp.Diff(originalPayload, reassembledPayload[header.ICMPv6HeaderSize:]); diff != "" {
+					t.Fatalf("reassembledPayload mismatch (-want +got):\n%s", diff)
+				}
+			}
+		})
 	}
 }