Skip to content

tp: extract IncrementalState to fix CustomState UAF#5593

Merged
LalitMaganti merged 3 commits into
mainfrom
dev/lalitm/tp-incremental-state-fix-uaf
Apr 29, 2026
Merged

tp: extract IncrementalState to fix CustomState UAF#5593
LalitMaganti merged 3 commits into
mainfrom
dev/lalitm/tp-incremental-state-fix-uaf

Conversation

@LalitMaganti
Copy link
Copy Markdown
Member

@LalitMaganti LalitMaganti commented Apr 25, 2026

Previously, each CustomState held a raw generation_ back-pointer that was re-pointed every time a new generation was created via OnNewTracePacketDefaults. The TraceSorter can hold a RefPtr to an older Generation G1 while a newer G2 is dropped after SEQ_INCREMENTAL_STATE_CLEARED. G1 stays alive but its generation_ then dangles to freed memory.

Fix this by by refactoring the generation/CustomState relationship so that lifetime is correct by construction:

  • New IncrementalState : RefCounted. Owns the per-incremental-state-interval data. A new IncrementalState is constructed only on SEQ_INCREMENTAL_STATE_CLEARED.

  • PacketSequenceStateGeneration becomes a thin per-trace_packet_defaults snapshot. All interned-data / custom-state / thread-descriptor accessors are forwarders to the IncrementalState.

  • CustomState is no longer RefCounted (it's owned uniquely by its IncrementalState). Its back-pointer is IncrementalState* and is set exactly once at lazy-allocation time inside IncrementalState::GetCustomState<T>. Because the IncrementalState owns the CustomState, the pointer is stable for the entire life of the CustomState making the UAF impossible.

  • OnNewTracePacketDefaultsnow constructs a new Generation that shares the same RefPtr<IncrementalState>. This removes a which is O(map-size).

  • OnPacketLoss walks the IncrementalState's CustomState array and clears any slot whose ClearOnPacketLoss() opted in .


Stack:

@github-actions
Copy link
Copy Markdown

github-actions Bot commented Apr 25, 2026

@LalitMaganti LalitMaganti force-pushed the dev/lalitm/tp-incremental-state-fix-uaf branch from ac3581a to 787b69b Compare April 25, 2026 17:04
LalitMaganti added a commit that referenced this pull request Apr 28, 2026
`PacketSequenceStateGeneration::OnPacketLoss()` previously mutated
`this` (flipping `is_incremental_state_valid_` and the track-event delta
state's `timestamps_valid_`) and returned `RefPtr(this)`. The
TraceSorter holds RefPtrs to the same generation for buffered packets
that were tokenized while the sequence was valid, so the in-place flip
was retroactively visible from those RefPtrs. No current reader observes
the mutation through a buffered RefPtr, but that is incidental —
anything added in the future that consults the generation's validity
from the parser side would see a stale, retroactively-invalidated view.

Make `OnPacketLoss()` return a new generation instead, sharing the
per-interval state (interned data, custom states, persistent thread
descriptor) with the previous one but with the validity bits cleared and
delta-encoded track-event state reset. Buffered packets retain the view
they had at tokenization time; the Builder swaps its current generation
to the new instance so subsequent tokenizations see the cleared state.

To support this, replace `TrackEventSequenceState::OnPacketLoss()`
(which mutated `timestamps_valid_` in place) with a value-returning form
that produces a fresh successor with persistent state preserved and
delta state reset. Incremental counter values and running reference
timestamps cannot be carried across the lost run of packets, so they are
dropped along with `timestamps_valid_`.

The full PSSG constructor's `trace_packet_defaults` parameter changes
from `TraceBlobView` to `std::optional<InternedMessageView>` so
OnPacketLoss can pass the existing defaults through without having to
reach inside the optional and reconstruct the view.

---

**Stack:**
- **#5588 — tp: stop mutating PacketSequenceStateGeneration on packet
loss** (this PR)
- #5590 — tp: split TrackEventSequenceState into descriptor + delta
state
- #5593 — tp: extract IncrementalState to fix CustomState UAF
LalitMaganti added a commit that referenced this pull request Apr 28, 2026
Refactor the per-sequence track-event state in two pieces with distinct
lifetimes:

- `TrackEventThreadDescriptor`: persistent pid/tid established by a
ThreadDescriptor packet. Held as a direct member of
`PacketSequenceStateGeneration` so it survives every transition,
including `SEQ_INCREMENTAL_STATE_CLEARED`. Real-world traces don't
always re-emit ThreadDescriptor after an incremental-state clear, so the
previously-known pid/tid must persist.

- `TrackEventSequenceState`: ephemeral delta state (running reference
timestamps + incremental counter values). Now a `CustomState`, so it is
shared across `trace_packet_defaults` transitions instead of being reset
on every defaults change. This is more permissive: a producer that emits
a thread descriptor, then changes defaults without changing the
timestamp clock, now gets correct delta calculations rather than a
`timestamps_valid_=false` skip.

All forwarder methods on `PacketSequenceStateGeneration`
(`IncrementAndGet*`, `track_event_timestamps_valid`,
`SetThreadDescriptor`) are removed. Callers now go through
`GetCustomState<>()` and `state.thread_descriptor()` directly, shrinking
the public surface of Generation. The `CustomState` base loses its leaky
pid/tid helpers in favour of a `thread_descriptor()` forwarder.

`CustomState` gains a virtual `ClearOnPacketLoss()` hook (default
false). `TrackEventSequenceState` overrides to return true so its
delta-encoded state is cleared on packet loss. The dependency from
generic generation code to specific CustomState subclasses is replaced
by this generic per-CustomState policy.

This is preparation for the upcoming `IncrementalState` extraction
(#5593) that fixes the `CustomState` back-pointer UAF.

---

**Stack:**
- #5588 — tp: stop mutating PacketSequenceStateGeneration on packet loss
- **#5590 — tp: split TrackEventSequenceState into descriptor + delta
state** (this PR)
- #5593 — tp: extract IncrementalState to fix CustomState UAF
Base automatically changed from dev/lalitm/tp-track-event-state-customstate to main April 28, 2026 15:46
Fixes the use-after-free flagged in the security report against
`PacketSequenceStateGeneration::CustomState`. Previously, each
CustomState held a raw `generation_` back-pointer that was re-pointed
every time a new generation was created via `OnNewTracePacketDefaults`
(see the now-removed `set_generation(this)` loop). The TraceSorter can
hold a RefPtr to an older Generation G1 while a newer G2 — to which
G1's shared CustomState had been re-pointed — is dropped after
SEQ_INCREMENTAL_STATE_CLEARED. G1 stays alive (and so does the
CustomState), but its `generation_` then dangles to freed memory; any
subsequent lookup through the CustomState dereferences the freed G2.

Refactor the generation/CustomState relationship so that lifetime is
correct by construction:

- New `IncrementalState : RefCounted`. Owns the per-incremental-state
  interval data: `interned_data_`, the array of CustomStates (now
  `unique_ptr` rather than `RefPtr`), and the persistent thread
  descriptor. A new IncrementalState is constructed only on
  SEQ_INCREMENTAL_STATE_CLEARED.

- `PacketSequenceStateGeneration` becomes a thin per-trace-packet-
  defaults snapshot: holds `RefPtr<IncrementalState>` (shared with
  sibling generations within the same interval), the defaults blob,
  and the validity flag. All interned-data / custom-state / thread-
  descriptor accessors are forwarders to the IncrementalState.

- `CustomState` is no longer RefCounted (it's owned uniquely by its
  IncrementalState). Its back-pointer is `IncrementalState*` and is
  set exactly once at lazy-allocation time inside
  `IncrementalState::GetCustomState<T>`. Because the IncrementalState
  owns the CustomState, the pointer is stable for the entire life of
  the CustomState — UAF impossible.

- `OnNewTracePacketDefaults` no longer copies the `InternedFieldMap`
  or re-points CustomStates; it constructs a new Generation that
  shares the same `RefPtr<IncrementalState>`. This removes a
  per-defaults-change copy that could be O(map-size) in the hot path.

- `OnPacketLoss` walks the IncrementalState's CustomState array and
  clears any slot whose `ClearOnPacketLoss()` opted in (TES today),
  then returns a new Generation referencing the same IncrementalState
  with the validity bit cleared.

Adds a regression test mirroring the report's scenario:
G1 → defaults → G2 → SEQ_INCREMENTAL_STATE_CLEARED → G3, then drop G2
while G1 remains pinned (as the sorter would). Pre-fix, accessing the
CustomState through G1 would UAF; the test exercises that path.
@LalitMaganti LalitMaganti force-pushed the dev/lalitm/tp-incremental-state-fix-uaf branch from b84a006 to fce67a8 Compare April 28, 2026 21:32
@LalitMaganti LalitMaganti merged commit 3578f39 into main Apr 29, 2026
23 checks passed
@LalitMaganti LalitMaganti deleted the dev/lalitm/tp-incremental-state-fix-uaf branch April 29, 2026 16:58
LalitMaganti added a commit that referenced this pull request May 6, 2026
tp: scope gpu_counter custom-groups cache to per-sequence state (#5742)
tp: extract IncrementalState to fix CustomState UAF (#5593)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants