fix(l3): a covered last fragment must complete; say what the timeout really does

Two review findings on the fragment reassembler.

1. A whole datagram could vanish. entry.total was assigned before addRange
   was asked, so a last fragment (MF=0) whose range was already covered by
   MF=1 fragments answered fragInsertDuplicate and returned nil — while the
   entry was already complete(). Nothing re-examined it, because every later
   fragment is a duplicate too, so it died at its deadline with all its bytes
   present. A duplicate now falls through to the completion check: it
   contributes no bytes (held bytes still win) but it does contribute the
   total length. This is what the documented first-wins policy always
   implied; the code just did not do it.

   The sender needed is non-conforming, so the old behaviour was safe rather
   than exploitable — but it contradicted the comment three screens up, and
   that comment is the next reader's only defence.

   Also closed positively: a last fragment declaring an end BELOW the bytes
   already held now poisons the datagram instead of quietly never completing.

2. The 5 s timeout was not a memory ceiling and the comment said it was.
   sweep ran only when a NEW key was created, so once fragmented traffic
   stopped, up to fragMaxEntries entries stayed resident indefinitely.

   Both halves are fixed, and the honest one is the comment. sweep now runs
   on EVERY fragment — an O(64) scan on a path that is already the rare one —
   which releases residue as soon as any fragment arrives instead of waiting
   for an unrelated new datagram. That still does not cover total silence, so
   fragTimeout now documents the guarantee the code actually keeps: bounded
   by fragMaxEntries/fragMaxTotalBytes at all times, released on the next
   fragment, NOT "freed within 5 s".

   No timer, deliberately: it would need a goroutine with a lifecycle tied to
   something returnDeviceWrapper has no teardown hook for, and a goroutine
   that must be stopped and might not be is a failure this project has
   already paid for — to reclaim at most ~1.1 MiB that only exists after
   fragmented traffic has already happened. What bounds growth is the byte
   and entry ceiling; this timeout's job is correctness, and for that a
   check driven by the arriving fragment is exact.

   The now-unreachable per-key deadline check is removed rather than left as
   dead defence in depth.

16 mutations, all red. M15 (duplicate returns early again) reds only the
buggy case while the control and the poison case stay green, so the test is
shown able to see both an assembled datagram and a lost one. M17 (sweep back
inside the new-key branch) reds the new test while both old timeout subtests
stay green.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BHw89tdWddzhjUc4bAH4tS

transport/wireguard/frag_reassembly.go

