store read size with the buffer; helpers

Author: illotum

buffer/buffer.go

@@ -29,83 +29,104 @@
 
 package buffer
 
-const (
-	MaxMessageSize = (1 << 16) - 1 // largest possible UDP datagram
-)
-
-// Source produces new Buffers.
-type Source interface {
-
-	// Get returns a Buffer of at least the requested size.
-	// Implementations may chose to return error if the request can not be fulfilled.
-	Get(size int) (*Buffer, error)
-}
+import "fmt"
 
 // Recycler holds state necessary for a correct Buffer return to its originating Source
 type Recycler interface {
 	Recycle(*Buffer)
 }
 
 // Buffer is a reusable slice of bytes of fixed length.
-// The returned Data slice must not be retained past Release.
+// Buffer or its Data must not be retained past Release.
 type Buffer struct {
+	// Data len tracks valid payload + offset (offset passed out of band).
+	// It starts equal to the requested size for new buffers, and can be adjusted.
+	// Data is read-only, exposed for convenience but should never be assigned to.
+	// Buffer methods maintain the invariant that Data[offset:len(Data)] is valid payload,
+	// SetLen and Shift let callers adjust the offset and length of the valid payload as needed.
 	data     []byte
 	recycler Recycler
 }
 
-// New creates a standalone Buffer.
+// New creates Buffer referencing the provided Recycler.
 func New(b []byte, recycler Recycler) *Buffer {
 	return &Buffer{data: b, recycler: recycler}
 }
 
-// Make creates a standalone Buffer with a new byte slice of the requested size.
+// Make creates Buffer with a new byte slice of the requested size.
 func Make(size int) *Buffer {
-	return &Buffer{data: make([]byte, size), recycler: nil}
+	buf, _ := DefaultSource.Get(size) // fragment pool never errors
+	return buf
+}
+
+// ShiftOffset moves the valid data in the Buffer from oldOffset to newOffset, and adjusts the length accordingly.
+func (b *Buffer) ShiftOffset(oldOffset, newOffset int) int {
+	packetLen := len(b.data) - oldOffset
+	if newOffset+packetLen > cap(b.data) || newOffset < 0 {
+		return -1
+	}
+	n := copy(b.data[newOffset:newOffset+packetLen], b.data[oldOffset:])
+	if n != packetLen {
+		panic(fmt.Sprintf("short copy: %d != %d", n, packetLen))
+	}
+	b.data = b.data[:newOffset+n]
+	return n
 }
 
-// Data returns the full underlying byte slice of the Buffer.
-func (b *Buffer) Data() []byte {
+func (b *Buffer) Bytes() []byte {
 	return b.data
 }
 
-// Release returns the Buffer to its originating Source for reuse.
-// The Buffer must not be used after calling Release.
-func (b *Buffer) Release() {
-	if b.recycler != nil {
-		clear(b.data)
-		b.recycler.Recycle(b)
-	}
+// SetLen sets the length of the valid data in the Buffer.
+// Intended to be used for truncating the valid data post read,
+// or extending post encryption. Does not check the capacity.
+func (b *Buffer) SetLen(l int) {
+	b.data = b.data[:l]
 }
 
-// ReleaseAll calls Release on each non-nil Buffer in the slice, and sets the slice elements to nil.
-func ReleaseAll(bs []*Buffer) {
-	for i := range bs {
-		if bs[i] != nil {
-			bs[i].Release()
-			bs[i] = nil
+// Ensure returns a Buffer of the requested len, with the valid data from the provided Buffer.
+// The returned Buffer may be the same as the provided Buffer if it has sufficient capacity, or a new Buffer otherwise.
+// Safe to call on a nil Buffer.
+func Ensure(b *Buffer, size int, src Source) (*Buffer, error) {
+	if src == nil {
+		src = DefaultSource
+	}
+	if b == nil {
+		return src.Get(size)
+	}
+	if size > cap(b.data) {
+		bb, err := src.Get(size)
+		if err != nil {
+			return nil, err
+		}
+		n := copy(bb.data, b.data)
+		if n != len(b.data) {
+			panic(fmt.Sprintf("short copy: %d != %d", n, len(b.data)))
 		}
+		Release(b)
+		return bb, nil
 	}
+	b.data = b.data[:size]
+	return b, nil
 }
 
-// Arena is a Buffer with an internal watermark for sequential allocations.
-// FIXME Arena needs a graceful fallback on overflow.
-type Arena struct {
-	*Buffer
-	watermark int
-}
-
-// Get returns a slice of the Arena's Buffer of the requested size, and advances the watermark.
-func (a *Arena) Get(size int) []byte {
-	if a.watermark+size > len(a.Buffer.Data()) {
-		panic("arena overflow") // or return a heap-allocated fallback
+// Release returns Buffer to its Source for reuse.
+// Safe to call on a nil Buffer.
+func Release(b *Buffer) {
+	if b == nil {
+		return
+	}
+	b.data = b.data[:cap(b.data)]
+	clear(b.data)
+	if b.recycler != nil {
+		b.recycler.Recycle(b)
 	}
-	b := a.Buffer.Data()[a.watermark : a.watermark+size]
-	a.watermark += size
-	return b
 }
 
-// Flush resets the Arena's watermark to zero, and clears the valid data in the Buffer.
-func (a *Arena) Flush() {
-	clear(a.Buffer.Data()[:a.watermark])
-	a.watermark = 0
+// ReleaseAll calls Release on each non-nil Buffer in the slice, and sets the slice elements to nil.
+func ReleaseAll(bs []*Buffer) {
+	for i := range bs {
+		Release(bs[i])
+		bs[i] = nil
+	}
 }

buffer/constants.go

@@ -0,0 +1,9 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+package buffer
+
+const (
+	MaxMessageSize = MaxSegmentSize
+)

buffer/constants_android.go

@@ -0,0 +1,13 @@
+//go:build android
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+const (
+	MaxSegmentSize    = 2200
+	MaxBytesPerSource = 4096 * MaxSegmentSize
+)

buffer/constants_default.go

@@ -0,0 +1,13 @@
+//go:build !android && !ios && !windows
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+const (
+	MaxSegmentSize    = (1 << 16) - 1 // largest possible UDP datagram
+	MaxBytesPerSource = 0             // Disable and allow for infinite memory growth
+)

buffer/constants_ios.go

@@ -0,0 +1,17 @@
+//go:build ios
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+// Fit within memory limits for iOS's Network Extension API, which has stricter requirements.
+// These are vars instead of consts, because heavier network extensions might want to reduce
+// them further.
+var (
+	MaxBytesPerSource = 1024 * MaxSegmentSize
+)
+
+const MaxSegmentSize = 1700

buffer/constants_windows.go

@@ -0,0 +1,13 @@
+//go:build windows
+
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+const (
+	MaxSegmentSize    = 2048 - 32 // largest possible UDP datagram
+	MaxBytesPerSource = 0         // Disable and allow for infinite memory growth
+)