netstack: Add a test for mixed Push/Consume

Not really designed to be used this way, but it works and it's been relied
upon. Add a test.

PiperOrigin-RevId: 371802756

1 files changed, 81 insertions, 53 deletions

diff --git a/pkg/tcpip/stack/packet_buffer_test.go b/pkg/tcpip/stack/packet_buffer_test.go
index a794065d1..bd4eb4fed 100644
--- a/pkg/tcpip/stack/packet_buffer_test.go
+++ b/pkg/tcpip/stack/packet_buffer_test.go
@@ -112,23 +112,13 @@ func TestPacketHeaderPush(t *testing.T) {
 			if got, want := pk.Size(), allHdrSize+len(test.data); got != want {
 				t.Errorf("After pk.Size() = %d, want %d", got, want)
 			}
-			checkData(t, pk, test.data)
-			checkViewEqual(t, "After pk.Views()", concatViews(pk.Views()...),
-				concatViews(test.link, test.network, test.transport, test.data))
-			// Check the after values for each header.
-			checkPacketHeader(t, "After pk.LinkHeader", pk.LinkHeader(), test.link)
-			checkPacketHeader(t, "After pk.NetworkHeader", pk.NetworkHeader(), test.network)
-			checkPacketHeader(t, "After pk.TransportHeader", pk.TransportHeader(), test.transport)
-			// Check the after values for PayloadSince.
-			checkViewEqual(t, "After PayloadSince(LinkHeader)",
-				PayloadSince(pk.LinkHeader()),
-				concatViews(test.link, test.network, test.transport, test.data))
-			checkViewEqual(t, "After PayloadSince(NetworkHeader)",
-				PayloadSince(pk.NetworkHeader()),
-				concatViews(test.network, test.transport, test.data))
-			checkViewEqual(t, "After PayloadSince(TransportHeader)",
-				PayloadSince(pk.TransportHeader()),
-				concatViews(test.transport, test.data))
+			// Check the after state.
+			checkPacketContents(t, "After ", pk, packetContents{
+				link:      test.link,
+				network:   test.network,
+				transport: test.transport,
+				data:      test.data,
+			})
 		})
 	}
 }
@@ -199,29 +189,13 @@ func TestPacketHeaderConsume(t *testing.T) {
 			if got, want := pk.Size(), len(test.data); got != want {
 				t.Errorf("After pk.Size() = %d, want %d", got, want)
 			}
-			// After state of pk.
-			var (
-				link      = test.data[:test.link]
-				network   = test.data[test.link:][:test.network]
-				transport = test.data[test.link+test.network:][:test.transport]
-				payload   = test.data[allHdrSize:]
-			)
-			checkData(t, pk, payload)
-			checkViewEqual(t, "After pk.Views()", concatViews(pk.Views()...), test.data)
-			// Check the after values for each header.
-			checkPacketHeader(t, "After pk.LinkHeader", pk.LinkHeader(), link)
-			checkPacketHeader(t, "After pk.NetworkHeader", pk.NetworkHeader(), network)
-			checkPacketHeader(t, "After pk.TransportHeader", pk.TransportHeader(), transport)
-			// Check the after values for PayloadSince.
-			checkViewEqual(t, "After PayloadSince(LinkHeader)",
-				PayloadSince(pk.LinkHeader()),
-				concatViews(link, network, transport, payload))
-			checkViewEqual(t, "After PayloadSince(NetworkHeader)",
-				PayloadSince(pk.NetworkHeader()),
-				concatViews(network, transport, payload))
-			checkViewEqual(t, "After PayloadSince(TransportHeader)",
-				PayloadSince(pk.TransportHeader()),
-				concatViews(transport, payload))
+			// Check the after state of pk.
+			checkPacketContents(t, "After ", pk, packetContents{
+				link:      test.data[:test.link],
+				network:   test.data[test.link:][:test.network],
+				transport: test.data[test.link+test.network:][:test.transport],
+				data:      test.data[allHdrSize:],
+			})
 		})
 	}
 }
@@ -252,6 +226,39 @@ func TestPacketHeaderConsumeDataTooShort(t *testing.T) {
 	})
 }
 
