Security: completed client tasks retained in JoinSet until shutdown — unbounded memory DoS

[coderabbitai]

Summary

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

The new graceful-shutdown implementation stores every accepted client task in a Tokio JoinSet, but the accept loop never calls join_next, try_join_next, or otherwise reaps completed tasks whilst the server is running. Completed client tasks remain owned by the JoinSet until they are joined, and this code only joins them after a shutdown signal. An unauthenticated remote attacker can repeatedly open a TCP connection, receive the banner, then disconnect; each handler exits, but its task record is retained in the JoinSet. Over time this creates unbounded memory growth and can exhaust memory or degrade service availability.

The previous code used detached tokio::spawn tasks, which did not retain all completed task handles in the main accept loop.

---

Evidence

src/main.rs (lines 103–134)

let (shutdown_tx, _) = broadcast::channel(1);
let mut join_set = JoinSet::new();
let mut shutdown = Box::pin(shutdown_signal());

loop {
    tokio::select! {
        _ = &mut shutdown => {
            println!("shutdown signal received");
            break;
        }
        res = listener.accept() => {
            let (socket, peer) = res?;
            let pool = pool.clone();
            let mut rx = shutdown_tx.subscribe();
            join_set.spawn(async move {
                if let Err(e) = handle_client(socket, peer, pool, rx).await {
                    eprintln!("connection error: {}", e);
                }
            });
        }
    }
}

// notify all tasks to shut down
let _ = shutdown_tx.send(());

while let Some(res) = join_set.join_next().await {
    if let Err(e) = res {
        eprintln!("task error: {}", e);
    }
}

• The JoinSet is created to own all per-client tasks, but is never drained during normal server operation.
• Each accepted connection is inserted into the JoinSet. The accept loop has no branch that joins or removes tasks that have already completed.
• Completed tasks are only joined after shutdown, so task records from short-lived attacker connections accumulate for the entire uptime of the daemon.
• Client handlers can finish normally when a client disconnects (lines 149–178), making it cheap for an unauthenticated attacker to create many completed-but-unjoined tasks.

---

Attack Path

1. Unauthenticated TCP attacker connects to the mxd listener (default 0.0.0.0:5500).
2. listener.accept() succeeds; join_set.spawn() creates a new task.
3. handle_client sends the MXD banner immediately and enters its command loop.
4. Attacker closes/disconnects; handle_client exits normally via EOF.
5. The accept loop is never woken to reap the completed task; the JoinSet entry persists.
6. Steps 1–5 are repeated; retained completed-task records accumulate without bound.
7. Memory is exhausted; the daemon becomes unavailable to legitimate clients.

---

Impact

Dimension	Assessment
Likelihood	High — unauthenticated, reachable over the primary TCP service; no credentials or expensive resources required beyond repeated TCP connections.
Impact	High — unbounded memory growth; service unavailability. No RCE, authentication bypass, or data disclosure.
Severity	High

---

Existing Controls

• Authentication is not required before a connection-handler task is spawned.
• No in-loop JoinSet draining is present.
• No repository evidence of connection rate limiting or maximum concurrent/completed task accounting.
• Default bind address is all interfaces on TCP port 5500.
• Shutdown handling drains tasks only after the daemon is already stopping.

---

Suggested Remediation Direction

Add a non-blocking try_join_next arm (or an explicit drain step) inside the accept loop so completed tasks are reaped promptly. Alternatively, use detached tokio::spawn with a separate atomic counter or semaphore to track active tasks for graceful shutdown, avoiding unbounded JoinSet accumulation entirely.

---

Assumptions / Blindspots

• The crate could not be built or dynamically exercised in the validation environment due to dependency/network restrictions; however, the source path and a dependency-free harness strongly support reachability.
• No deployment manifests or IaC files were present to confirm whether real deployments expose the default public bind address or restrict it with a firewall.
• The exact memory growth rate per retained completed Tokio task was not measured; time-to-exhaustion is workload- and host-dependent.

[leynos]

@CodeRabbit plan

[coderabbitai]

Coding Plan

Summary

• Introduce a non-blocking try_join_next drain loop after each successful accept to reap completed client tasks promptly.
• Preserve the existing post-shutdown await_spawned_tasks call for graceful shutdown of in-flight tasks.
• Add a dedicated unit test confirming tasks are reaped during normal operation, not only at shutdown.

Design Choices

Design Choice 1: Reaping strategy

Options Considered:

1. Non-blocking drain via try_join_next — after each accept, loop over try_join_next() to reap all completed tasks immediately; minimal code change, preserves existing graceful-shutdown semantics.
2. Detached spawn with atomic counter — replace JoinSet with bare tokio::spawn and an AtomicUsize tracking active tasks; requires restructuring shutdown drain logic and introduces new synchronisation primitives.

Chosen Option: 1

Rationale: Option 1 — try_join_next is idiomatic Tokio, requires no new primitives, and fits naturally into the existing select loop without altering shutdown behaviour.

💡 User Tips

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

Implementation Steps

Reap Completed Tasks During Accept Loop

This phase eliminates the unbounded memory growth by draining finished tasks as part of normal operation, ensuring completed task records do not accumulate until shutdown.

