relay: Batch forwarded packets on the relay host

interface.go

@@ -384,6 +384,8 @@ type rxContext struct {
 	hostmapCache map[uint32]*HostInfo
 	lhh          *LightHouseHandler
 	ctCache      *firewall.ConntrackCacheTicker
+	relayBatch   *batch.SendBatch
+	relayRemote  netip.AddrPort
 }
 
 func newRxContext(f *Interface, q int) *rxContext {
@@ -414,6 +416,7 @@ func (f *Interface) listenOut(i int) {
 	}
 
 	flusher := func() {
+		f.flushRelayBatch(rxc)
 		if err := f.batchers[i].Flush(); err != nil {
 			f.l.Error("Failed to flush tun coalescer", "error", err)
 		}

outside.go

@@ -25,7 +25,7 @@ var ErrOutOfWindow = errors.New("out of window packet")
 
 // readOutsidePackets processes one received underlay packet.
 // Message payloads are decrypted IN PLACE, so packet must stay untouched
-// by the caller until the batcher for queue q has been flushed
+// by the caller until the TUN and relay batches for queue q have been flushed.
 func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxContext) {
 	h := rxc.h
 	err := h.Parse(packet)
@@ -77,6 +77,12 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 		f.messageMetrics.Rx(h.Type, h.Subtype, 1)
 	}
 
+	// Earlier forwards must leave before this packet can generate a response.
+	isMessageRelay := h.Type == header.Message && h.Subtype == header.MessageRelay
+	if !isMessageRelay {
+		f.flushRelayBatch(rxc)
+	}
+
 	// Unencrypted packets
 	switch h.Type {
 	case header.Handshake:
@@ -88,9 +94,6 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 		return
 	}
 
-	// Relay packets are special
-	isMessageRelay := (h.Type == header.Message && h.Subtype == header.MessageRelay)
-
 	var hostinfo *HostInfo
 	if isMessageRelay {
 		hostinfo = f.hostMap.QueryRelayIndex(h.RemoteIndex)
@@ -102,6 +105,8 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 	// recvError if necessary
 	if hostinfo == nil || hostinfo.ConnectionState == nil {
 		if !via.IsRelayed {
+			// An unknown relay index can also produce a response.
+			f.flushRelayBatch(rxc)
 			f.maybeSendRecvError(via.UdpAddr, h.RemoteIndex)
 		}
 		return
@@ -233,8 +238,10 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 				// Forward this packet through the relay tunnel, rebuilding it in place.
 				// Encode overwrites the old outer header, and the new AEAD tag lands where the old one was
 				fwdBuf := packet[:0]
-				//todo it would potentially be nice to batch these
-				f.SendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true, rxc.q)
+				out, err := f.prepareSendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true)
+				if err == nil {
+					f.queueRelayPacket(out, targetHI.GetRemote(), rxc)
+				}
 			case TerminalType:
 				hostinfo.logger(f.l).Error("Unexpected Relay Type of Terminal")
 				return
@@ -259,6 +266,27 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 	}
 }
 
