Handshake context registry leaks on aborted handshakes

[coderabbitai]

Summary

Severity: High
Affected commit: f77d939
Raised via: PR #323, requested by @leynos

An unauthenticated remote memory-leak / availability bug exists in the Wireframe handshake context handoff introduced in commit f77d939. Registry insertion happens before a fallible network write, whilst cleanup is only performed by app construction on the normal successful path.

---

Background

Commit f77d939 replaces thread-local handshake storage with a process-global Tokio task-ID-keyed HashMap used to bridge the asynchronous handshake callback and the synchronous app factory.

scope_current_context(Some(...)) inserts the connection context into this global registry before the scoped future is polled, but never removes the entry when the future returns an error. In the handshake success path that scoped future writes the 8-byte success reply and only later reaches the app-factory path that consumes the registry entry via take_current_context(). If a remote client sends a valid 12-byte TRTP preamble and then closes/resets the socket so write_handshake_reply(...).await? returns an I/O error, the connection setup aborts before the app factory can run, leaving the registry entry permanently retained. Repeating this pre-authentication flow can grow the global HashMap without bound and exhaust server memory.

The previous thread-local implementation did not create an unbounded per-task global map; stale state was at most per worker thread and overwritten.

---

Attack Path

1. Remote unauthenticated TCP client connects to the Wireframe listener on the configured bind address.
2. Client sends a valid 12-byte TRTP/version-1 preamble.
3. HotlinePreamble validation succeeds; the success callback fires.
4. scope_current_context inserts ConnectionContext into the global task-ID registry.
5. write_handshake_reply fails because the client reset/closed the connection.
6. The error aborts setup; the app factory is never reached.
7. take_current_context() is never called; the registry entry is retained.
8. Steps 1–7 are repeated without bound, exhausting process memory.

---

Evidence

src/wireframe/connection.rs

• Lines 116–130 — registry() introduces a process-global HashMap<Id, ConnectionContext>. Entries persist until explicitly removed.
• Lines 154–164 — scope_current_context stores the context in the global registry before awaiting the scoped future, but provides no finally/drop cleanup when that future returns an error.
• Lines 193–201 — take_current_context() is the only removal path; it is reached by the app factory on the normal setup path, not when handshake setup aborts before app construction.

src/wireframe/handshake.rs (lines 79–82)

scope_current_context is called with the cloned context before write_handshake_reply(...).await?. If the client disconnects or resets, the write fails, the function exits with an error, and the registry entry is never consumed.

src/server/wireframe/mod.rs (lines 190–200)

build_app_context calls take_current_context() to consume and remove the stored context. This function is only invoked after successful handshake setup; aborted handshakes do not reach this cleanup path.

---

Impact

Dimension	Assessment
Likelihood	High — requires only network reachability, a valid 12-byte TRTP preamble, and a client-controlled reset. Pre-authentication; default bind is all interfaces on TCP 5500.
Impact	High — unbounded process-global map growth; repeated unauthenticated connections exhaust daemon/host memory (DoS). No RCE, data disclosure, or authentication bypass.
Severity	High

---

Existing Controls

• Preamble validation requires the TRTP protocol ID and supported version before the vulnerable path.
• A 5-second preamble timeout limits idle sockets but does not clean entries inserted after a valid preamble.
• Authentication and authorisation are not reached before the vulnerable handshake path.
• Deployments using only the legacy runtime are not affected.

---

Suggested Remediation Direction

Ensure scope_current_context removes the registry entry unconditionally when the scoped future completes, whether with success or error. Introduce a drop-guard that calls take_registry_context() when it goes out of scope, so the guard fires on both the success and error paths.

---

Assumptions / Blindspots

• The external wireframe crate implementation was not inspected from source in this static pass; reachability relies on repository server construction and comments about preamble callbacks and app factory ordering.
• No deployment manifests confirm whether mxd-wireframe-server is exposed publicly or sits behind a private proxy.
• The exact number of aborted handshakes needed for host-level exhaustion depends on allocator behaviour, Tokio task-ID behaviour, memory limits, and operational connection controls outside the repository.