+// This is a very obscure use-case seen in the code that verifies packets
+// before sending them out. It tries to parse the headers to verify.
+// PacketHeader was initially not designed to mix Push() and Consume(), but it
+// works and it's been relied upon. Include a test here.
+func TestPacketHeaderPushConsumeMixed(t *testing.T) {
+	link := makeView(10)
+	network := makeView(20)
+	data := makeView(30)
+
+	initData := append([]byte(nil), network...)
+	initData = append(initData, data...)
+	pk := NewPacketBuffer(PacketBufferOptions{
+		ReserveHeaderBytes: len(link),
+		Data:               buffer.NewViewFromBytes(initData).ToVectorisedView(),
+	})
+
+	// 1. Consume network header
+	gotNetwork, ok := pk.NetworkHeader().Consume(len(network))
+	if !ok {
+		t.Fatalf("pk.NetworkHeader().Consume(%d) = _, false; want _, true", len(network))
+	}
+	checkViewEqual(t, "gotNetwork", gotNetwork, network)
+
+	// 2. Push link header
+	copy(pk.LinkHeader().Push(len(link)), link)
+
+	checkPacketContents(t, "" /* prefix */, pk, packetContents{
+		link:    link,
+		network: network,
+		data:    data,
+	})
+}
+
 func TestPacketHeaderPushCalledAtMostOnce(t *testing.T) {
 	const headerSize = 10
 
@@ -494,6 +501,37 @@ func TestPacketBufferData(t *testing.T) {
 	}
 }
 
+type packetContents struct {
+	link      buffer.View
+	network   buffer.View
+	transport buffer.View
+	data      buffer.View
+}
+
+func checkPacketContents(t *testing.T, prefix string, pk *PacketBuffer, want packetContents) {
+	t.Helper()
+	// Headers.
+	checkPacketHeader(t, prefix+"pk.LinkHeader", pk.LinkHeader(), want.link)
+	checkPacketHeader(t, prefix+"pk.NetworkHeader", pk.NetworkHeader(), want.network)
+	checkPacketHeader(t, prefix+"pk.TransportHeader", pk.TransportHeader(), want.transport)
+	// Data.
+	checkData(t, pk, want.data)
+	// Whole packet.
+	checkViewEqual(t, prefix+"pk.Views()",
+		concatViews(pk.Views()...),
+		concatViews(want.link, want.network, want.transport, want.data))
+	// PayloadSince.
+	checkViewEqual(t, prefix+"PayloadSince(LinkHeader)",
+		PayloadSince(pk.LinkHeader()),
+		concatViews(want.link, want.network, want.transport, want.data))
+	checkViewEqual(t, prefix+"PayloadSince(NetworkHeader)",
+		PayloadSince(pk.NetworkHeader()),
+		concatViews(want.network, want.transport, want.data))
+	checkViewEqual(t, prefix+"PayloadSince(TransportHeader)",
+		PayloadSince(pk.TransportHeader()),
+		concatViews(want.transport, want.data))
+}
+
 func checkInitialPacketBuffer(t *testing.T, pk *PacketBuffer, opts PacketBufferOptions) {
 	t.Helper()
 	reserved := opts.ReserveHeaderBytes
@@ -510,19 +548,9 @@ func checkInitialPacketBuffer(t *testing.T, pk *PacketBuffer, opts PacketBufferO
 	if got, want := pk.Size(), len(data); got != want {
 		t.Errorf("Initial pk.Size() = %d, want %d", got, want)
 	}
-	checkData(t, pk, data)
-	checkViewEqual(t, "Initial pk.Views()", concatViews(pk.Views()...), data)
-	// Check the initial values for each header.
-	checkPacketHeader(t, "Initial pk.LinkHeader", pk.LinkHeader(), nil)
-	checkPacketHeader(t, "Initial pk.NetworkHeader", pk.NetworkHeader(), nil)
-	checkPacketHeader(t, "Initial pk.TransportHeader", pk.TransportHeader(), nil)
-	// Check the initial valies for PayloadSince.
-	checkViewEqual(t, "Initial PayloadSince(LinkHeader)",
-		PayloadSince(pk.LinkHeader()), data)
-	checkViewEqual(t, "Initial PayloadSince(NetworkHeader)",
-		PayloadSince(pk.NetworkHeader()), data)
-	checkViewEqual(t, "Initial PayloadSince(TransportHeader)",
-		PayloadSince(pk.TransportHeader()), data)
+	checkPacketContents(t, "Initial ", pk, packetContents{
+		data: data,
+	})
 }
 
 func checkPacketHeader(t *testing.T, name string, h PacketHeader, want []byte) {