From 825c185dc56251bd330124ef773c6653e3887579 Mon Sep 17 00:00:00 2001
From: Ting-Yu Wang <anivia@google.com>
Date: Fri, 29 Jan 2021 17:35:27 -0800
Subject: Make fragmentation return a reassembled PacketBuffer

This allows later decoupling of the backing network buffer implementation.

PiperOrigin-RevId: 354643297
---
 pkg/tcpip/stack/BUILD                   |  1 +
 pkg/tcpip/stack/packet_buffer.go        |  6 ++++++
 pkg/tcpip/stack/packet_buffer_unsafe.go | 19 +++++++++++++++++++
 3 files changed, 26 insertions(+)
 create mode 100644 pkg/tcpip/stack/packet_buffer_unsafe.go

(limited to 'pkg/tcpip/stack')

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{}))
-- 
cgit v1.2.3