[leynos]

@CodeRabbit plan

[coderabbitai]

Coding Plan

Summary

• Remove pre-insertion from scope_current_context; the registry is now populated exclusively by explicit store_current_context calls after successful handshake writes.
• Failed handshakes never insert into the registry, eliminating the leak vector.
• Update one existing test to call store_current_context explicitly, and add a regression test confirming aborted handshakes leave no orphaned entries.

Design Choices

Design Choice 1: Registry insertion should only occur for successfully completed handshakes

Options Considered:

1. Add a RAII drop-guard that calls take_registry_context() unconditionally when the scope exits, then "defuse" it inside store_current_context to prevent double-removal on success.
2. Remove the pre-insertion from scope_current_context entirely; rely solely on the explicit store_current_context call that occurs after the handshake write succeeds.

Chosen Option: 2

Rationale: Removing pre-insertion is simpler, introduces no new types, and is semantically correct—the registry should represent only successfully completed handshakes, not attempted ones.

💡 User Tips

Regenerate the plan with different choices with @coderabbitai <feedback>.

Implementation Steps

Phase 1: Eliminate Pre-Insertion from scope_current_context

Remove the registry insertion from scope_current_context so that entries are only created by the explicit store_current_context call after a successful handshake write. This ensures failed handshakes never leave orphaned entries.

Task 1: Remove Pre-Insertion Logic

Remove the store_registry_context call from scope_current_context in src/wireframe/connection.rs.

• Delete the conditional block (lines 158–160) that calls store_registry_context(seed_context) when context is Some.
• Retain the CONNECTION_CONTEXT.scope(...) wrapper—the task-local slot is still required during the scoped future's execution.
• The function should simply scope the context into the task-local and await the inner future; registry population is now exclusively the responsibility of store_current_context.

Task 2: Update Module Documentation

Update the module-level docstring in src/wireframe/connection.rs to reflect the new semantics.

• Clarify that scope_current_context scopes the context into task-local storage only.
• Document that registry insertion requires an explicit store_current_context call after success, ensuring failed handshakes do not leave orphaned entries.
• Remove any mention of "pre-insertion" or "mirroring at scope start."

🤖 Prompt for AI agents

Patch `src/wireframe/connection.rs` to eliminate the pre-insertion of
`ConnectionContext` into the global registry from `scope_current_context`. The
goal is to ensure the registry is populated exclusively by explicit
`store_current_context` calls on the success path.

- In `scope_current_context`, delete the conditional block (lines 158–160) that
calls `store_registry_context(seed_context)` when `context` is `Some`.
- Retain the `CONNECTION_CONTEXT.scope(...)` wrapper unchanged—the task-local
slot must still be populated for in-scope access.
- The function body should only scope the context into the task-local and await
the inner future; no registry write should occur here.
- Update the module-level docstring to:
- State that `scope_current_context` scopes the context into task-local storage
only.
- State that registry insertion requires an explicit `store_current_context`
call after a successful operation.
  - Remove all references to "pre-insertion" or "mirroring at scope start."

Phase 2: Align Tests with New Semantics

Update existing unit tests that relied on pre-insertion behaviour, and add a regression test to verify aborted handshakes do not leak registry entries.

Task 1: Update `scopes_context_for_post_handshake_handoff` Test

Modify the test in src/wireframe/connection.rs to explicitly store the context rather than relying on pre-insertion.

• Change the empty async {} block to call store_current_context(scoped_context) before completing.
• The test now mirrors the real success-path behaviour: context is stored explicitly after any fallible work, then retrieved after the scope exits via take_current_context().

Task 2: Verify `reports_registry_entry_count` Test Remains Correct

Confirm the existing test continues to pass without modification.

• This test already calls store_current_context inside the scoped future, so it does not depend on pre-insertion.
• Run the test to verify it passes with the updated scope_current_context implementation.

Task 3: Add Regression Test for Aborted Handshake

Create a new test in src/wireframe/connection.rs that verifies no registry leak occurs when the scoped future fails.