Task 1: Add Non-Blocking Drain in `accept_connections`

Insert a while let Some(result) = join_set.try_join_next() loop immediately after the call to handle_accept_result inside the res = listener.accept() arm of the tokio::select!. This loop must log any JoinError to stderr (mirroring await_spawned_tasks behaviour) and continue until try_join_next returns None.

• Update src/server/legacy/mod.rs: within the res = listener.accept() arm, add a non-blocking drain loop after handle_accept_result returns
• Reuse the same error-logging pattern as await_spawned_tasks (eprintln!("task error: {e}")) for consistency
• Leave the post-loop await_spawned_tasks call unchanged so in-flight tasks are still awaited at shutdown

Task 2: Add Unit Test for In-Loop Reaping

Create a test that verifies completed tasks are reaped before shutdown. The test must spawn a client, complete a handshake, disconnect, then confirm the JoinSet length has decreased (or Arc refcount has dropped) without sending the shutdown signal.

• Add a new test to src/server/legacy/unit_tests.rs named completed_tasks_are_reaped_before_shutdown
• Set up a loopback listener and connect a client; call handle_accept_result to spawn into the JoinSet
• Complete the handshake and close the client connection; yield to allow the handler task to exit
• Invoke the new drain logic (either via a helper or by calling handle_accept_result with a second dummy accept to trigger the drain path)
• Assert the JoinSet is empty (or Arc::strong_count equals the expected baseline) before any shutdown signal is sent

🤖 Prompt for AI agents

Implement the non-blocking task-reaping fix in `src/server/legacy/mod.rs` and
add a corresponding unit test in `src/server/legacy/unit_tests.rs`. The goal is
to drain completed `JoinSet` entries on every accept iteration, eliminating
unbounded memory growth from accumulated task records.

**Task 1 — Add non-blocking drain loop in `accept_connections`:**
- Open `src/server/legacy/mod.rs` and locate the `res = listener.accept()` arm
inside the `tokio::select!` loop within `accept_connections`
- Immediately after the call to `handle_accept_result`, insert a `while let
Some(result) = join_set.try_join_next()` loop
- Inside the loop, log any `JoinError` using `eprintln!("task error: {e}")` to
mirror the existing pattern in `await_spawned_tasks`, then continue until
`try_join_next` returns `None`
- Do NOT remove or alter the existing `await_spawned_tasks` call that runs after
the select loop exits — it must remain to handle in-flight tasks at shutdown

**Task 2 — Add unit test `completed_tasks_are_reaped_before_shutdown`:**
- Add the new test to `src/server/legacy/unit_tests.rs`
- Set up a loopback listener and connect a client; call `handle_accept_result`
to spawn the handler into the `JoinSet`
- Complete the handshake and close the client connection; yield execution to
allow the handler task to exit
- Invoke the drain logic (via a helper or by triggering `handle_accept_result`
with a second dummy accept to exercise the drain path)
- Assert the `JoinSet` is empty (or `Arc::strong_count` equals the expected
baseline) **before** any shutdown signal is sent

Research

The server implementation lives in src/server/legacy/mod.rs, not src/main.rs (which is a thin shim). The accept_connections function runs a tokio::select! loop over a shutdown signal and listener.accept(). Accepted connections are spawned into a JoinSet<()> via handle_accept_result → spawn_client_handler. The await_spawned_tasks helper drains the JoinSet but is only called after the shutdown signal breaks the loop. Existing tests in src/server/legacy/unit_tests.rs verify Arc refcount behaviour and task completion, and the test_support feature exposes handle_accept_result for integration testing.

---

🚀 Next Steps

🤖 All AI agent prompts combined

Implement the non-blocking task-reaping fix in `src/server/legacy/mod.rs` and
add a corresponding unit test in `src/server/legacy/unit_tests.rs`. The goal is
to drain completed `JoinSet` entries on every accept iteration, eliminating
unbounded memory growth from accumulated task records.

**Task 1 — Add non-blocking drain loop in `accept_connections`:**
- Open `src/server/legacy/mod.rs` and locate the `res = listener.accept()` arm
inside the `tokio::select!` loop within `accept_connections`
- Immediately after the call to `handle_accept_result`, insert a `while let
Some(result) = join_set.try_join_next()` loop
- Inside the loop, log any `JoinError` using `eprintln!("task error: {e}")` to
mirror the existing pattern in `await_spawned_tasks`, then continue until
`try_join_next` returns `None`
- Do NOT remove or alter the existing `await_spawned_tasks` call that runs after
the select loop exits — it must remain to handle in-flight tasks at shutdown

**Task 2 — Add unit test `completed_tasks_are_reaped_before_shutdown`:**
- Add the new test to `src/server/legacy/unit_tests.rs`
- Set up a loopback listener and connect a client; call `handle_accept_result`
to spawn the handler into the `JoinSet`
- Complete the handshake and close the client connection; yield execution to
allow the handler task to exit
- Invoke the drain logic (via a helper or by triggering `handle_accept_result`
with a second dummy accept to exercise the drain path)
- Assert the `JoinSet` is empty (or `Arc::strong_count` equals the expected
baseline) **before** any shutdown signal is sent

💡 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!