@@ -57,6 +57,10 @@ import (
 // i.e. ~1.1 MiB with an adversary trying his hardest, and in practice a few
 // kilobytes: a fragmented datagram lives for the microseconds between its
 // fragments arriving in the same read batch.
+//
+// These are the ceilings that hold WHILE traffic flows and also AFTER it stops:
+// expiry is traffic-driven, not timed, so the residue after the last fragment
+// is bounded by these numbers rather than by fragTimeout. Read fragTimeout.
 const (
 	// fragMaxEntries is the maximum number of datagrams under reassembly at
 	// once. A 65th datagram evicts the oldest one rather than growing the map.
@@ -82,9 +86,33 @@ const (
 	// purpose. Exceeding it poisons the entry (see fragEntry.broken).
 	fragMaxRanges = 64
 
-	// fragTimeout is how long an incomplete datagram may hold memory. The timer
-	// starts at the FIRST fragment and is never refreshed, so a sender cannot
-	// keep an entry alive by dribbling fragments into it.
+	// fragTimeout is how long an incomplete datagram stays USABLE. Past it the
+	// entry is refused and released. The clock starts at the FIRST fragment and
+	// is never refreshed, so a sender cannot keep an entry alive by dribbling
+	// fragments into it.
+	//
+	// READ THIS BEFORE TREATING IT AS A MEMORY GUARANTEE. There is no timer.
+	// Expiry is driven by traffic: the sweep runs inside reassemble, so an
+	// expired entry is released ON THE NEXT FRAGMENT that reaches this endpoint,
+	// not five seconds from now. When fragmented traffic stops completely, up to
+	// fragMaxEntries expired entries (bounded by the ceilings above, so ~1.1 MiB
+	// worst case, and in practice a few KiB) stay resident until it resumes or
+	// the endpoint is torn down.
+	//
+	// That is a deliberate trade, not an oversight:
+	//
+	//   - What bounds growth is fragMaxEntries/fragMaxTotalBytes, not this
+	//     value. The residue is capped no matter how long the silence lasts.
+	//   - This timeout's real job is CORRECTNESS — never assembling a datagram
+	//     from a stale head — and for that a check driven by the arriving
+	//     fragment is exact, because the arriving fragment is the only thing
+	//     that could make a stale entry matter.
+	//   - The alternative is a goroutine per endpoint with a lifecycle tied to
+	//     something returnDeviceWrapper has no teardown hook for today. A
+	//     goroutine that must be stopped and might not be is a failure this
+	//     project has already paid for; waking one forever to reclaim at most
+	//     ~1.1 MiB that only exists after fragmented traffic has already
+	//     happened is the wrong side of the trade on a 128 MB router.
 	//
 	// Linux uses 30 s (net.ipv4.ipfrag_time). 5 s is chosen instead because the
 	// real inter-fragment gap on this path is sub-millisecond (the peer emits
@@ -424,8 +452,11 @@ func validFragmentExtent(start int, length int, more bool) bool {
 // headroom and as the wireguard-go write offset.
 //
 // It returns the completed datagram — a freshly allocated, fully owned buffer
-// laid out as prefix bytes of headroom followed by the packet — or nil when
-// the datagram is still incomplete, was a duplicate, or was refused.
+// laid out as prefix bytes of headroom followed by the packet — or nil when the
+// datagram is still incomplete or was refused. A duplicate fragment is not
+// automatically nil: it contributes no bytes, but a last fragment whose range
+// is already held still tells us the total length and can complete the
+// datagram.
 func (r *fragmentReassembler) reassemble(packet []byte, prefix int) []byte {
 	info, ok := parseFragment(packet)
 	if !ok {
@@ -436,13 +467,16 @@ func (r *fragmentReassembler) reassemble(packet []byte, prefix int) []byte {
 	defer r.access.Unlock()
 
 	now := r.now()
+	// Sweep on EVERY fragment, not only when a new key appears. It is an
+	// O(fragMaxEntries) scan on a path that is already the rare one, and it is
+	// the single mechanism behind both halves of the timeout: an expired entry
+	// is unusable AND unallocated. See fragTimeout for what this does and does
+	// not guarantee — there is no timer, so "expired" means "released on the
+	// next fragment", not "released 5 s from now".
+	r.sweep(now)
+
 	entry := r.entries[info.key]
-	if entry != nil && now >= entry.deadline {
-		r.remove(entry)
-		entry = nil
-	}
 	if entry == nil {
-		r.sweep(now)
 		if r.entries == nil {
 			r.entries = make(map[fragKey]*fragEntry)
 		}
@@ -465,6 +499,11 @@ func (r *fragmentReassembler) reassemble(packet []byte, prefix int) []byte {
 			r.poison(entry)
 			return nil
 		}
+		if len(entry.ranges) > 0 && entry.ranges[len(entry.ranges)-1].end > end {
+			// Bytes are already held past the end this fragment declares.
+			r.poison(entry)
+			return nil
+		}
 		entry.total = end
 	} else if entry.total >= 0 && end > entry.total {
 		// A fragment claiming bytes past the declared end.
@@ -472,19 +511,29 @@ func (r *fragmentReassembler) reassemble(packet []byte, prefix int) []byte {
 		return nil
 	}
 
+	// A duplicate contributes no bytes — the held ones win, see addRange — but
+	// it can still be the fragment that COMPLETES the datagram: a last fragment
+	// (MF=0) whose range is already covered contributes only the total length.
+	// So a duplicate falls through to the completion check instead of returning
+	// here. Returning early stranded such a datagram forever: every later
+	// fragment is a duplicate too, so nothing would re-examine the entry and it
+	// died at its deadline with all its bytes present.
+	duplicate := false
 	switch entry.addRange(info.start, end) {
 	case fragInsertDuplicate:
-		return nil
+		duplicate = true
 	case fragInsertConflict:
 		r.poison(entry)
 		return nil
 	}
 
-	if !r.grow(entry, end) {
-		r.poison(entry)
-		return nil
+	if !duplicate {
+		if !r.grow(entry, end) {
+			r.poison(entry)
+			return nil
+		}
+		copy(entry.payload[info.start:end], info.payload)
 	}
-	copy(entry.payload[info.start:end], info.payload)
 
 	if info.unfragmentable != nil && entry.header == nil {
 		stored := make([]byte, len(info.unfragmentable))

transport/wireguard/frag_reassembly_lx_test.go

@@ -586,6 +586,107 @@ func TestReturnFragmentDuplicateAccepted(t *testing.T) {
 	}
 }
 
+// TestReturnFragmentCoveredLastFragment pins the review finding that a whole
+// datagram could vanish: the last fragment (MF=0) carries the total length, and
+// when its range is already covered by fragments that arrived with MF=1, it adds
+// no bytes — but it is still the fragment that completes the datagram. Returning
+// early on "duplicate" stranded it forever, because every later fragment is a
+// duplicate too, so nothing ever re-examined the entry.
+//
+// The senders below are non-conforming, which is why the old behaviour was safe
+// rather than exploitable — but it contradicted the documented first-wins
+// policy, and the documentation is the next reader's only defence.
+//
+// Each subtest is paired with its own control so the instrument is shown able to
+// report BOTH outcomes: an assembled datagram and a lost one.
+func TestReturnFragmentCoveredLastFragment(t *testing.T) {
+	t.Parallel()
+
+	t.Run("a covered last fragment completes the datagram", func(t *testing.T) {
+		t.Parallel()
+		datagram := ipv4Datagram(61, fragTestPayload(240))
+		wrapper, _, returnPath, _ := newFragTestWrapper(nil)
+		// Everything, but wrongly flagged as "more to come"...
+		writePackets(t, wrapper, ipv4Fragment(datagram, 0, 240, true))
+		require.Empty(t, returnPath.payloads(), "MF=1 alone must not complete anything")
+		// ...then a last fragment that adds no new bytes, only the length.
+		writePackets(t, wrapper, ipv4Fragment(datagram, 160, 80, false))
+		delivered := returnPath.payloads()
+		require.Len(t, delivered, 1, "a datagram whose bytes were all present was never delivered")
+		require.Equal(t, datagram, delivered[0])
+		require.Empty(t, wrapper.reassembler.entries)
+		require.Zero(t, wrapper.reassembler.bytes)
+	})
+
+	t.Run("control: a covered last fragment does not complete a datagram with a hole", func(t *testing.T) {
+		t.Parallel()
+		datagram := ipv4Datagram(62, fragTestPayload(240))
+		wrapper, _, returnPath, _ := newFragTestWrapper(nil)
+		writePackets(t, wrapper, ipv4Fragment(datagram, 0, 80, true))
+		writePackets(t, wrapper, ipv4Fragment(datagram, 160, 80, true))
+		// Covered by the range above, and declares total = 240 — but [80,160)
+		// is missing, so nothing may be delivered.
+		writePackets(t, wrapper, ipv4Fragment(datagram, 160, 80, false))
+		require.Empty(t, returnPath.payloads(), "an incomplete datagram was delivered")
+
+		// Filling the hole delivers it: the same instrument says yes.
+		writePackets(t, wrapper, ipv4Fragment(datagram, 80, 80, true))
+		delivered := returnPath.payloads()
+		require.Len(t, delivered, 1)
+		require.Equal(t, datagram, delivered[0])
+	})
+
+	t.Run("bytes held past the declared end poison the datagram", func(t *testing.T) {
+		t.Parallel()
+		datagram := ipv4Datagram(63, fragTestPayload(240))
+		wrapper, _, returnPath, _ := newFragTestWrapper(nil)
+		writePackets(t, wrapper, ipv4Fragment(datagram, 0, 240, true))
+		// Declares the datagram ends at 160 while 240 bytes are already held.
+		// The offset must be non-zero: an IPv4 fragment at offset 0 with MF=0
+		// is not a fragment at all, so that shape can never reach the cache.
+		writePackets(t, wrapper, ipv4Fragment(datagram, 80, 80, false))
+		require.Zero(t, wrapper.reassembler.bytes, "the poisoned entry must release its buffers")
+		writePackets(t, wrapper, ipv4Fragment(datagram, 160, 80, false))
+		require.Empty(t, returnPath.payloads(), "a self-contradicting datagram was assembled")
+	})
+}
+
+// TestReturnFragmentSweepIsPerCall pins the second review finding: the sweep
+// must run on EVERY fragment, not only when a new key is created. Otherwise an
+// entry that expired while another key stayed alive keeps its memory until some
+// unrelated new datagram happens to appear.
+func TestReturnFragmentSweepIsPerCall(t *testing.T) {
+	t.Parallel()
+	stale := ipv4Datagram(71, fragTestPayload(3000))
+	staleFragments := ipv4Fragments(stale, 1000)
+	live := ipv4Datagram(72, fragTestPayload(3000))
+	liveFragments := ipv4Fragments(live, 1000)
+
+	wrapper, _, _, clock := newFragTestWrapper(nil)
+
+	writePackets(t, wrapper, staleFragments[0]) // deadline t+5s
+	staleBytes := wrapper.reassembler.bytes
+	require.Positive(t, staleBytes)
+
+	clock.advance(3 * time.Second)
+	writePackets(t, wrapper, liveFragments[0]) // deadline t+8s
+	require.Len(t, wrapper.reassembler.entries, 2)
+	liveBytes := wrapper.reassembler.bytes - staleBytes
+	require.Positive(t, liveBytes)
+
+	// t+6s: the stale entry has expired, the live one has not. This fragment
+	// creates NO new key, so a sweep that only runs on key creation never sees
+	// the expired entry.
+	clock.advance(3 * time.Second)
+	writePackets(t, wrapper, liveFragments[1])
+
+	require.Len(t, wrapper.reassembler.entries, 1, "the expired entry survived a call that created no new key")
+	for key := range wrapper.reassembler.entries {
+		require.Equal(t, uint32(72), key.id)
+	}
+	require.Equal(t, liveBytes, wrapper.reassembler.bytes, "the expired entry's bytes were never released")
+}
+
 // TestReturnFragmentEntryEviction pins requirement 1 (entry ceiling): the
 // (fragMaxEntries+1)-th datagram evicts the OLDEST, and the evicted one can no
 // longer be completed.