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/iobuf"
 	"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 []iobuf.View,
 	eps []Endpoint,
 ) (n int, err error) {
 	msgs := s.getMessages()
+	iobuf.EnsureAllocated(bufs)
 	for i := range bufs {
-		(*msgs)[i].Buffers[0] = bufs[i]
+		(*msgs)[i].Buffers[0] = bufs[i].Bytes
 		(*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].Bytes = bufs[i].Bytes[:msg.N]
+		if len(bufs[i].Bytes) == 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 []iobuf.View, 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 []iobuf.View, 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/iobuf"
 	"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([]iobuf.View, 1)
+	bufs[0] = iobuf.View{Bytes: 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/iobuf"
 )
 
 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 []iobuf.View, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
-	n, ep, err := bind.v4.Receive(bufs[0], &bind.isOpen)
-	sizes[0] = n
+	iobuf.EnsureAllocated(bufs[:1])
+	n, ep, err := bind.v4.Receive(bufs[0].Bytes, &bind.isOpen)
+	bufs[0].Bytes = bufs[0].Bytes[:n]
 	eps[0] = ep
 	return 1, err
 }
 
-func (bind *WinRingBind) receiveIPv6(bufs [][]byte, sizes []int, eps []Endpoint) (int, error) {
+func (bind *WinRingBind) receiveIPv6(bufs []iobuf.View, eps []Endpoint) (int, error) {
 	bind.mu.RLock()
 	defer bind.mu.RUnlock()
-	n, ep, err := bind.v6.Receive(bufs[0], &bind.isOpen)
-	sizes[0] = n
+	iobuf.EnsureAllocated(bufs[:1])
+	n, ep, err := bind.v6.Receive(bufs[0].Bytes, &bind.isOpen)
+	bufs[0].Bytes = bufs[0].Bytes[: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/iobuf"
 )
 
 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 []iobuf.View, 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
+			iobuf.EnsureAllocated(bufs[:1])
+			n := copy(bufs[0].Bytes, rx)
+			bufs[0].Bytes = bufs[0].Bytes[:n]
 			eps[0] = c.target6
 			return 1, nil
 		}

conn/conn.go

@@ -13,19 +13,20 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+
+	"github.com/tailscale/wireguard-go/iobuf"
 )
 
 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 []iobuf.View, 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/iobuf"
 )
 
 func TestPrettyName(t *testing.T) {
 	var (
-		recvFunc ReceiveFunc = func(bufs [][]byte, sizes []int, eps []Endpoint) (n int, err error) { return }
+		recvFunc ReceiveFunc = func(bufs []iobuf.View, eps []Endpoint) (n int, err error) { return }
 	)
 
 	const want = "TestPrettyName"

device/channels.go

@@ -93,7 +93,7 @@ func (device *Device) flushInboundQueue(q *autodrainingInboundQueue) {
 		case elemsContainer := <-q.c:
 			elemsContainer.Lock()
 			for _, elem := range elemsContainer.elems {
-				device.PutMessageBuffer(elem.buffer)
+				elem.buffer.Release()
 				device.PutInboundElement(elem)
 			}
 			device.PutInboundElementsContainer(elemsContainer)
@@ -126,7 +126,7 @@ func (device *Device) flushOutboundQueue(q *autodrainingOutboundQueue) {
 		case elemsContainer := <-q.c:
 			elemsContainer.Lock()
 			for _, elem := range elemsContainer.elems {
-				device.PutMessageBuffer(elem.buffer)
+				elem.buffer.Release()
 				device.PutOutboundElement(elem)
 			}
 			device.PutOutboundElementsContainer(elemsContainer)

device/constants.go

@@ -7,6 +7,8 @@ package device
 
 import (
 	"time"
+
+	"github.com/tailscale/wireguard-go/iobuf"
 )
 
 /* Specification constants */
@@ -27,9 +29,9 @@ const (
 )
 
 const (
-	MinMessageSize = MessageKeepaliveSize                                                      // minimum size of transport message (keepalive)
-	MaxMessageSize = MaxSegmentSize                                                            // maximum size of transport message
-	MaxContentSize = MaxSegmentSize - MessageTransportSize - MessageEncapsulatingTransportSize // maximum size of transport message content
+	MinMessageSize = MessageKeepaliveSize                                                         // minimum size of transport message (keepalive)
+	MaxMessageSize = iobuf.MaxReadSize                                                            // maximum size of transport message
+	MaxContentSize = iobuf.MaxReadSize - MessageTransportSize - MessageEncapsulatingTransportSize // maximum size of transport message content
 )
 
 /* Implementation constants */

device/device.go

@@ -74,7 +74,6 @@ type Device struct {
 	pool struct {
 		inboundElementsContainer  *waitpool.WaitPool
 		outboundElementsContainer *waitpool.WaitPool
-		messageBuffers            *waitpool.WaitPool
 		inboundElements           *waitpool.WaitPool
 		outboundElements          *waitpool.WaitPool
 	}

device/device_test.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/tailscale/wireguard-go/conn"
 	"github.com/tailscale/wireguard-go/conn/bindtest"
+	"github.com/tailscale/wireguard-go/iobuf"
 	"github.com/tailscale/wireguard-go/tun"
 	"github.com/tailscale/wireguard-go/tun/tuntest"
 )
@@ -437,7 +438,7 @@ type fakeTUNDeviceSized struct {
 }
 
 func (t *fakeTUNDeviceSized) File() *os.File { return nil }
-func (t *fakeTUNDeviceSized) Read(bufs [][]byte, sizes []int, offset int) (n int, err error) {
+func (t *fakeTUNDeviceSized) Read(bufs []iobuf.View, offset int) (n int, err error) {
 	return 0, nil
 }
 func (t *fakeTUNDeviceSized) Write(bufs [][]byte, offset int) (int, error) { return 0, nil }