relay: Let WriteBatch handle destination changes and drop the flush before responses

interface.go

@@ -384,8 +384,8 @@ type rxContext struct {
 	hostmapCache map[uint32]*HostInfo
 	lhh          *LightHouseHandler
 	ctCache      *firewall.ConntrackCacheTicker
-	relayBatch   *batch.SendBatch
-	relayRemote  netip.AddrPort
+	// relayBatch holds forwarded relay packets, which borrow their receive buffers until flushed
+	relayBatch *batch.SendBatch
 }
 
 func newRxContext(f *Interface, q int) *rxContext {
@@ -398,6 +398,7 @@ func newRxContext(f *Interface, q int) *rxContext {
 		hostmapCache: map[uint32]*HostInfo{},
 		lhh:          f.lightHouse.NewRequestHandler(),
 		ctCache:      firewall.NewConntrackCacheTicker(f.ctx, f.l, f.conntrackCacheTimeout),
+		relayBatch:   batch.NewSendBatch(f.writers[q], batch.SendBatchCap, 0),
 	}
 }
 
@@ -416,7 +417,7 @@ func (f *Interface) listenOut(i int) {
 	}
 
 	flusher := func() {
-		f.flushRelayBatch(rxc)
+		f.flushSendBatch(rxc.relayBatch, i)
 		if err := f.batchers[i].Flush(); err != nil {
 			f.l.Error("Failed to flush tun coalescer", "error", err)
 		}

outside.go

@@ -77,12 +77,6 @@ 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:
@@ -94,6 +88,9 @@ 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)
@@ -105,8 +102,6 @@ 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
@@ -240,7 +235,11 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 				fwdBuf := packet[:0]
 				out, err := f.prepareSendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true)
 				if err == nil {
-					f.queueRelayPacket(out, targetHI.GetRemote(), rxc)
+					// Borrows the receive buffer, the listener's flush callback drains it before the buffer is reused
+					rxc.relayBatch.Commit(out, targetHI.GetRemote())
+					if rxc.relayBatch.Len() >= batch.SendBatchCap {
+						f.flushSendBatch(rxc.relayBatch, rxc.q)
+					}
 				}
 			case TerminalType:
 				hostinfo.logger(f.l).Error("Unexpected Relay Type of Terminal")
@@ -266,27 +265,6 @@ 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

@@ -176,7 +176,7 @@ func TestRelayReceiveBatch(t *testing.T) {
 	}
 }
 
-func TestRelayBatchDestinationChange(t *testing.T) {
+func TestRelayBatchMixedDestinations(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}
@@ -196,10 +196,8 @@ func TestRelayBatchDestinationChange(t *testing.T) {
 	}
 	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)
-	}
+	require.Len(t, p.conn.writes, 1, "WriteBatch splits runs by destination, the batch does not need to")
+	require.Equal(t, remotes, p.conn.writes[0].remotes)
 }
 
 func TestRelayBatchCapacity(t *testing.T) {
@@ -217,53 +215,6 @@ func TestRelayBatchCapacity(t *testing.T) {
 	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) {