relay: Batch forwarded packets on the relay host

udp/conn.go

@@ -41,7 +41,7 @@ type Conn interface {
 	//
 	// Returns the number of packets successfully written. A destination the kernel rejects costs only
 	// its own packet, so a short count means some peers were undeliverable, not that the batch failed.
-	// Not safe for concurrent use on the same Conn.
+	// Safe for concurrent use on the same Conn.
 	WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error)
 	// ReloadConfig applies the listen settings that can change, r is nil at startup
 	ReloadConfig(c *config.C, r *config.Reload)

udp/udp_linux_writebatch.go

@@ -12,6 +12,7 @@ import (
 	"net/netip"
 	"strconv"
 	"strings"
+	"sync"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
@@ -20,7 +21,7 @@ import (
 // batchWriter owns the sendmmsg(2)/UDP-GSO transmit path for a StdConn: the
 // scratch WriteBatch packs mmsghdr entries into, plus the GSO capability
 // state probed at socket creation. Each queue has its own StdConn and
-// batchWriter, so no locking is needed.
+// batchWriter; mu serializes batches from the TUN reader and relay forwarding.
 //
 // Terminology, smallest to largest:
 //
@@ -39,6 +40,7 @@ import (
 //	chunk    the entries packed for one sendmmsg call, at most MaxWriteBatch.
 //	batch    the caller's whole bufs/addrs pair, processed as one or more chunks.
 type batchWriter struct {
+	mu   sync.Mutex
 	fd   int
 	isV4 bool
 
@@ -192,6 +194,9 @@ func parseRelease(r string) (major, minor int) {
 // Returns the number of packets sent. An error means the call itself failed.
 // A short count means some destinations were undeliverable.
 func (w *batchWriter) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
 	if len(bufs) != len(addrs) {
 		return 0, fmt.Errorf("WriteBatch: len(bufs)=%d != len(addrs)=%d", len(bufs), len(addrs))
 	}

udp/udp_linux_writebatch_concurrent_test.go

@@ -0,0 +1,61 @@
+//go:build linux && !android && !e2e_testing
+
+package udp
+
+import (
+	"net/netip"
+	"runtime"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// The TUN reader and relay forwarding share a socket. Exercise their concurrent
+// batches against the actual GSO planner and scratch, without relying on a NIC.
+func TestWriteBatchConcurrent(t *testing.T) {
+	const batches = 64
+	w := newRewindTestWriter()
+	sent := make(chan [][]byte, 2*batches)
+	w.sendFn = func(start, n int) (int, error) {
+		runtime.Gosched()
+		sent <- capturePrepared(w, start, n)
+		return n, nil
+	}
+	conn := &StdConn{bw: w}
+	remote := netip.MustParseAddrPort("192.0.2.1:4242")
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	for producer := range 2 {
+		wg.Go(func() {
+			<-start
+			for i := range batches {
+				var bufs [][]byte
+				var addrs []netip.AddrPort
+				for j := range 4 {
+					bufs = append(bufs, []byte{byte(producer), byte(i), byte(j), 0})
+					addrs = append(addrs, remote)
+				}
+				n, err := conn.WriteBatch(bufs, addrs)
+				if err != nil || n != len(bufs) {
+					t.Errorf("WriteBatch = %d, %v; want %d, nil", n, err, len(bufs))
+					return
+				}
+			}
+		})
+	}
+	close(start)
+	wg.Wait()
+	close(sent)
+	var counts [2]int
+	for packets := range sent {
+		require.Len(t, packets, 4)
+		producer := packets[0][0]
+		require.Less(t, int(producer), len(counts))
+		for j, packet := range packets {
+			require.Equal(t, []byte{producer, byte(counts[producer]), byte(j), 0}, packet)
+		}
+		counts[producer]++
+	}
+	require.Equal(t, [2]int{batches, batches}, counts)
+}