From 3d9eebf44268def961e987f9353a4699bac8286d Mon Sep 17 00:00:00 2001 From: Aashish Thapa Date: Sun, 26 Jul 2026 16:50:42 -0500 Subject: [PATCH] fix(arti): relay until both halves close to avoid truncating responses --- localPackages/Arti/arti-bitchat/src/socks.rs | 195 +++++++++++++++++-- 1 file changed, 174 insertions(+), 21 deletions(-) diff --git a/localPackages/Arti/arti-bitchat/src/socks.rs b/localPackages/Arti/arti-bitchat/src/socks.rs index 4795a7d9..860a117c 100644 --- a/localPackages/Arti/arti-bitchat/src/socks.rs +++ b/localPackages/Arti/arti-bitchat/src/socks.rs @@ -7,7 +7,7 @@ use std::net::SocketAddr; use std::sync::Arc; use arti_client::{TorClient, IntoTorAddr}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::TcpStream; use tor_rtcompat::PreferredRuntime; @@ -168,31 +168,53 @@ pub async fn handle_socks_connection( ]; stream.write_all(&reply).await?; - // Bidirectional copy - let (mut client_read, mut client_write) = stream.into_split(); - let (mut tor_read, mut tor_write) = tor_stream.split(); + relay(stream, tor_stream).await +} - let client_to_tor = async { - tokio::io::copy(&mut client_read, &mut tor_write).await - }; - let tor_to_client = async { - tokio::io::copy(&mut tor_read, &mut client_write).await - }; +/// Relay bytes between the local SOCKS `client` and the `tor` stream until the +/// response is fully delivered. +/// +/// Tor streams are not half-closeable: shutting down the Tor writer emits a +/// RELAY_END cell that tears the whole stream down at the exit. So a client +/// half-close (EOF on client->tor) must not be propagated as a shutdown, or a +/// response still in flight is lost at the exit. Only the client's write half +/// is closed on EOF; the Tor stream's END is deferred until both halves drop at +/// the end of the relay. (`copy_bidirectional` cannot be used here: it shuts +/// down the peer writer at the first EOF, which for Tor is exactly this bug.) +async fn relay(client: C, tor: T) -> io::Result<()> +where + C: AsyncRead + AsyncWrite + Unpin, + T: AsyncRead + AsyncWrite + Unpin, +{ + let (mut client_read, mut client_write) = tokio::io::split(client); + let (mut tor_read, mut tor_write) = tokio::io::split(tor); - tokio::select! { - result = client_to_tor => { - if let Err(e) = result { - tracing::debug!("Client to Tor copy error: {}", e); - } - } - result = tor_to_client => { - if let Err(e) = result { - tracing::debug!("Tor to client copy error: {}", e); + // Request path: on EOF, flush the Tor writer but do not shut it down. + let upload = async { + tokio::io::copy(&mut client_read, &mut tor_write).await?; + tor_write.flush().await + }; + // Response path: on EOF (the exit sent END, or the destination finished), + // half-close the client's write side and finish. + let download = async { + tokio::io::copy(&mut tor_read, &mut client_write).await?; + client_write.shutdown().await + }; + tokio::pin!(upload, download); + + // The relay ends when the response path finishes or either path errors. The + // request path finishing only stops that copy; it never cuts the response + // short. + let mut upload_done = false; + loop { + tokio::select! { + result = &mut download => return result, + result = &mut upload, if !upload_done => { + result?; + upload_done = true; } } } - - Ok(()) } async fn send_reply(stream: &mut TcpStream, rep: u8) -> io::Result<()> { @@ -206,3 +228,134 @@ async fn send_reply(stream: &mut TcpStream, rep: u8) -> io::Result<()> { ]; stream.write_all(&reply).await } + +#[cfg(test)] +mod tests { + use super::*; + use std::pin::Pin; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::{Context, Poll}; + use std::time::Duration; + use tokio::io::ReadBuf; + + // The client half-closes its write side after sending a request; the far + // side replies only afterward. A relay that stops at the first EOF loses + // the response. The timeout turns that into a failure rather than a hang. + #[tokio::test] + async fn relay_delivers_response_after_client_half_close() { + let (mut client, a) = tokio::io::duplex(64); + let (mut server, b) = tokio::io::duplex(64); + + let server_task = tokio::spawn(async move { + let mut req = [0u8; 3]; + server.read_exact(&mut req).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + server.write_all(b"RESPONSE").await.unwrap(); + server.shutdown().await.unwrap(); + }); + + let client_task = tokio::spawn(async move { + client.write_all(b"REQ").await.unwrap(); + client.shutdown().await.unwrap(); + let mut got = Vec::new(); + let _ = + tokio::time::timeout(Duration::from_secs(2), client.read_to_end(&mut got)).await; + got + }); + + relay(a, b).await.unwrap(); + + assert_eq!(client_task.await.unwrap(), b"RESPONSE"); + server_task.await.unwrap(); + } + + // A Tor writer shutdown emits a stream-killing END, so the relay must never + // shut down the Tor side before that side has itself reached EOF. This wraps + // the Tor stream and flags a shutdown that happens too early — the property + // `copy_bidirectional` violates and the asymmetric relay upholds. + struct FlagEarlyShutdown { + inner: S, + read_eof: bool, + early_shutdown: Arc, + } + + impl AsyncRead for FlagEarlyShutdown { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + let before = buf.filled().len(); + let result = Pin::new(&mut this.inner).poll_read(cx, buf); + if let Poll::Ready(Ok(())) = &result { + if buf.filled().len() == before { + this.read_eof = true; + } + } + result + } + } + + impl AsyncWrite for FlagEarlyShutdown { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if !this.read_eof { + this.early_shutdown.store(true, Ordering::SeqCst); + } + Pin::new(&mut this.inner).poll_shutdown(cx) + } + } + + #[tokio::test] + async fn relay_does_not_shut_down_tor_writer_before_response() { + let (mut client, a) = tokio::io::duplex(64); + let (mut server, b) = tokio::io::duplex(64); + + let early_shutdown = Arc::new(AtomicBool::new(false)); + let tor = FlagEarlyShutdown { + inner: b, + read_eof: false, + early_shutdown: early_shutdown.clone(), + }; + + let server_task = tokio::spawn(async move { + let mut req = [0u8; 3]; + server.read_exact(&mut req).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + server.write_all(b"RESPONSE").await.unwrap(); + server.shutdown().await.unwrap(); + }); + + let client_task = tokio::spawn(async move { + client.write_all(b"REQ").await.unwrap(); + client.shutdown().await.unwrap(); + let mut got = Vec::new(); + let _ = + tokio::time::timeout(Duration::from_secs(2), client.read_to_end(&mut got)).await; + got + }); + + relay(a, tor).await.unwrap(); + + assert!( + !early_shutdown.load(Ordering::SeqCst), + "relay shut down the Tor writer before the response arrived" + ); + assert_eq!(client_task.await.unwrap(), b"RESPONSE"); + server_task.await.unwrap(); + } +}