Make fragmentation return a reassembled PacketBuffer

This allows later decoupling of the backing network buffer implementation.

PiperOrigin-RevId: 354643297

3 files changed, 26 insertions, 0 deletions

diff --git a/pkg/tcpip/stack/BUILD b/pkg/tcpip/stack/BUILD
index bb30556cf..ee23c9b98 100644
--- a/pkg/tcpip/stack/BUILD
+++ b/pkg/tcpip/stack/BUILD
@@ -72,6 +72,7 @@ go_library(
         "nud.go",
         "packet_buffer.go",
         "packet_buffer_list.go",
+        "packet_buffer_unsafe.go",
         "pending_packets.go",
         "rand.go",
         "registration.go",
diff --git a/pkg/tcpip/stack/packet_buffer.go b/pkg/tcpip/stack/packet_buffer.go
index 9d4fc3e48..4f013b212 100644
--- a/pkg/tcpip/stack/packet_buffer.go
+++ b/pkg/tcpip/stack/packet_buffer.go
@@ -187,6 +187,12 @@ func (pk *PacketBuffer) Size() int {
 	return pk.HeaderSize() + pk.Data.Size()
 }
 
+// MemSize returns the estimation size of the pk in memory, including backing
+// buffer data.
+func (pk *PacketBuffer) MemSize() int {
+	return pk.HeaderSize() + pk.Data.MemSize() + packetBufferStructSize
+}
+
 // Views returns the underlying storage of the whole packet.
 func (pk *PacketBuffer) Views() []buffer.View {
 	// Optimization for outbound packets that headers are in pk.header.
diff --git a/pkg/tcpip/stack/packet_buffer_unsafe.go b/pkg/tcpip/stack/packet_buffer_unsafe.go
new file mode 100644
index 000000000..ee3d47270
--- /dev/null
+++ b/pkg/tcpip/stack/packet_buffer_unsafe.go
@@ -0,0 +1,19 @@
+// Copyright 2021 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 stack
+
+import "unsafe"
+
+const packetBufferStructSize = int(unsafe.Sizeof(PacketBuffer{}))