From a446b45d4d05ba9216cf79c53000799fb1201cbe Mon Sep 17 00:00:00 2001
From: Lucas Manning <lucasmanning@google.com>
Date: Thu, 3 Oct 2024 12:08:47 -0700
Subject: [PATCH] Ensure views returned by PullUp are owned exclusively by
 their packet.

PiperOrigin-RevId: 681977034
---
 pkg/buffer/buffer.go                | 7 +++++--
 pkg/tcpip/network/ipv4/ipv4.go      | 4 +++-
 pkg/tcpip/network/ipv6/ipv6.go      | 4 ++++
 pkg/tcpip/transport/tcp/segment.go  | 2 +-
 pkg/tcpip/transport/udp/endpoint.go | 4 +++-
 5 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/pkg/buffer/buffer.go b/pkg/buffer/buffer.go
index 3e6bc6dd13..0f352c0dc8 100644
--- a/pkg/buffer/buffer.go
+++ b/pkg/buffer/buffer.go
@@ -318,8 +318,11 @@ func (b *Buffer) PullUp(offset, length int) (View, bool) {
 		if x := curr.Intersect(tgt); x.Len() == tgt.Len() {
 			// buf covers the whole requested target range.
 			sub := x.Offset(-curr.begin)
-			// Don't increment the reference count of the underlying chunk. Views
-			// returned by PullUp are explicitly unowned and read only
+			if v.sharesChunk() {
+				old := v.chunk
+				v.chunk = v.chunk.Clone()
+				old.DecRef()
+			}
 			new := View{
 				read:  v.read + sub.begin,
 				write: v.read + sub.end,
diff --git a/pkg/tcpip/network/ipv4/ipv4.go b/pkg/tcpip/network/ipv4/ipv4.go
index e2721a4db6..9e9067f386 100644
--- a/pkg/tcpip/network/ipv4/ipv4.go
+++ b/pkg/tcpip/network/ipv4/ipv4.go
@@ -868,7 +868,9 @@ func (e *endpoint) HandlePacket(pkt *stack.PacketBuffer) {
 			return
 		}
 	}
-
+	// CheckPrerouting can modify the backing storage of the packet, so refresh
+	// the header.
+	h = header.IPv4(pkt.NetworkHeader().Slice())
 	e.handleValidatedPacket(h, pkt, e.nic.Name() /* inNICName */)
 }
 
diff --git a/pkg/tcpip/network/ipv6/ipv6.go b/pkg/tcpip/network/ipv6/ipv6.go
index c61a4e338c..cc86028e1f 100644
--- a/pkg/tcpip/network/ipv6/ipv6.go
+++ b/pkg/tcpip/network/ipv6/ipv6.go
@@ -1140,6 +1140,9 @@ func (e *endpoint) HandlePacket(pkt *stack.PacketBuffer) {
 		}
 	}
 
+	// CheckPrerouting can modify the backing storage of the packet, so refresh
+	// the header.
+	h = header.IPv6(pkt.NetworkHeader().Slice())
 	e.handleValidatedPacket(h, pkt, e.nic.Name() /* inNICName */)
 }
 
@@ -1495,6 +1498,7 @@ func (e *endpoint) processExtensionHeaders(h header.IPv6, pkt *stack.PacketBuffe
 		routerAlert       *header.IPv6RouterAlertOption
 	)
 	for {
+		h := header.IPv6(pkt.NetworkHeader().Slice())
 		if done, err := e.processExtensionHeader(&it, &pkt, h, &routerAlert, &hasFragmentHeader, forwarding); err != nil || done {
 			return err
 		}
diff --git a/pkg/tcpip/transport/tcp/segment.go b/pkg/tcpip/transport/tcp/segment.go
index 6de583daf5..a05a0767d1 100644
--- a/pkg/tcpip/transport/tcp/segment.go
+++ b/pkg/tcpip/transport/tcp/segment.go
@@ -130,7 +130,7 @@ func newIncomingSegment(id stack.TransportEndpointID, clock tcpip.Clock, pkt *st
 	s.window = seqnum.Size(hdr.WindowSize())
 	s.rcvdTime = clock.NowMonotonic()
 	s.dataMemSize = pkt.MemSize()
-	s.pkt = pkt.IncRef()
+	s.pkt = pkt.Clone()
 	s.csumValid = csumValid
 
 	if !s.pkt.RXChecksumValidated {
diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go
index f8e3057970..785ea47b2e 100644
--- a/pkg/tcpip/transport/udp/endpoint.go
+++ b/pkg/tcpip/transport/udp/endpoint.go
@@ -952,7 +952,9 @@ func (e *endpoint) HandlePacket(id stack.TransportEndpointID, pkt *stack.PacketB
 			Addr: id.LocalAddress,
 			Port: hdr.DestinationPort(),
 		},
-		pkt: pkt.IncRef(),
+		// We need to clone the packet because ReadTo modifies the write index of
+		// the underlying buffer. Clone does not copy the data, just the metadata.
+		pkt: pkt.Clone(),
 	}
 	e.rcvList.PushBack(packet)
 	e.rcvBufSize += pkt.Data().Size()