+// queueRelayPacket borrows a receive buffer until the batch is flushed. Destination changes and the send batch
+// limit end a run, preserving arrival order while allowing WriteBatch to use UDP segmentation offload.
+func (f *Interface) queueRelayPacket(packet []byte, remote netip.AddrPort, rxc *rxContext) {
+	if rxc.relayBatch == nil {
+		rxc.relayBatch = batch.NewSendBatch(f.writers[rxc.q], batch.SendBatchCap, 0)
+	} else if rxc.relayRemote != remote {
+		f.flushRelayBatch(rxc)
+	}
+	rxc.relayRemote = remote
+	rxc.relayBatch.Commit(packet, remote)
+	if rxc.relayBatch.Len() >= batch.SendBatchCap {
+		f.flushRelayBatch(rxc)
+	}
+}
+
+func (f *Interface) flushRelayBatch(rxc *rxContext) {
+	if rxc.relayBatch != nil {
+		f.flushSendBatch(rxc.relayBatch, rxc.q)
+	}
+}
+
 // closeTunnel closes a tunnel locally, it does not send a closeTunnel packet to the remote
 func (f *Interface) closeTunnel(hostInfo *HostInfo) {
 	// Hold the hostmap lock across both the delete and the lighthouse eviction,

outside_relay_test.go

@@ -0,0 +1,291 @@
+package nebula
+
+import (
+	"bytes"
+	"net/netip"
+	"sync"
+	"testing"
+
+	"github.com/flynn/noise"
+	"github.com/gaissmai/bart"
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/noiseutil"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/test"
+	"github.com/slackhq/nebula/udp"
+	"github.com/stretchr/testify/require"
+)
+
+type relayBatchWrite struct {
+	packets [][]byte
+	remotes []netip.AddrPort
+	batched bool
+}
+
+type relayBatchConn struct {
+	udp.NoopConn
+	read   func(udp.EncReader, func())
+	writes []relayBatchWrite
+}
+
+func (c *relayBatchConn) ListenOut(r udp.EncReader, flush func()) error {
+	c.read(r, flush)
+	return nil
+}
+
+func (c *relayBatchConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	w := relayBatchWrite{remotes: append([]netip.AddrPort(nil), addrs...), batched: true}
+	for _, b := range bufs {
+		w.packets = append(w.packets, bytes.Clone(b))
+	}
+	c.writes = append(c.writes, w)
+	return len(bufs), nil
+}
+
+func (c *relayBatchConn) WriteTo(b []byte, addr netip.AddrPort) error {
+	c.writes = append(c.writes, relayBatchWrite{packets: [][]byte{bytes.Clone(b)}, remotes: []netip.AddrPort{addr}})
+	return nil
+}
+
+func (c *relayBatchConn) packets() [][]byte {
+	var packets [][]byte
+	for _, w := range c.writes {
+		packets = append(packets, w.packets...)
+	}
+	return packets
+}
+
+type relayBatchPath struct {
+	f                  *Interface
+	conn               *relayBatchConn
+	source, target     *HostInfo
+	incoming, outgoing *Relay
+	sourceKey          noiseutil.CipherState
+	cipher             noise.CipherFunc
+	counter            uint64
+}
+
+// Separate cipher instances let the reference SendVia path use the same keys and
+// counters without rewinding an AEAD instance (including FIPS nonce checks).
+func relayBatchKey(cipher noise.CipherFunc, id byte) noiseutil.CipherState {
+	suite := noise.NewCipherSuite(noise.DH25519, cipher, noise.HashSHA256)
+	return noiseutil.NewCipherState(noise.UnsafeNewCipherState(suite, [32]byte{id}, 0), cipher)
+}
+
+func newRelayBatchPath(t *testing.T, cipher noise.CipherFunc) *relayBatchPath {
+	t.Helper()
+	l := test.NewLogger()
+	conn := &relayBatchConn{}
+	f := &Interface{
+		ctx: t.Context(), l: l, hostMap: newHostMap(l), outside: conn,
+		writers: []udp.Conn{conn}, lightHouse: &LightHouse{l: l},
+		myVpnNetworksTable: &bart.Lite{}, metricTxDropped: metrics.NewCounter(),
+		batchers: []*batch.MultiCoalescer{batch.NewMultiCoalescer(&capturingTun{}, l)},
+	}
+	f.connectionManager = &connectionManager{
+		intf: f, relayUsed: map[uint32]struct{}{}, relayUsedLock: &sync.RWMutex{},
+	}
+	newHost := func(addr string, remote string, id byte) *HostInfo {
+		h := &HostInfo{
+			vpnAddrs: []netip.Addr{netip.MustParseAddr(addr)},
+			ConnectionState: &ConnectionState{
+				eKey: relayBatchKey(cipher, id), dKey: relayBatchKey(cipher, id), window: NewBits(ReplayWindow),
+			},
+			relayState: RelayState{relayForByAddr: map[netip.Addr]*Relay{}, relayForByIdx: map[uint32]*Relay{}},
+		}
+		u := netip.MustParseAddrPort(remote)
+		h.remote.Store(&u)
+		f.hostMap.Hosts[h.vpnAddrs[0]] = h
+		return h
+	}
+	source := newHost("192.0.2.1", "198.51.100.1:4242", 1)
+	target := newHost("192.0.2.2", "198.51.100.2:4242", 2)
+	incoming := &Relay{Type: ForwardingType, State: Established, LocalIndex: 301, RemoteIndex: 302, PeerAddr: target.vpnAddrs[0]}
+	outgoing := &Relay{Type: ForwardingType, State: Established, LocalIndex: 303, RemoteIndex: 304, PeerAddr: source.vpnAddrs[0]}
+	source.relayState.InsertRelay(incoming.PeerAddr, incoming.LocalIndex, incoming)
+	target.relayState.InsertRelay(outgoing.PeerAddr, outgoing.LocalIndex, outgoing)
+	f.hostMap.Relays[incoming.LocalIndex] = source
+	f.hostMap.Indexes[100] = source
+	return &relayBatchPath{f: f, conn: conn, source: source, target: target, incoming: incoming, outgoing: outgoing, sourceKey: relayBatchKey(cipher, 1), cipher: cipher}
+}
+
+func (p *relayBatchPath) wrap(t *testing.T, payload []byte) []byte {
+	t.Helper()
+	p.counter++
+	b := header.Encode(make([]byte, header.Len), header.Version, header.Message, header.MessageRelay, p.incoming.LocalIndex, p.counter)
+	b = append(b, payload...)
+	b, err := p.sourceKey.EncryptDanger(b, b, nil, p.counter, make([]byte, 12))
+	require.NoError(t, err)
+	return b
+}
+
+func (p *relayBatchPath) reference(payloads [][]byte) [][]byte {
+	conn := &relayBatchConn{}
+	f := &Interface{writers: []udp.Conn{conn}, connectionManager: p.f.connectionManager, l: p.f.l}
+	h := &HostInfo{ConnectionState: &ConnectionState{eKey: relayBatchKey(p.cipher, 2)}}
+	u := p.target.GetRemote()
+	h.remote.Store(&u)
+	for _, payload := range payloads {
+		f.SendVia(h, p.outgoing, payload, make([]byte, 12), make([]byte, 0, len(payload)+header.MaxOverhead), false, 0)
+	}
+	return conn.packets()
+}
+
+func TestRelayReceiveBatch(t *testing.T) {
+	for _, cipher := range []noise.CipherFunc{noiseutil.CipherAESGCM, noise.CipherChaChaPoly} {
+		t.Run(cipher.CipherName(), func(t *testing.T) {
+			p := newRelayBatchPath(t, cipher)
+			const count = 8
+			var payloads, wire [][]byte
+			for i := range count * 2 {
+				payload := bytes.Repeat([]byte{byte(i)}, 64)
+				payloads = append(payloads, payload)
+				wire = append(wire, p.wrap(t, payload))
+			}
+			want := p.reference(payloads)
+			// Model a listener reusing the same slots for the next receive batch.
+			buffers := make([][]byte, count)
+			for i := range buffers {
+				buffers[i] = make([]byte, len(wire[i]))
+			}
+			p.conn.read = func(r udp.EncReader, flush func()) {
+				for b := range 2 {
+					for i := range count {
+						copy(buffers[i], wire[b*count+i])
+						r(p.source.GetRemote(), buffers[i])
+						require.Len(t, p.conn.writes, b, "forwarding waits for the batch boundary")
+					}
+					flush()
+					require.Len(t, p.conn.writes, b+1, "flush must send before receive slots are reused")
+					for _, buffer := range buffers {
+						clear(buffer)
+					}
+				}
+			}
+			p.f.listenOut(0)
+			require.Equal(t, want, p.conn.packets())
+			for _, w := range p.conn.writes {
+				require.True(t, w.batched)
+				require.Len(t, w.packets, count)
+				for _, remote := range w.remotes {
+					require.Equal(t, p.target.GetRemote(), remote)
+				}
+			}
+		})
+	}
+}
+
+func TestRelayBatchDestinationChange(t *testing.T) {
+	p := newRelayBatchPath(t, noise.CipherChaChaPoly)
+	first, second := p.target.GetRemote(), netip.MustParseAddrPort("198.51.100.3:4242")
+	remotes := []netip.AddrPort{first, first, second, second, first}
+	var payloads, wire [][]byte
+	for i := range remotes {
+		payload := []byte{byte(i)}
+		payloads = append(payloads, payload)
+		wire = append(wire, p.wrap(t, payload))
+	}
+	want := p.reference(payloads)
+	p.conn.read = func(r udp.EncReader, flush func()) {
+		for i, remote := range remotes {
+			p.target.remote.Store(&remote)
+			r(p.source.GetRemote(), wire[i])
+		}
+		flush()
+	}
+	p.f.listenOut(0)
+	require.Equal(t, want, p.conn.packets())
+	require.Len(t, p.conn.writes, 3)
+	for i, wantRemotes := range [][]netip.AddrPort{{first, first}, {second, second}, {first}} {
+		require.Equal(t, wantRemotes, p.conn.writes[i].remotes)
+	}
+}
+
+func TestRelayBatchCapacity(t *testing.T) {
+	p := newRelayBatchPath(t, noise.CipherChaChaPoly)
+	p.conn.read = func(r udp.EncReader, flush func()) {
+		for i := range batch.SendBatchCap + 1 {
+			r(p.source.GetRemote(), p.wrap(t, []byte{byte(i)}))
+		}
+		require.Len(t, p.conn.writes, 1)
+		require.Len(t, p.conn.writes[0].packets, batch.SendBatchCap)
+		flush()
+	}
+	p.f.listenOut(0)
+	require.Len(t, p.conn.writes, 2)
+	require.Len(t, p.conn.writes[1].packets, 1)
+}
+
+func TestRelayBatchBeforeResponse(t *testing.T) {
+	for _, unknownRelay := range []bool{false, true} {
+		name := "test request"
+		if unknownRelay {
+			name = "unknown relay"
+		}
+		t.Run(name, func(t *testing.T) {
+			p := newRelayBatchPath(t, noise.CipherChaChaPoly)
+			first, last := p.wrap(t, []byte("first")), p.wrap(t, []byte("last"))
+			p.counter++
+			request := header.Encode(make([]byte, header.Len), header.Version, header.Test, header.TestRequest, 100, p.counter)
+			request, err := p.sourceKey.EncryptDanger(request, request, []byte("ping"), p.counter, make([]byte, 12))
+			require.NoError(t, err)
+			if unknownRelay {
+				request = header.Encode(make([]byte, header.Len), header.Version, header.Message, header.MessageRelay, 999, p.counter)
+			}
+			p.conn.read = func(r udp.EncReader, flush func()) {
+				r(p.source.GetRemote(), first)
+				require.Empty(t, p.conn.writes)
+				r(p.source.GetRemote(), request)
+				require.Len(t, p.conn.writes, 2)
+				r(p.source.GetRemote(), last)
+				flush()
+			}
+			p.f.listenOut(0)
+			require.Len(t, p.conn.writes, 3)
+			require.True(t, p.conn.writes[0].batched)
+			require.False(t, p.conn.writes[1].batched)
+			require.True(t, p.conn.writes[2].batched)
+			require.Equal(t, p.source.GetRemote(), p.conn.writes[1].remotes[0])
+			h := &header.H{}
+			require.NoError(t, h.Parse(p.conn.writes[1].packets[0]))
+			if unknownRelay {
+				require.Equal(t, header.RecvError, h.Type)
+				require.EqualValues(t, 999, h.RemoteIndex)
+			} else {
+				require.Equal(t, header.Test, h.Type)
+				require.Equal(t, header.TestReply, h.Subtype)
+				reply := p.conn.writes[1].packets[0]
+				out, err := p.sourceKey.DecryptDanger(nil, reply[:header.Len], reply[header.Len:], h.MessageCounter, make([]byte, 12))
+				require.NoError(t, err)
+				require.Equal(t, []byte("ping"), out)
+			}
+		})
+	}
+}
+
+func TestRelayBatchRemovedTarget(t *testing.T) {
+	for _, remove := range []string{"host", "relay", "disestablished"} {
+		t.Run(remove, func(t *testing.T) {
+			p := newRelayBatchPath(t, noise.CipherChaChaPoly)
+			first, removed := p.wrap(t, []byte("first")), p.wrap(t, []byte("removed"))
+			want := p.reference([][]byte{[]byte("first")})
+			p.conn.read = func(r udp.EncReader, flush func()) {
+				r(p.source.GetRemote(), first)
+				switch remove {
+				case "host":
+					delete(p.f.hostMap.Hosts, p.target.vpnAddrs[0])
+				case "relay":
+					delete(p.target.relayState.relayForByAddr, p.source.vpnAddrs[0])
+				case "disestablished":
+					p.target.relayState.UpdateRelayForByIpState(p.source.vpnAddrs[0], Disestablished)
+				}
+				r(p.source.GetRemote(), removed)
+				flush()
+			}
+			p.f.listenOut(0)
+			require.Len(t, p.conn.writes, 1)
+			require.Equal(t, want, p.conn.packets(), "only the packet accepted before removal is sent")
+		})
+	}
+}

overlay/batch/tx_batch.go

@@ -10,8 +10,8 @@ type batchWriter interface {
 }
 
 // SendBatch accumulates encrypted UDP packets and flushes them via WriteBatch.
-// One SendBatch is owned by each listenIn goroutine; no locking is needed.
-// Slots are backed by an Arena (see its docs)
+// Each reader goroutine owns its own SendBatch; no locking is needed.
+// Slots can borrow receive buffers or be backed by an Arena (see its docs).
 type SendBatch struct {
 	out   batchWriter
 	bufs  [][]byte

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)
+}