tun: Prepare batch counter reservation for staged encryption

Build on the inside-packet filtering split with shared preparation and encryption helpers for inline and reserved sends. Reserve each batch’s per-tunnel counter range before encryption, including relay counters, without enabling staged processing yet. Test reserved encryption against a real receiver and check segmentation counts.

connection_state_test.go

@@ -23,11 +23,20 @@ import (
 // runTestHandshake runs a complete IX handshake between two freshly-built
 // peers and returns the initiator and responder Results. Used to produce
 // real cipher states for tests that need to exercise post-handshake glue.
-func runTestHandshake(t *testing.T) (initR, respR *handshake.Result) {
+func runTestHandshake(t testing.TB) (initR, respR *handshake.Result) {
+	t.Helper()
+	return runTestHandshakeWith(t, noise.CipherChaChaPoly)
+}
+
+// runTestHandshakeWith is runTestHandshake with the data plane cipher picked by the caller.
+func runTestHandshakeWith(t testing.TB, cipher noise.CipherFunc) (initR, respR *handshake.Result) {
 	t.Helper()
 
+	// Valid around now rather than the cert_test default, which is fixed at process start and lapses a minute
+	// later, inside a long -count run
+	now := time.Now()
 	ca, _, caKey, _ := ct.NewTestCaCert(
-		cert.Version2, cert.Curve_CURVE25519, time.Time{}, time.Time{}, nil, nil, nil,
+		cert.Version2, cert.Curve_CURVE25519, now.Add(-time.Minute), now.Add(time.Hour), nil, nil, nil,
 	)
 	caPool := ct.NewTestCAPool(ca)
 
@@ -40,7 +49,7 @@ func runTestHandshake(t *testing.T) (initR, respR *handshake.Result) {
 		require.NoError(t, err)
 		hsBytes, err := c.MarshalForHandshakes()
 		require.NoError(t, err)
-		ncs := noise.NewCipherSuite(noise.DH25519, noise.CipherChaChaPoly, noise.HashSHA256)
+		ncs := noise.NewCipherSuite(noise.DH25519, cipher, noise.HashSHA256)
 		cred := handshake.NewCredential(c, hsBytes, priv, ncs)
 		return func(v cert.Version) *handshake.Credential {
 			if v == cert.Version2 {

flow_test.go

@@ -0,0 +1,256 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"encoding/binary"
+	"io"
+	"math"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/flynn/noise"
+	"github.com/gaissmai/bart"
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/cert"
+	"github.com/slackhq/nebula/firewall"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/iputil"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/overlay/tio"
+	"github.com/slackhq/nebula/test"
+	"github.com/slackhq/nebula/udp"
+	"github.com/stretchr/testify/require"
+)
+
+// flowTestMSS is the segment payload size of the synthetic TSO superpackets.
+const flowTestMSS = 1200
+
+var (
+	flowTestLocal  = netip.MustParseAddr("10.0.0.1")
+	flowTestRemote = netip.MustParseAddr("10.0.0.2")
+	flowTestUnder  = netip.MustParseAddrPort("192.0.2.2:4242")
+)
+
+// flowTCPSuperpacket builds an IPv4/TCP TSO superpacket from local to remote carrying segs segments. Segment i's
+// payload starts with the flow id, the superpacket seq and i, so a receiver can check what it got and in which order.
+func flowTCPSuperpacket(flow uint16, seq uint32, segs int) tio.Packet {
+	const hdr = 40
+	pay := segs * flowTestMSS
+	b := make([]byte, hdr+pay)
+	b[0] = 0x45
+	binary.BigEndian.PutUint16(b[2:4], uint16(min(hdr+pay, 65535)))
+	b[8] = 64
+	b[9] = iputil.IPProtocolTCP
+	copy(b[12:16], flowTestLocal.AsSlice())
+	copy(b[16:20], flowTestRemote.AsSlice())
+	binary.BigEndian.PutUint16(b[20:22], 10000+flow)
+	binary.BigEndian.PutUint16(b[22:24], 80)
+	binary.BigEndian.PutUint32(b[24:28], seq*uint32(pay))
+	b[32] = 0x50
+	b[33] = 0x10 // ACK
+	binary.BigEndian.PutUint16(b[34:36], 65535)
+	for i := range segs {
+		p := b[hdr+i*flowTestMSS:]
+		binary.BigEndian.PutUint16(p[0:2], flow)
+		binary.BigEndian.PutUint32(p[2:6], seq)
+		binary.BigEndian.PutUint16(p[6:8], uint16(i))
+	}
+	if segs == 1 {
+		return tio.Packet{Bytes: b}
+	}
+	return tio.Packet{Bytes: b, GSO: tio.GSOInfo{Size: flowTestMSS, HdrLen: hdr, CsumStart: 20, Proto: tio.GSOProtoTCP}}
+}
+
+// decodedSeg identifies one segment by the marker flowTCPSuperpacket put in its payload.
+type decodedSeg struct {
+	flow uint16
+	seq  uint32
+	idx  uint16
+}
+
+func decodeSeg(p []byte) decodedSeg {
+	return decodedSeg{flow: binary.BigEndian.Uint16(p[0:2]), seq: binary.BigEndian.Uint32(p[2:6]), idx: binary.BigEndian.Uint16(p[6:8])}
+}
+
+// flowTestBatches builds n reads of per superpackets of segs segments, across flows interleaved flows, and the
+// segments they carry in order.
+func flowTestBatches(n, per, flows, segs int) ([][]tio.Packet, []decodedSeg) {
+	var batches [][]tio.Packet
+	var want []decodedSeg
+	seq := make([]uint32, flows)
+	for range n {
+		var b []tio.Packet
+		for j := range per {
+			flow := uint16(j % flows)
+			b = append(b, flowTCPSuperpacket(flow, seq[flow], segs))
+			for s := range segs {
+				want = append(want, decodedSeg{flow: flow, seq: seq[flow], idx: uint16(s)})
+			}
+			seq[flow]++
+		}
+		batches = append(batches, b)
+	}
+	return batches, want
+}
+
+// flowTestPeer is one end of a tunnel: the interface, the hostinfo for the other end, and its capture points.
+type flowTestPeer struct {
+	f        *Interface
+	hostinfo *HostInfo
+	conn     *flowCaptureConn
+	tun      *flowTunCapture
+}
+
+// newFlowTestPair builds two interfaces joined by one real tunnel (keys from an actual handshake). sender sends from
+// flowTestLocal to flowTestRemote, receiver is the far side.
+func newFlowTestPair(t testing.TB, cipher noise.CipherFunc) (sender, receiver *flowTestPeer) {
+	t.Helper()
+	initR, respR := runTestHandshakeWith(t, cipher)
+	sci, err := newConnectionStateFromResult(initR)
+	require.NoError(t, err)
+	rci, err := newConnectionStateFromResult(respR)
+	require.NoError(t, err)
+	sender = newFlowTestPeer(t, flowTestLocal, flowTestRemote, sci, 100, 200)
+	receiver = newFlowTestPeer(t, flowTestRemote, flowTestLocal, rci, 200, 100)
+	return sender, receiver
+}
+
+func newFlowTestPeer(t testing.TB, me, them netip.Addr, ci *ConnectionState, localIdx, remoteIdx uint32) *flowTestPeer {
+	l := test.NewLogger()
+	hm := newHostMap(l)
+	hm.preferredRanges.Store(&[]netip.Prefix{})
+
+	myNets, myAddrs := &bart.Lite{}, &bart.Lite{}
+	myNets.Insert(netip.PrefixFrom(me, 24))
+	myAddrs.Insert(netip.PrefixFrom(me, 32))
+	myCert := &dummyCert{version: cert.Version1, networks: []netip.Prefix{netip.PrefixFrom(me, 24)}}
+	theirCert := &dummyCert{version: cert.Version1, networks: []netip.Prefix{netip.PrefixFrom(them, 24)}, name: "peer"}
+	ci.peerCert = &cert.CachedCertificate{Certificate: theirCert, InvertedGroups: map[string]struct{}{}}
+
+	fw := NewFirewall(l, time.Minute, time.Minute, time.Minute, myCert)
+	require.NoError(t, fw.AddRule(false, firewall.ProtoAny, 0, 0, []string{"any"}, "", "", "", "", ""))
+	require.NoError(t, fw.AddRule(true, firewall.ProtoAny, 0, 0, []string{"any"}, "", "", "", "", ""))
+
+	conn, tun := &flowCaptureConn{}, &flowTunCapture{}
+	lh := newTestLighthouse()
+	lh.remoteAllowList.Store(&RemoteAllowList{})
+
+	f := &Interface{
+		hostMap:            hm,
+		outside:            conn,
+		firewall:           fw,
+		lightHouse:         lh,
+		pki:                &PKI{},
+		myVpnNetworksTable: myNets,
+		myVpnAddrsTable:    myAddrs,
+		routines:           1,
+		writers:            []udp.Conn{conn},
+		messageMetrics:     &MessageMetrics{txExhausted: metrics.NewCounter()},
+		metricTxDropped:    metrics.NilCounter{},
+		l:                  l,
+	}
+	f.ctx = t.Context()
+	f.pki.caPool.Store(cert.NewCAPool())
+	// Never promote or requery, those would put extra test packets on the wire
+	f.tryPromoteEvery.Store(math.MaxUint32)
+	f.reQueryEvery.Store(math.MaxUint32)
+	f.batchers = []*batch.MultiCoalescer{batch.NewMultiCoalescer(tun, l)}
+	f.connectionManager = &connectionManager{
+		intf:          f,
+		relayUsed:     map[uint32]struct{}{},
+		relayUsedLock: &sync.RWMutex{},
+		trafficTimer:  NewLockingTimerWheel[uint32](time.Second, time.Minute),
+	}
+	f.handshakeManager = NewHandshakeManager(l, hm, lh, conn, defaultHandshakeConfig)
+	f.handshakeManager.f = f
+
+	hostinfo := &HostInfo{
+		ConnectionState: ci,
+		vpnAddrs:        []netip.Addr{them},
+		localIndexId:    localIdx,
+		remoteIndexId:   remoteIdx,
+		relayState:      RelayState{relayForByAddr: map[netip.Addr]*Relay{}, relayForByIdx: map[uint32]*Relay{}},
+	}
+	hostinfo.remote.Store(&flowTestUnder)
+	hostinfo.buildNetworks(myNets, theirCert)
+	hm.unlockedAddHostInfo(hostinfo, f)
+
+	return &flowTestPeer{f: f, hostinfo: hostinfo, conn: conn, tun: tun}
+}
+
+// flowCaptureConn records every packet sent, in order, as copies.
+type flowCaptureConn struct {
+	udp.NoopConn
+	mu    sync.Mutex
+	pkts  [][]byte
+	addrs []netip.AddrPort
+	// delay, when set, slows every WriteBatch down so batches pile up in flight
+	delay time.Duration
+}
+
+func (c *flowCaptureConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	time.Sleep(c.delay)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for i, b := range bufs {
+		c.pkts = append(c.pkts, append([]byte(nil), b...))
+		c.addrs = append(c.addrs, addrs[i])
+	}
+	return len(bufs), nil
+}
+
+func (c *flowCaptureConn) WriteTo(b []byte, addr netip.AddrPort) error {
+	_, err := c.WriteBatch([][]byte{b}, []netip.AddrPort{addr})
+	return err
+}
+
+func (c *flowCaptureConn) captured() [][]byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([][]byte(nil), c.pkts...)
+}
+
+// flowTunCapture is a tun queue that records every plaintext write.
+type flowTunCapture struct {
+	mu     sync.Mutex
+	writes [][]byte
+	// delay, when set, slows every write down
+	delay time.Duration
+}
+
+func (c *flowTunCapture) Read() ([]tio.Packet, error) { return nil, io.EOF }
+func (c *flowTunCapture) Close() error                { return nil }
+func (c *flowTunCapture) Write(b []byte) (int, error) {
+	time.Sleep(c.delay)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.writes = append(c.writes, append([]byte(nil), b...))
+	return len(b), nil
+}
+
+// decryptWire checks every wire packet decrypts on the receiver, in order, with strictly increasing counters, and
+// returns what each one carried. It uses its own replay window and copies, so the wire can still be fed to the
+// receiver afterwards.
+func decryptWire(t testing.TB, receiver *flowTestPeer, wire [][]byte) []decodedSeg {
+	t.Helper()
+	peer := receiver.hostinfo.ConnectionState
+	ci := &ConnectionState{dKey: peer.dKey, window: NewBits(peer.window.length)}
+	nb := make([]byte, 12)
+	var last uint64
+	out := make([]decodedSeg, 0, len(wire))
+	for i, pkt := range wire {
+		h := &header.H{}
+		require.NoError(t, h.Parse(pkt))
+		require.Equal(t, header.Message, h.Type)
+		require.Equal(t, receiver.hostinfo.localIndexId, h.RemoteIndex)
+		require.Greater(t, h.MessageCounter, last, "wire packet %d: counter %d after %d", i, h.MessageCounter, last)
+		last = h.MessageCounter
+		plain, err := ci.Decrypt(receiver.f.l, h.MessageCounter, append([]byte(nil), pkt...), nb)
+		require.NoError(t, err, "wire packet %d", i)
+		out = append(out, decodeSeg(plain[40:]))
+	}
+	return out
+}

flow_tx.go

@@ -0,0 +1,233 @@
+package nebula
+
+import (
+	"errors"
+	"net/netip"
+
+	"github.com/slackhq/nebula/firewall"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/noiseutil"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/overlay/tio"
+)
+
+var errFlowSegmentOverrun = errors.New("superpacket has more segments than counters reserved for it")
+
+// txItem is one firewall-approved inside packet (or superpacket) with everything needed to encrypt it. A nonzero
+// counter is the first of the segs counters reserved for it on ci, relayCounter the same on the relay tunnel when
+// relayHI is set. Zero counters are allocated per segment while encrypting, as inline sends do.
+type txItem struct {
+	pkt          tio.Packet
+	hostinfo     *HostInfo
+	ci           *ConnectionState
+	remote       netip.AddrPort
+	segs         uint64
+	counter      uint64
+	relayHI      *HostInfo
+	relay        *Relay
+	relayCounter uint64
+}
+
+// prepareTxItem does the per-packet work sendInsideMessage does before it encrypts, and resolves where the packet goes.
+func (f *Interface) prepareTxItem(it *txItem, hostinfo *HostInfo, pkt tio.Packet) bool {
+	ci := hostinfo.ConnectionState
+	if ci.eKey == nil {
+		return false
+	}
+
+	f.recordOutbound(hostinfo)
+	it.pkt, it.hostinfo, it.ci = pkt, hostinfo, ci
+	it.remote = hostinfo.GetRemote()
+	if it.remote.IsValid() {
+		return true
+	}
+
+	it.relayHI, it.relay = f.findRelay(hostinfo)
+	if it.relayHI == nil || it.relay == nil {
+		return false
+	}
+	it.remote = it.relayHI.GetRemote()
+	return true
+}
+
+// encryptTxItem segments and encrypts it into sb, for inline and staged sends alike.
+func (f *Interface) encryptTxItem(it *txItem, nb []byte, sb *batch.SendBatch) {
+	reserved := it.counter != 0
+	c, rc, end := it.counter, it.relayCounter, it.counter+it.segs
+	err := tio.SegmentSuperpacket(it.pkt, func(seg []byte) error {
+		if reserved && c >= end {
+			// More segments than counters reserved for them, never encrypt past the reservation
+			return errFlowSegmentOverrun
+		}
+		if it.relayHI == nil {
+			// header + plaintext + AEAD tag (16 bytes for both AES-GCM and ChaCha20-Poly1305)
+			if out := f.encryptTxSegment(it, c, seg, nb, sb.Reserve(header.Len+len(seg)+16)); out != nil {
+				sb.Commit(out, it.remote)
+			}
+		} else {
+			f.encryptTxSegmentVia(it, c, rc, seg, nb, sb)
+		}
+		if reserved {
+			c++
+			rc++
+		}
+		return nil
+	})
+	if err != nil {
+		it.hostinfo.logger(f.l).Error("Failed to segment superpacket for send", "error", err)
+	}
+}
+
+// encryptTxSegment encrypts one segment into out with counter c, or with the next counter of the tunnel when c is
+// zero. It returns nil when encryption failed, which is already logged.
+func (f *Interface) encryptTxSegment(it *txItem, c uint64, seg, nb, out []byte) []byte {
+	if c != 0 {
+		// Reserved counters are unique, so the encryption needs no lock
+		return f.encryptMessage(header.Message, 0, it.hostinfo, it.ci, c, seg, nb, out, it.remote)
+	}
+	if noiseutil.EncryptLockNeeded {
+		it.ci.writeLock.Lock()
+	}
+	c = it.ci.messageCounter.Add(1)
+	out = f.encryptMessage(header.Message, 0, it.hostinfo, it.ci, c, seg, nb, out, it.remote)
+	if noiseutil.EncryptLockNeeded {
+		it.ci.writeLock.Unlock()
+	}
+	return out
+}
+
+// encryptTxSegmentVia encrypts one segment for the relayed tunnel, then wraps it for the relay with counter rc, both
+// allocated on demand when zero like encryptTxSegment does.
+func (f *Interface) encryptTxSegmentVia(it *txItem, c, rc uint64, seg, nb []byte, sb *batch.SendBatch) {
+	//relay header + header + plaintext + AEAD tag + relay tag
+	scratch := sb.Reserve(len(seg) + header.MaxOverhead)
+	out := f.encryptTxSegment(it, c, seg, nb, scratch[header.Len:])
+	if out == nil {
+		return
+	}
+	toSend, err := f.prepareSendVia(it.relayHI, it.relay, out, nb, scratch, true, rc)
+	if err != nil {
+		return
+	}
+	remote := it.remote
+	if rc == 0 {
+		// Unreserved sends follow the relay's current remote, like sendNoMetrics
+		remote = it.relayHI.GetRemote()
+	}
+	sb.Commit(toSend, remote)
+}
+
+// segmentCount reports how many segments tio.SegmentSuperpacket will cut pkt into, without touching its bytes.
+func segmentCount(pkt tio.Packet) uint64 {
+	if !pkt.GSO.IsSuperpacket() {
+		return 1
+	}
+	n := (len(pkt.Bytes) - int(pkt.GSO.HdrLen) + int(pkt.GSO.Size) - 1) / int(pkt.GSO.Size)
+	return uint64(max(n, 1))
+}
+
+// txTunnel accumulates one tunnel's counter needs across a batch so the batch reserves them with one atomic add.
+type txTunnel struct {
+	ci        *ConnectionState
+	hostinfo  *HostInfo
+	n         uint64
+	base      uint64
+	exhausted bool
+}
+
+// txReserver turns one tun read into txItems with message counters reserved for them. Only the reader goroutine
+// collects and reserves counters.
+type txReserver struct {
+	f       *Interface
+	tunnels []txTunnel
+}
+
+// collect runs the filter over pkts and appends the sendable ones to items, tallying each tunnel's counter needs.
+func (r *txReserver) collect(items []txItem, pkts []tio.Packet, fwPacket *firewall.ParsedPacket, rejectBuf []byte, q int, localCache firewall.ConntrackCache) []txItem {
+	f := r.f
+	r.tunnels = r.tunnels[:0]
+	for _, pkt := range pkts {
+		hostinfo := f.filterInsidePacket(pkt, fwPacket, rejectBuf, q, localCache)
+		if hostinfo == nil {
+			continue
+		}
+		var it txItem
+		if !f.prepareTxItem(&it, hostinfo, pkt) {
+			continue
+		}
+		it.segs = segmentCount(pkt)
+		it.counter = r.need(it.hostinfo, it.ci, it.segs)
+		if it.relayHI != nil {
+			// The relay bookkeeping inline sends do per segment in prepareSendVia, once per item
+			f.connectionManager.OutNoRebind(it.relayHI)
+			f.connectionManager.RelayUsed(it.relay.LocalIndex)
+			it.relayCounter = r.need(it.relayHI, it.relayHI.ConnectionState, it.segs)
+		}
+		items = append(items, it)
+	}
+	return items
+}
+
+// need adds n counters to ci's tally for this batch and returns the offset of the first one within it.
+func (r *txReserver) need(hostinfo *HostInfo, ci *ConnectionState, n uint64) uint64 {
+	t := r.tunnel(ci)
+	if t == nil {
+		r.tunnels = append(r.tunnels, txTunnel{ci: ci, hostinfo: hostinfo})
+		t = &r.tunnels[len(r.tunnels)-1]
+	}
+	off := t.n
+	t.n += n
+	return off
+}
+
+// tunnel finds ci's tally, a batch rarely touches more than a couple of tunnels so a scan beats a map.
+func (r *txReserver) tunnel(ci *ConnectionState) *txTunnel {
+	for i := range r.tunnels {
+		if r.tunnels[i].ci == ci {
+			return &r.tunnels[i]
+		}
+	}
+	return nil
+}
+
+// reserve takes the counters collect tallied and stamps them onto items, dropping the items of any tunnel whose
+// counter is exhausted. It returns the items that are left.
+func (r *txReserver) reserve(items []txItem) []txItem {
+	f := r.f
+	// One atomic add per tunnel reserves the whole batch's counters before any worker sees it, so counter order on a
+	// tunnel is dispatch order, which is the order the batches are flushed in. Reservation and the sender's order are
+	// both taken on the reader goroutine, so there is no window for another batch to slip in between.
+	for i := range r.tunnels {
+		t := &r.tunnels[i]
+		end := t.ci.messageCounter.Add(t.n)
+		t.base = end - t.n + 1
+		if end >= RejectAfterMessages {
+			t.ci.messageCounter.Store(RejectAfterMessages)
+			t.exhausted = true
+			f.messageMetrics.TxExhausted(int64(t.n))
+			if t.base <= RejectAfterMessages {
+				t.hostinfo.logger(f.l).Error("Dropping outbound packets, tunnel message counter is exhausted")
+			}
+		}
+	}
+
+	n := 0
+	for _, it := range items {
+		t := r.tunnel(it.ci)
+		if t.exhausted {
+			continue
+		}
+		it.counter += t.base
+		if it.relayHI != nil {
+			rt := r.tunnel(it.relayHI.ConnectionState)
+			if rt.exhausted {
+				continue
+			}
+			it.relayCounter += rt.base
+		}
+		items[n] = it
+		n++
+	}
+	clear(items[n:])
+	return items[:n]
+}

flow_tx_test.go

@@ -0,0 +1,44 @@
+//go:build linux && !android && !e2e_testing
+
+package nebula
+
+import (
+	"testing"
+
+	"github.com/flynn/noise"
+	"github.com/slackhq/nebula/firewall"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/overlay/tio"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReservedEncryptionMatchesInline encrypts a batch with counters reserved up front: the receiver accepts every
+// segment, in order.
+func TestReservedEncryptionMatchesInline(t *testing.T) {
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly)
+	packets, want := flowTestBatches(1, 3, 3, 9)
+	r := txReserver{f: sender.f}
+	items := r.collect(nil, packets[0], &firewall.ParsedPacket{}, make([]byte, mtu), 0, nil)
+	require.Equal(t, uint64(len(want)), r.tunnels[0].n)
+	sb := batch.NewSendBatch(sender.conn, 128, 65536)
+	for _, it := range r.reserve(items) {
+		sender.f.encryptTxItem(&it, make([]byte, 12), sb)
+	}
+	sender.f.flushSendBatch(sb, 0)
+	require.Equal(t, want, decryptWire(t, receiver, sender.conn.captured()))
+}
+
+// TestSegmentCount checks segmentCount predicts exactly the segments tio.SegmentSuperpacket cuts, odd tails included.
+func TestSegmentCount(t *testing.T) {
+	require.Equal(t, uint64(1), segmentCount(tio.Packet{Bytes: make([]byte, 60)}))
+	for _, segs := range []int{1, 2, 3, 45} {
+		for _, tail := range []int{0, 1, flowTestMSS - 1} {
+			pkt := flowTCPSuperpacket(0, 0, segs)
+			pkt.Bytes = pkt.Bytes[:len(pkt.Bytes)-tail]
+			pkt.GSO = tio.GSOInfo{Size: flowTestMSS, HdrLen: 40, CsumStart: 20, Proto: tio.GSOProtoTCP}
+			want, got := segmentCount(pkt), uint64(0)
+			require.NoError(t, tio.SegmentSuperpacket(pkt, func([]byte) error { got++; return nil }))
+			require.Equal(t, want, got, "segs %d tail %d", segs, tail)
+		}
+	}
+}

inside.go

@@ -133,31 +133,6 @@ func (f *Interface) filterInsidePacket(pkt tio.Packet, fwPacket *firewall.Parsed
 	return nil
 }
 
-func (f *Interface) sendInsideEncrypt(hostinfo *HostInfo, ci *ConnectionState, seg, scratch, nb []byte) []byte {
-	if noiseutil.EncryptLockNeeded {
-		ci.writeLock.Lock()
-	}
-	c := ci.messageCounter.Add(1)
-
-	out := header.Encode(scratch, header.Version, header.Message, 0, hostinfo.remoteIndexId, c)
-
-	out, encErr := ci.eKey.EncryptDanger(out, out, seg, c, nb)
-	if noiseutil.EncryptLockNeeded {
-		ci.writeLock.Unlock()
-	}
-	if encErr != nil {
-		hostinfo.logger(f.l).Error("Failed to encrypt outgoing packet",
-			"error", encErr,
-			"udpAddr", hostinfo.GetRemote(),
-			"counter", c,
-		)
-		// Skip this segment; the rest of the superpacket can still go out. TCP will retransmit anything we drop here.
-		return nil
-	}
-
-	return out
-}
-
 // sendInsideMessage encrypts a firewall-approved inside packet (or every
 // segment of a TSO/USO superpacket) into the caller's batch slot for
 // later sendmmsg flush. Segmentation is fused with encryption here so the
@@ -165,88 +140,46 @@ func (f *Interface) sendInsideEncrypt(hostinfo *HostInfo, ci *ConnectionState, s
 // scratch arena: SegmentSuperpacket builds each segment's plaintext in
 // segScratch[:segLen] in turn, and we encrypt directly into a fresh SendBatch slot.
 func (f *Interface) sendInsideMessage(hostinfo *HostInfo, pkt tio.Packet, nb []byte, sendBatch *batch.SendBatch) {
-	ci := hostinfo.ConnectionState
-	if ci.eKey == nil {
-		return
+	var it txItem
+	if f.prepareTxItem(&it, hostinfo, pkt) {
+		f.encryptTxItem(&it, nb, sendBatch)
 	}
+}
 
-	// One traffic-out mark covers every segment of the superpacket; doing it
-	// per segment in sendInsideEncrypt paid an atomic store up to ~45 extra
-	// times per TSO packet, inside writeLock under boring crypto.
-	//
-	// We rebound since this tunnel last sent, ask the lighthouse to get the far side punching at us again
+// recordOutbound marks traffic out on hostinfo, once per packet or superpacket, and when we rebound since this tunnel
+// last sent, asks the lighthouse to get the far side punching at us again.
+func (f *Interface) recordOutbound(hostinfo *HostInfo) {
 	if f.connectionManager.Out(hostinfo) {
 		f.lightHouse.QueryServer(hostinfo.vpnAddrs[0])
 		if f.l.Enabled(context.Background(), slog.LevelDebug) {
-			hostinfo.logger(f.l).Debug("Lighthouse update triggered for punch due to rebind epoch",
-				"vpnAddrs", hostinfo.vpnAddrs,
-			)
+			hostinfo.logger(f.l).Debug("Lighthouse update triggered for punch due to rebind epoch", "vpnAddrs", hostinfo.vpnAddrs)
 		}
 	}
+}
 
-	remote := hostinfo.GetRemote()
-	if !remote.IsValid() { //the relay path
-		//first, find our relay hostinfo:
-		var relayHostInfo *HostInfo
-		var relay *Relay
-		var err error
-		for _, relayIP := range hostinfo.relayState.CopyRelayIps() {
-			relayHostInfo, relay, err = f.hostMap.QueryVpnAddrsRelayFor(hostinfo.vpnAddrs, relayIP)
-			if err != nil {
-				hostinfo.relayState.DeleteRelay(relayIP)
-				hostinfo.logger(f.l).Info("sendNoMetrics failed to find HostInfo",
-					"relay", relayIP,
-					"error", err,
-				)
-				continue
-			}
-			break
+// findRelay returns the first relay that can reach hostinfo, forgetting the ones that no longer resolve.
+func (f *Interface) findRelay(hostinfo *HostInfo) (*HostInfo, *Relay) {
+	for _, relayIP := range hostinfo.relayState.CopyRelayIps() {
+		via, relay, err := f.hostMap.QueryVpnAddrsRelayFor(hostinfo.vpnAddrs, relayIP)
+		if err == nil {
+			return via, relay
 		}
-		if relayHostInfo == nil || relay == nil {
-			//failure already logged
-			return
-		}
-
-		err = tio.SegmentSuperpacket(pkt, func(seg []byte) error {
-			//relay header + header + plaintext + AEAD tag + relay tag
-			scratch := sendBatch.Reserve(len(seg) + header.MaxOverhead)
-
-			innerPacket := f.sendInsideEncrypt(hostinfo, ci, seg, scratch[header.Len:], nb)
-			if innerPacket == nil {
-				return nil
-			}
-
-			//now we need to do a relay-encrypt:
-			toSend, err := f.prepareSendVia(relayHostInfo, relay, innerPacket, nb, scratch, true)
-			if err != nil {
-				//already logged
-				return nil
-			}
-
-			sendBatch.Commit(toSend, relayHostInfo.GetRemote())
-			return nil
-		})
-		if err != nil {
-			hostinfo.logger(f.l).Error("Failed to segment superpacket for relay send", "error", err)
-		}
-		return
+		hostinfo.relayState.DeleteRelay(relayIP)
+		hostinfo.logger(f.l).Info("sendNoMetrics failed to find HostInfo", "relay", relayIP, "error", err)
 	}
+	return nil, nil
+}
 
-	err := tio.SegmentSuperpacket(pkt, func(seg []byte) error {
-		// header + plaintext + AEAD tag (16 bytes for both AES-GCM and ChaCha20-Poly1305)
-		scratch := sendBatch.Reserve(header.Len + len(seg) + 16)
-
-		out := f.sendInsideEncrypt(hostinfo, ci, seg, scratch, nb)
-		if out == nil {
-			return nil
-		}
-
-		sendBatch.Commit(out, remote)
-		return nil
-	})
+// encryptMessage encodes the header for counter c and encrypts p behind it into out. The caller allocates c and
+// holds writeLock when the cipher requires encryption in nonce order.
+func (f *Interface) encryptMessage(t header.MessageType, st header.MessageSubType, hostinfo *HostInfo, ci *ConnectionState, c uint64, p, nb, out []byte, remote netip.AddrPort) []byte {
+	out = header.Encode(out, header.Version, t, st, hostinfo.remoteIndexId, c)
+	out, err := ci.eKey.EncryptDanger(out, out, p, c, nb)
 	if err != nil {
-		hostinfo.logger(f.l).Error("Failed to segment superpacket for send", "error", err)
+		hostinfo.logger(f.l).Error("Failed to encrypt outgoing packet", "error", err, "udpAddr", remote, "counter", c)
+		return nil
 	}
+	return out
 }
 
 func (f *Interface) rejectInside(packet []byte, out []byte, q int) {
@@ -455,27 +388,33 @@ func (f *Interface) prepareSendVia(via *HostInfo,
 	nb,
 	out []byte,
 	nocopy bool,
+	c uint64,
 ) ([]byte, error) {
-	if noiseutil.EncryptLockNeeded {
+	// A nonzero c was reserved by the caller, who also did the connection manager bookkeeping
+	reserved := c != 0
+	locked := !reserved && noiseutil.EncryptLockNeeded
+	if locked {
 		// NOTE: for goboring AESGCMTLS we need to lock because of the nonce check
 		via.ConnectionState.writeLock.Lock()
 	}
-	c, ok := via.ConnectionState.NextMessageCounter()
-	if !ok {
-		if noiseutil.EncryptLockNeeded {
-			via.ConnectionState.writeLock.Unlock()
+	if !reserved {
+		var ok bool
+		c, ok = via.ConnectionState.NextMessageCounter()
+		if !ok {
+			if locked {
+				via.ConnectionState.writeLock.Unlock()
+			}
+			f.dropExhausted(via, c, "Dropping outbound relay packets, tunnel message counter is exhausted")
+			return nil, fmt.Errorf("tunnel message counter is exhausted")
 		}
-		f.dropExhausted(via, c, "Dropping outbound relay packets, tunnel message counter is exhausted")
-		return nil, fmt.Errorf("tunnel message counter is exhausted")
+		f.connectionManager.OutNoRebind(via)
 	}
-
 	out = header.Encode(out, header.Version, header.Message, header.MessageRelay, relay.RemoteIndex, c)
-	f.connectionManager.OutNoRebind(via)
 
 	// Authenticate the header and payload, but do not encrypt for this message type.
 	// The payload consists of the inner, unencrypted Nebula header, as well as the end-to-end encrypted payload.
 	if len(out)+len(ad)+via.ConnectionState.eKey.Overhead() > cap(out) {
-		if noiseutil.EncryptLockNeeded {
+		if locked {
 			via.ConnectionState.writeLock.Unlock()
 		}
 		via.logger(f.l).Error("SendVia out buffer not large enough for relay",
@@ -499,14 +438,16 @@ func (f *Interface) prepareSendVia(via *HostInfo,
 
 	var err error
 	out, err = via.ConnectionState.eKey.EncryptDanger(out, out, nil, c, nb)
-	if noiseutil.EncryptLockNeeded {
+	if locked {
 		via.ConnectionState.writeLock.Unlock()
 	}
 	if err != nil {
 		via.logger(f.l).Info("Failed to EncryptDanger in sendVia", "error", err)
 		return nil, err
 	}
-	f.connectionManager.RelayUsed(relay.LocalIndex)
+	if !reserved {
+		f.connectionManager.RelayUsed(relay.LocalIndex)
+	}
 	return out, nil
 }
 
@@ -519,7 +460,7 @@ func (f *Interface) prepareSendVia(via *HostInfo,
 // out is a buffer used to store the result of the Encrypt operation
 // q indicates which writer to use to send the packet.
 func (f *Interface) SendVia(via *HostInfo, relay *Relay, ad, nb, out []byte, nocopy bool, q int) {
-	toSend, err := f.prepareSendVia(via, relay, ad, nb, out, nocopy)
+	toSend, err := f.prepareSendVia(via, relay, ad, nb, out, nocopy, 0)
 	if err != nil {
 		// already logged by prepareSendVia
 		return
@@ -573,34 +514,19 @@ func (f *Interface) sendNoMetrics(t header.MessageType, st header.MessageSubType
 		return
 	}
 
-	//l.WithField("trace", string(debug.Stack())).Error("out Header ", &Header{Version, t, st, 0, hostinfo.remoteIndexId, c}, p)
-	out = header.Encode(out, header.Version, t, st, hostinfo.remoteIndexId, c)
-	// A closing tunnel is torn down right after this, so skip the connection manager entirely: no point recording
-	// traffic or asking the lighthouse for a punch. Otherwise, if we rebound since this tunnel last sent, ask the
-	// lighthouse to get the far side punching at us again.
-	if t != header.CloseTunnel && f.connectionManager.Out(hostinfo) {
-		f.lightHouse.QueryServer(hostinfo.vpnAddrs[0])
-		if f.l.Enabled(context.Background(), slog.LevelDebug) {
-			f.l.Debug("Lighthouse update triggered for punch due to rebind epoch",
-				"vpnAddrs", hostinfo.vpnAddrs,
-			)
-		}
+	// A closing tunnel needs neither traffic tracking nor another punch.
+	if t != header.CloseTunnel {
+		f.recordOutbound(hostinfo)
 	}
-
-	var err error
-	out, err = ci.eKey.EncryptDanger(out, out, p, c, nb)
+	out = f.encryptMessage(t, st, hostinfo, ci, c, p, nb, out, remote)
 	if noiseutil.EncryptLockNeeded {
 		ci.writeLock.Unlock()
 	}
-	if err != nil {
-		hostinfo.logger(f.l).Error("Failed to encrypt outgoing packet",
-			"error", err,
-			"udpAddr", remote,
-			"counter", c,
-		)
+	if out == nil {
 		return
 	}
 
+	var err error
 	if remote.IsValid() {
 		err = f.writers[q].WriteTo(out, remote)
 		if err != nil {
@@ -622,19 +548,8 @@ func (f *Interface) sendNoMetrics(t header.MessageType, st header.MessageSubType
 			}
 		}
 	} else {
-		// Try to send via a relay
-		for _, relayIP := range hostinfo.relayState.CopyRelayIps() {
-			relayHostInfo, relay, err := f.hostMap.QueryVpnAddrsRelayFor(hostinfo.vpnAddrs, relayIP)
-			if err != nil {
-				hostinfo.relayState.DeleteRelay(relayIP)
-				hostinfo.logger(f.l).Info("sendNoMetrics failed to find HostInfo",
-					"relay", relayIP,
-					"error", err,
-				)
-				continue
-			}
-			f.SendVia(relayHostInfo, relay, out, nb, fullOut[:header.Len+len(out)], true, q)
-			break
+		if via, relay := f.findRelay(hostinfo); via != nil && relay != nil {
+			f.SendVia(via, relay, out, nb, fullOut[:header.Len+len(out)], true, q)
 		}
 	}
 }