tun, conn, device: allocate buffers in the I/O devices

Previously crypto device maintained a batch of preallocated
to the MaxMessageSize buffers that the I/O only needs to
read into.

This change inverts the buffer ownership. A (wrapped) nil
pointer is passed into I/O. Device expects a backing
array to be allocated, and a slice of read data cut to
offset+size returned from the read within the same wrapper.

The wrapper is defined in buffer.Buffer, and carries
the backing reference and an optional Recycler implementation
to return the backing for reuse.

I/O is encouraged to implement a buffer management solution
that works best for its host OS. A shared sync.Pool is
provided as a default option.

Updates tailscale/corp#36989

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

Author: illotum

conn/bind_std.go

@@ -16,6 +16,7 @@ import (
 	"sync"
 	"syscall"
 
+	"github.com/tailscale/wireguard-go/device/buffer"
 	"golang.org/x/net/ipv4"
 	"golang.org/x/net/ipv6"
 )
@@ -233,13 +234,13 @@ func (s *StdNetBind) receiveIP(
 	br batchReader,
 	conn *net.UDPConn,
 	rxOffload bool,
-	bufs [][]byte,
-	sizes []int,
+	bufs []buffer.Buffer,
 	eps []Endpoint,
 ) (n int, err error) {
 	msgs := s.getMessages()
+	buffer.EnsureAllocated(bufs)
 	for i := range bufs {
-		(*msgs)[i].Buffers[0] = bufs[i]
+		(*msgs)[i].Buffers[0] = bufs[i].Data
 		(*msgs)[i].OOB = (*msgs)[i].OOB[:cap((*msgs)[i].OOB)]
 	}
 	defer s.putMessages(msgs)
@@ -271,8 +272,8 @@ func (s *StdNetBind) receiveIP(
 	}
 	for i := 0; i < numMsgs; i++ {
 		msg := &(*msgs)[i]
-		sizes[i] = msg.N
-		if sizes[i] == 0 {
+		bufs[i].Data = bufs[i].Data[:msg.N]
+		if len(bufs[i].Data) == 0 {
 			continue
 		}
 		addrPort := msg.Addr.(*net.UDPAddr).AddrPort()
@@ -284,14 +285,14 @@ func (s *StdNetBind) receiveIP(
 }
 
 func (s *StdNetBind) makeReceiveIPv4(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {
-	return func(bufs [][]byte, sizes []int, eps []Endpoint) (n int, err error) {
-		return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)
+	return func(bufs []buffer.Buffer, eps []Endpoint) (n int, err error) {
+		return s.receiveIP(pc, conn, rxOffload, bufs, eps)
 	}
 }
 
 func (s *StdNetBind) makeReceiveIPv6(pc *ipv6.PacketConn, conn *net.UDPConn, rxOffload bool) ReceiveFunc {
-	return func(bufs [][]byte, sizes []int, eps []Endpoint) (n int, err error) {
-		return s.receiveIP(pc, conn, rxOffload, bufs, sizes, eps)
+	return func(bufs []buffer.Buffer, eps []Endpoint) (n int, err error) {
+		return s.receiveIP(pc, conn, rxOffload, bufs, eps)
 	}
 }
 

conn/bind_std_test.go

@@ -5,6 +5,7 @@ import (
 	"net"
 	"testing"
 
+	"github.com/tailscale/wireguard-go/device/buffer"
 	"golang.org/x/net/ipv6"
 )
 
@@ -15,15 +16,14 @@ func TestStdNetBindReceiveFuncAfterClose(t *testing.T) {
 		t.Fatal(err)
 	}
 	bind.Close()
-	bufs := make([][]byte, 1)
-	bufs[0] = make([]byte, 1)
-	sizes := make([]int, 1)
+	bufs := make([]buffer.Buffer, 1)
+	bufs[0] = buffer.Buffer{Data: make([]byte, 1)}
 	eps := make([]Endpoint, 1)
 	for _, fn := range fns {
 		// The ReceiveFuncs must not access conn-related fields on StdNetBind
 		// unguarded. Close() nils the conn-related fields resulting in a panic
 		// if they violate the mutex.
-		fn(bufs, sizes, eps)
+		fn(bufs, eps)
 	}
 }
 

conn/bind_windows.go

@@ -18,6 +18,7 @@ import (
 	"golang.org/x/sys/windows"
 
 	"github.com/tailscale/wireguard-go/conn/winrio"
+	"github.com/tailscale/wireguard-go/device/buffer"
 )
 
 const (
@@ -416,20 +417,22 @@ retry:
 	return n, &ep, nil
 }
 
-func (bind *WinRingBind) receiveIPv4(bufs [][]byte, sizes []int, eps []Endpoint) (int, error) {
+func (bind *WinRingBind) receiveIPv4(bufs []buffer.Buffer, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
-	n, ep, err := bind.v4.Receive(bufs[0], &bind.isOpen)
-	sizes[0] = n
+	buffer.EnsureAllocated(bufs[:1])
+	n, ep, err := bind.v4.Receive(bufs[0].Data, &bind.isOpen)
+	bufs[0].Data = bufs[0].Data[:n]
 	eps[0] = ep
 	return 1, err
 }
 
-func (bind *WinRingBind) receiveIPv6(bufs [][]byte, sizes []int, eps []Endpoint) (int, error) {
+func (bind *WinRingBind) receiveIPv6(bufs []buffer.Buffer, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
-	n, ep, err := bind.v6.Receive(bufs[0], &bind.isOpen)
-	sizes[0] = n
+	buffer.EnsureAllocated(bufs[:1])
+	n, ep, err := bind.v6.Receive(bufs[0].Data, &bind.isOpen)
+	bufs[0].Data = bufs[0].Data[:n]
 	eps[0] = ep
 	return 1, err
 }

conn/bindtest/bindtest.go

@@ -13,6 +13,7 @@ import (
 	"os"
 
 	"github.com/tailscale/wireguard-go/conn"
+	"github.com/tailscale/wireguard-go/device/buffer"
 )
 
 type ChannelBind struct {
@@ -94,13 +95,14 @@ func (c *ChannelBind) BatchSize() int { return 1 }
 func (c *ChannelBind) SetMark(mark uint32) error { return nil }
 
 func (c *ChannelBind) makeReceiveFunc(ch chan []byte) conn.ReceiveFunc {
-	return func(bufs [][]byte, sizes []int, eps []conn.Endpoint) (n int, err error) {
+	return func(bufs []buffer.Buffer, eps []conn.Endpoint) (n int, err error) {
 		select {
 		case <-c.closeSignal:
 			return 0, net.ErrClosed
 		case rx := <-ch:
-			copied := copy(bufs[0], rx)
-			sizes[0] = copied
+			buffer.EnsureAllocated(bufs[:1])
+			n := copy(bufs[0].Data, rx)
+			bufs[0].Data = bufs[0].Data[:n]
 			eps[0] = c.target6
 			return 1, nil
 		}

conn/conn.go

@@ -13,19 +13,20 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+
+	"github.com/tailscale/wireguard-go/device/buffer"
 )
 
 const (
 	IdealBatchSize = 128 // maximum number of packets handled per read and write
 )
 
-// A ReceiveFunc receives at least one packet from the network and writes them
-// into packets. On a successful read it returns the number of elements of
-// sizes, packets, and endpoints that should be evaluated. Some elements of
-// sizes may be zero, and callers should ignore them. Callers must pass a sizes
-// and eps slice with a length greater than or equal to the length of packets.
-// These lengths must not exceed the length of the associated Bind.BatchSize().
-type ReceiveFunc func(packets [][]byte, sizes []int, eps []Endpoint) (n int, err error)
+// A ReceiveFunc receives at least one packet from the network into bufs.
+// On a successful read it returns the number of elements of bufs and eps
+// that should be evaluated. Callers must pass an eps slice with a length
+// greater than or equal to the length of bufs. These lengths must not
+// exceed the length of the associated Bind.BatchSize().
+type ReceiveFunc func(bufs []buffer.Buffer, eps []Endpoint) (n int, err error)
 
 // A Bind listens on a port for both IPv6 and IPv4 UDP traffic.
 //

conn/conn_test.go

@@ -7,11 +7,13 @@ package conn
 
 import (
 	"testing"
+
+	"github.com/tailscale/wireguard-go/device/buffer"
 )
 
 func TestPrettyName(t *testing.T) {
 	var (
-		recvFunc ReceiveFunc = func(bufs [][]byte, sizes []int, eps []Endpoint) (n int, err error) { return }
+		recvFunc ReceiveFunc = func(bufs []buffer.Buffer, eps []Endpoint) (n int, err error) { return }
 	)
 
 	const want = "TestPrettyName"

device/buffer/buffer.go

@@ -0,0 +1,65 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+// Package buffer provides pooled packet buffers for the I/O pipeline.
+// Each [Buffer] carries one packet and a recycle function that returns
+// its backing storage to the originating pool on [Release].
+package buffer
+
+import "unsafe"
+
+// RecycleHandle is the opaque reference to the [Buffer] backing arrays.
+type RecycleHandle unsafe.Pointer
+
+// Recycler returns a backing array for reuse.
+// The argument is the Backing of the [Buffer] being released.
+type Recycler interface {
+	Recycle(RecycleHandle)
+}
+
+// RecycleFunc is a function adapter for Recycler.
+type RecycleFunc func(RecycleHandle)
+
+func (f RecycleFunc) Recycle(ptr RecycleHandle) { f(ptr) }
+
+// Buffer is the packet envelope. Meant to be a value type,
+// allocated once per goroutine and reused across read cycles.
+type Buffer struct {
+	// Recycle returns the backing array to its pool.
+	// Nil for unmanaged/external Buffers.
+	Recycler Recycler
+
+	// Backing is the raw pointer to the byte array,
+	// anonymized to allow varying array and recycler implementations.
+	// Zero for external/unmanaged Buffers.
+	Backing RecycleHandle
+
+	// Data holds the bounded packet data, sliced from the backing array,
+	// and may be re-sliced by the caller. Do not append() on this slice.
+	// Nil for uninitialized Buffers.
+	Data []byte
+}
+
+// Release returns the backing data to its pool and zeros the Buffer.
+func (b *Buffer) Release() {
+	if b.Recycler != nil {
+		b.Recycler.Recycle(b.Backing)
+	}
+	*b = Buffer{}
+}
+
+// Claim transfers ownership: returns a copy of the Buffer and zeros the source.
+func (b *Buffer) Claim() Buffer {
+	c := *b
+	*b = Buffer{}
+	return c
+}
+
+// ReleaseAll releases each Buffer in the slice.
+func ReleaseAll(bufs []Buffer) {
+	for i := range bufs {
+		bufs[i].Release()
+	}
+}

device/buffer/constants.go

@@ -0,0 +1,10 @@
+/* SPDX-License-Identifier: MIT
+ *
+ * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
+ */
+
+package buffer
+
+const (
+	MaxBufferSize = MaxReadSize // the largest buffer that I/O may attempt to read or write.
+)

device/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 (
+	MaxReadSize      = 2200
+	MaxPooledBuffers = 4096
+)

device/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 (
+	MaxReadSize      = (1 << 16) - 1
+	MaxPooledBuffers = 0 // Disable and allow for infinite memory growth
+)