tun, device: allocate buffers in the edge devices

Signed-off-by: Alex Valiushko <alexvaliushko@tailscale.com>
Change-Id: I58908d9d3fd09441e9378a74b0ee19136a6a6964

Author: illotum

buffer/buffer.go

@@ -0,0 +1,117 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+// Package buffer implements a reusable buffer abstraction.
+//
+// Wireguard-go's data processing is constrained by both the hosts API,
+// and the transformations performed during encapsulation:
+//
+//  1. Encryption requires tail- and headroom for extra headers and padding.
+//     Available via winrio, and pread(2).
+//  2. Systems are moving towards coalesced reads for both TCP and UDP.
+//     The read data has no gaps for individual slices.
+//  3. crypto.AEAD interface requires a contiguous dst []byte for Sealing.
+//     So we can't use scatter-gather to inject the gaps.
+//
+// Until one of these three conditions is changed, the encryption strategy
+// is to copy on read into buffers with the required gaps.
+// The buffers are right-sized for the packet to avoid memory inflation.
+// To recycle said allocations, each buffer carries a recycle function
+// that routes it back to its originating pool.
+//
+// Decryption shrinks each fragment instead of growing, so buffers can pass
+// through the pipeline without copying till the egress coalescence.
+// Depending on the chosen head of the coalescence, there may or may be no room
+// and reallocation is a necessary fallback until we start passing
+// buffers in batches.
+
+package buffer
+
+import "fmt"
+
+// Recycler holds state necessary for a correct Buffer return to its originating Source
+type Recycler interface {
+	Recycle(*Buffer)
+}
+
+// RecycleFunc adapts arbitrary closures to the Recycler interface.
+type RecycleFunc func(*Buffer)
+
+func (f RecycleFunc) Recycle(b *Buffer) {
+	f(b)
+}
+
+// Buffer is a reusable slice of bytes of fixed length.
+// Buffer or its Data must not be retained past Release.
+type Buffer struct {
+	data     []byte
+	recycler Recycler
+}
+
+// New creates Buffer referencing the provided data and Recycler.
+func New(data []byte, recycler Recycler) *Buffer {
+	return &Buffer{data: data, recycler: recycler}
+}
+
+// Bytes returns the valid data in the Buffer.
+func (b *Buffer) Bytes() []byte {
+	return b.data
+}
+
+// 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]
+}
+
+// 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. Performs no nil check on the Source or bounds check on the size.
+func Ensure(b *Buffer, size int, src Source) *Buffer {
+	if b == nil {
+		return src.Get(size)
+	}
+	if size > cap(b.data) {
+		bb := src.Get(size)
+		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
+	}
+	b.SetLen(size)
+	return b
+}
+
+// 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)]
+	if b.recycler != nil {
+		b.recycler.Recycle(b)
+	}
+}
+
+// Claim takes ownership of the Buffer at *slot, setting *slot to nil.
+// Returns the claimed Buffer (may be nil if slot was already nil).
+func Claim(slot **Buffer) *Buffer {
+	b := *slot
+	*slot = nil
+	return b
+}
+
+// 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/buffer_test.go

@@ -0,0 +1,112 @@
+package buffer
+
+import (
+	"testing"
+)
+
+type countingRecycler struct {
+	count int
+}
+
+func (r *countingRecycler) Recycle(b *Buffer) {
+	r.count++
+}
+
+func TestEnsure(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		buf      *Buffer
+		size     int
+		wantKept bool
+	}{
+		{"nil buffer", nil, 100, false},
+		{"sufficient cap", New(make([]byte, 200), &countingRecycler{}), 100, true},
+		{"insufficient cap", New(make([]byte, 100), &countingRecycler{}), 200, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var r *countingRecycler
+			if tc.buf != nil {
+				r = tc.buf.recycler.(*countingRecycler)
+			}
+			buf := Ensure(tc.buf, tc.size, NewPoolSource())
+			if len(buf.Bytes()) != tc.size {
+				t.Errorf("len = %d, want %d", len(buf.Bytes()), tc.size)
+			}
+			if tc.wantKept != (buf == tc.buf) {
+				t.Errorf("kept same buffer: got %v, want %v", buf == tc.buf, tc.wantKept)
+			}
+			if r != nil && !tc.wantKept && r.count != 1 {
+				t.Errorf("recycler not called, count = %d, want 1", r.count)
+			}
+		})
+	}
+}
+
+func TestRelease(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		buf  *Buffer
+	}{
+		{"nil buffer", nil},
+		{"nil recycler", New(nil, nil)},
+		{"non-nil recycler", New(make([]byte, 10), &countingRecycler{})},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var r *countingRecycler
+			if tc.buf != nil {
+				r, _ = tc.buf.recycler.(*countingRecycler)
+			}
+			Release(tc.buf) // must not panic
+			if r != nil && r.count != 1 {
+				t.Errorf("recycler not called, count = %d, want 1", r.count)
+			}
+		})
+	}
+
+}
+
+func TestClaim(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		slot *Buffer
+	}{
+		{"nil slot", nil},
+		{"non-nil slot", New(make([]byte, 10), nil)},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			want := tc.slot
+			claimed := Claim(&tc.slot)
+			if tc.slot != nil {
+				t.Error("slot should be nil after Claim")
+			}
+			if claimed != want {
+				t.Error("claimed buffer does not match original")
+			}
+		})
+	}
+}
+
+func TestReleaseAll(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		bufs []*Buffer
+	}{
+		{"nil slice", nil},
+		{"empty slice", []*Buffer{}},
+		{"slice with nil buffers", []*Buffer{nil, nil}},
+		{"slice with non-nil buffers", []*Buffer{New(make([]byte, 10), nil), New(make([]byte, 20), nil)}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			l := len(tc.bufs)
+			ReleaseAll(tc.bufs) // must not panic
+			if len(tc.bufs) != l {
+				t.Errorf("length of bufs changed: got %d, want %d", len(tc.bufs), l)
+			}
+			for i, b := range tc.bufs {
+				if b != nil {
+					t.Errorf("bufs[%d] should be nil after ReleaseAll", i)
+				}
+			}
+		})
+	}
+}

buffer/constants.go

@@ -0,0 +1,10 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+const (
+	MaxBufferSize = MaxSegmentSize // the largest buffer that callers may request from a Source.
+)

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 // largest possible Android read
+	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 Unix read
+	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 // largest possible iOS read

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 Windows read
+	MaxBytesPerSource = 0         // Disable and allow for infinite memory growth
+)