Replace main with always-masque-relay branch contents - #63
Conversation
| let (fwd_tx, mut fwd_rx) = | ||
| tokio::sync::mpsc::channel::<Bytes>(RELAY_FORWARD_CHANNEL_CAPACITY); | ||
|
|
||
| // Reader: UDP socket → channel (never blocked by stream writes) | ||
| let reader_handle = tokio::spawn(async move { | ||
| let mut buf = vec![0u8; 65536]; | ||
| loop { | ||
| match socket.recv_from(&mut buf).await { | ||
| Ok((len, source)) => { | ||
| let payload = Bytes::copy_from_slice(&buf[..len]); | ||
| tracing::trace!( | ||
| session_id, source = %source, len, | ||
| "Stream relay: received UDP from target" | ||
| ); | ||
| let datagram = | ||
| UncompressedDatagram::new(VarInt::from_u32(0), source, payload); | ||
| let encoded = datagram.encode(); | ||
| stats.record_bytes(encoded.len() as u64); | ||
| stats.record_datagram(); | ||
| if fwd_tx.send(encoded).await.is_err() { | ||
| break; // writer closed | ||
| } | ||
| } | ||
| Err(e) => { | ||
| tracing::debug!(session_id, error = %e, "UDP recv error"); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
| }); | ||
|
|
||
| tokio::select! { | ||
| // TODO: Rate limiting — check_rate_limit should be called in both | ||
| // directions to enforce the per-session bandwidth_limit from | ||
| // RelaySessionConfig. Currently the stream path bypasses rate | ||
| // limiting entirely. Requires passing the session's rate limiter | ||
| // into this loop. | ||
| // | ||
| // Direction 1: UDP → Stream (target → relay → client) | ||
| _ = async { | ||
| let mut buf = vec![0u8; 65536]; | ||
| loop { | ||
| match socket.recv_from(&mut buf).await { | ||
| Ok((len, source)) => { | ||
| let payload = Bytes::copy_from_slice(&buf[..len]); | ||
| tracing::trace!( | ||
| session_id, source = %source, len, | ||
| "Stream relay: received UDP from target" | ||
| ); | ||
|
|
||
| let datagram = UncompressedDatagram::new( | ||
| VarInt::from_u32(0), source, payload, | ||
| ); | ||
| let encoded = datagram.encode(); | ||
|
|
||
| // Write length-prefixed frame to stream | ||
| let frame_len = encoded.len() as u32; | ||
| if let Err(e) = send_stream.write_all(&frame_len.to_be_bytes()).await { | ||
| tracing::debug!(session_id, error = %e, "Stream write error (length)"); | ||
| break; | ||
| } | ||
| if let Err(e) = send_stream.write_all(&encoded).await { | ||
| tracing::debug!(session_id, error = %e, "Stream write error (data)"); | ||
| break; | ||
| } | ||
| // Writer: channel → QUIC stream (paced by stream flow control) | ||
| let writer_handle = tokio::spawn(async move { | ||
| let mut keepalive = tokio::time::interval(RELAY_KEEPALIVE_INTERVAL); | ||
| keepalive.tick().await; // skip immediate first tick | ||
|
|
||
| stats.record_bytes(encoded.len() as u64); | ||
| stats.record_datagram(); | ||
| loop { | ||
| tokio::select! { | ||
| item = fwd_rx.recv() => { | ||
| let Some(encoded) = item else { break }; | ||
| let frame_len = encoded.len() as u32; | ||
| if let Err(e) = send_stream.write_all(&frame_len.to_be_bytes()).await { | ||
| tracing::debug!(session_id, error = %e, "Stream write error (length)"); | ||
| break; | ||
| } | ||
| Err(e) => { | ||
| tracing::debug!(session_id, error = %e, "UDP recv error"); | ||
| if let Err(e) = send_stream.write_all(&encoded).await { | ||
| tracing::debug!(session_id, error = %e, "Stream write error (data)"); | ||
| break; | ||
| } | ||
| } | ||
| _ = keepalive.tick() => { | ||
| if let Err(e) = send_stream.write_all(&keepalive_bytes).await { | ||
| tracing::debug!(session_id, error = %e, "Keepalive write error"); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
| } => {}, | ||
| } | ||
| }); | ||
|
|
||
| tokio::select! { | ||
| _ = reader_handle => {}, |
There was a problem hiding this comment.
Leaked background tasks after
select! exits in run_stream_forwarding_loop
reader_handle and writer_handle are spawned detached tasks. When tokio::select! terminates early (e.g. the Direction-2 branch exits), the JoinHandles are dropped but the tasks continue running — dropping a JoinHandle in Tokio does not abort the spawned task. After close_session releases the session's Arc<UdpSocket>, reader_handle still holds its own Arc clone and keeps draining UDP traffic. This prevents the OS from reclaiming the socket file-descriptor until the reader task finishes naturally, which may require an I/O error that never comes. Consider aborting the handles explicitly when select! exits:
tokio::select! {
_ = &mut reader_handle => { writer_handle.abort(); },
_ = &mut writer_handle => { reader_handle.abort(); },
_ = async { /* Direction 2 */ } => {
reader_handle.abort();
writer_handle.abort();
},
}Prompt To Fix With AI
This is a comment left during a code review.
Path: src/masque/relay_server.rs
Line: 1065-1126
Comment:
**Leaked background tasks after `select!` exits in `run_stream_forwarding_loop`**
`reader_handle` and `writer_handle` are spawned detached tasks. When `tokio::select!` terminates early (e.g. the Direction-2 branch exits), the `JoinHandle`s are **dropped** but the tasks continue running — dropping a `JoinHandle` in Tokio does not abort the spawned task. After `close_session` releases the session's `Arc<UdpSocket>`, `reader_handle` still holds its own `Arc` clone and keeps draining UDP traffic. This prevents the OS from reclaiming the socket file-descriptor until the reader task finishes naturally, which may require an I/O error that never comes. Consider aborting the handles explicitly when select! exits:
```rust
tokio::select! {
_ = &mut reader_handle => { writer_handle.abort(); },
_ = &mut writer_handle => { reader_handle.abort(); },
_ = async { /* Direction 2 */ } => {
reader_handle.abort();
writer_handle.abort();
},
}
```
How can I resolve this? If you propose a fix, please make it concise.| pub fn unregister_upstream_relay(&self, relay: SocketAddr) { | ||
| self.upstream_relay_ips.write().remove(&relay.ip()); | ||
| } |
There was a problem hiding this comment.
unregister_upstream_relay silently drops other relays on the same IP
The comment on lines 148-153 acknowledges this: "If multiple upstream relays share the same IP (unusual but possible over different ports) this removes the IP entirely." However the recommended remediation — "callers should register after removal if any remain" — is fragile and error-prone. In practice, a node could have two active relay sessions to two different ports on the same public IP (e.g., two relay servers behind the same NAT). Deregistering one would unblock clients from that IP for the other, violating the upstream-IP policy. A reference-counted or multi-value set would be safer here.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/masque/ip_policy.rs
Line: 152-154
Comment:
**`unregister_upstream_relay` silently drops other relays on the same IP**
The comment on lines 148-153 acknowledges this: "If multiple upstream relays share the same IP (unusual but possible over different ports) this removes the IP entirely." However the recommended remediation — "callers should register after removal if any remain" — is fragile and error-prone. In practice, a node could have two active relay sessions to two different ports on the same public IP (e.g., two relay servers behind the same NAT). Deregistering one would unblock clients from that IP for the other, violating the upstream-IP policy. A reference-counted or multi-value set would be safer here.
How can I resolve this? If you propose a fix, please make it concise.
Benchmark ResultsPerformance Comparison
SummaryConfiguration
|
87f92d0 to
ec16340
Compare
Summary
This PR supersedes
origin/mainwith the contents ofmick/always-masque-relay-rebased. It uses a-s oursmerge so the resulting tree is exactly this branch's tree — none of the 26 commits onmainthat are not in this branch are applied.Why a
-s oursmergeorigin/mainas merged, so git tooling treats the branches as reconciled afterward.Merge instructions
Use "Create a merge commit" when merging this PR. Squash or rebase will discard the
-s ourssemantics and may reintroduce the superseded changes.Test plan
mainafter merge matches this branch (git diff main origin/mick/always-masque-relay-rebasedis empty)Greptile Summary
This PR replaces
mainwith themick/always-masque-relay-rebasedbranch using a-s oursmerge, bringing in a large set of MASQUE relay, BBRv2 congestion control, IP-diversity policy, bounded-channel backpressure, and NAT-traversal improvements. The relay infrastructure is the most significant new surface; two issues need attention before this lands:close_session(relay_server.rs): the session is markedClosedunder the first write lock, released, then removed under a second write lock. Any concurrent caller (e.g.session_count(),cleanup_expired_sessions) observes a phantom closed-but-present session in between, which can cause inflated session counts againstmax_sessionsand spurious double-close attempts.run_stream_forwarding_loop:reader_handleandwriter_handleare spawned as detached tasks. Whentokio::select!exits early, theJoinHandles are dropped but the tasks keep running, holding theirArc<UdpSocket>clone and preventing FD reclamation for the lifetime of the task.Confidence Score: 3/5
Two P1 defects in the relay server's session lifecycle need fixes before merging to avoid ghost sessions and FD leaks under load.
The double write-lock TOCTOU in close_session causes incorrect session counts against the max_sessions limit and can confuse cleanup logic. The leaked JoinHandle background tasks in run_stream_forwarding_loop cause UDP socket FDs to linger after session teardown. Both are present-state defects on the hot relay path, not speculative. The rest of the PR (BBRv2, bounded channels, IP policy, NAT traversal) looks well-structured.
src/masque/relay_server.rs — close_session (double lock) and run_stream_forwarding_loop (task leak)
Important Files Changed
Sequence Diagram
sequenceDiagram participant Client as QUIC Client participant RS as MasqueRelayServer participant RSocket as MasqueRelaySocket participant UDP as Bound UDP Socket participant Target as Target Peer Client->>RS: CONNECT-UDP (handle_connect_request) RS->>UDP: bind(INADDR_ANY:0) RS->>Client: ConnectUdpResponse(advertised_addr) Note over RS: run_stream_forwarding_loop RS->>RS: spawn reader_handle (UDP→channel) RS->>RS: spawn writer_handle (channel→stream) par Direction 2: Client→Target Client->>RSocket: length-prefixed frame (stream) RSocket->>RS: recv_stream.read_exact RS->>UDP: send_to(target) UDP->>Target: raw UDP and Direction 1: Target→Client Target->>UDP: raw UDP UDP->>RS: reader_handle.recv_from RS->>RSocket: fwd_tx.send(encoded) RSocket->>Client: send_stream.write_all (writer_handle) end Note over RS: On any direction exit RS->>RS: close_session(session_id) RS--xClient: session cleaned upComments Outside Diff (2)
src/masque/relay_server.rs, line 1199-1222 (link)close_sessionclose_sessionacquires a write lock to callsession.close()(marking the session asClosed), releases the lock, then immediately re-acquires a second write lock to actually remove the session from the map. Between these two critical sections any concurrent reader can observe aClosed-state session that is still present insessions. Callers likesession_count()return an inflated count — which is checked againstmax_sessionsinhandle_connect_request— andcleanup_expired_sessionscan attempt to re-close the same session. Combine the two operations in a single lock acquisition:Prompt To Fix With AI
src/masque/relay_server.rs, line 496-502 (link)stats.current_active_sessions()is read (Ordering::Relaxed) and compared tomax_sessions, butrecord_session_terminatedusesfetch_sub(alsoRelaxed). Becauseactive_sessionsinMasqueRelayStatsis decremented byrecord_session_terminated, which is called only at the end ofclose_sessionafter both write locks, there is already a window wherecurrent_active_sessions()over-counts active sessions. This compounds with the double-lock TOCTOU noted above. No immediate fix needed beyond theclose_sessionfix, but worth noting that the stats counter and the actualsessionsmap can diverge ifclose_sessionerrors between the two write locks (the first lock marks-then-releases, but if?returned early the stat was never decremented).Prompt To Fix With AI
Prompt To Fix All With AI
Reviews (1): Last reviewed commit: "merge: supersede origin/main — keep bran..." | Re-trigger Greptile