• Spawn a task that calls scope_current_context(Some(context), async { Err(io::Error::...) }).
• After the future completes with an error, assert take_current_context() returns None.
• Assert registry_count_for_metadata reports 0 entries for the test's metadata.
• Mark the test #[serial] to avoid interference with other registry tests.

🤖 Prompt for AI agents

Update the unit tests in `src/wireframe/connection.rs` to align with the new
`scope_current_context` semantics and add a regression test for the
aborted-handshake leak scenario.

**Update `scopes_context_for_post_handshake_handoff`:**
- Replace the empty `async {}` block inside the scoped future with a call to
`store_current_context(scoped_context)` before the future resolves.
- Ensure the test still asserts that `take_current_context()` returns the
expected context after the scope exits.

**Verify `reports_registry_entry_count`:**
- Confirm this test already calls `store_current_context` inside its scoped
future and does not rely on pre-insertion.
- Run it against the updated implementation and confirm it passes without
modification.

**Add regression test `aborted_handshake_does_not_leak_registry_entry` (or
equivalent):**
- Mark the test `#[serial]` to prevent registry interference with other tests.
- Spawn a task that calls `scope_current_context(Some(context), async {
Err(io::Error::...) })`.
- After the future resolves with an error, assert `take_current_context()`
returns `None`.
- Assert `registry_count_for_metadata` reports 0 entries for the test's
metadata.

Research

The Wireframe handshake uses a dual-storage design in src/wireframe/connection.rs: a Tokio task-local slot for in-scope access and a global Mutex<HashMap<Id, ConnectionContext>> registry for post-scope hand-off to the synchronous app factory. scope_current_context currently pre-inserts into both stores before polling the scoped future. The success handler in src/wireframe/handshake.rs wraps the handshake reply write inside this scope and explicitly calls store_current_context after a successful write—making the pre-insertion redundant on the success path and the source of the leak on the failure path.

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

Patch `src/wireframe/connection.rs` to eliminate the pre-insertion of
`ConnectionContext` into the global registry from `scope_current_context`. The
goal is to ensure the registry is populated exclusively by explicit
`store_current_context` calls on the success path.

- In `scope_current_context`, delete the conditional block (lines 158–160) that
calls `store_registry_context(seed_context)` when `context` is `Some`.
- Retain the `CONNECTION_CONTEXT.scope(...)` wrapper unchanged—the task-local
slot must still be populated for in-scope access.
- The function body should only scope the context into the task-local and await
the inner future; no registry write should occur here.
- Update the module-level docstring to:
- State that `scope_current_context` scopes the context into task-local storage
only.
- State that registry insertion requires an explicit `store_current_context`
call after a successful operation.
  - Remove all references to "pre-insertion" or "mirroring at scope start."
===============================================================================

Task: 2

Update the unit tests in `src/wireframe/connection.rs` to align with the new
`scope_current_context` semantics and add a regression test for the
aborted-handshake leak scenario.

**Update `scopes_context_for_post_handshake_handoff`:**
- Replace the empty `async {}` block inside the scoped future with a call to
`store_current_context(scoped_context)` before the future resolves.
- Ensure the test still asserts that `take_current_context()` returns the
expected context after the scope exits.

**Verify `reports_registry_entry_count`:**
- Confirm this test already calls `store_current_context` inside its scoped
future and does not rely on pre-insertion.
- Run it against the updated implementation and confirm it passes without
modification.

**Add regression test `aborted_handshake_does_not_leak_registry_entry` (or
equivalent):**
- Mark the test `#[serial]` to prevent registry interference with other tests.
- Spawn a task that calls `scope_current_context(Some(context), async {
Err(io::Error::...) })`.
- After the future resolves with an error, assert `take_current_context()`
returns `None`.
- Assert `registry_count_for_metadata` reports 0 entries for the test's
metadata.

💡 Iterate on the plan with: `@coderabbitai` <feedback>

Example Feedback
- `@coderabbitai` You can skip phase 3. Add a simple unit test case for phase 2.
- `@coderabbitai` For design choice 1 go ahead with option 3 and replan.

---

💬 Have feedback or questions? Drop into our discord!