From 0df51c3995b5c6dc655b42eaff33529f17b56d52 Mon Sep 17 00:00:00 2001 From: clover caruso Date: Mon, 5 Oct 2026 21:48:11 -0700 Subject: [PATCH] fix: reconnect Live Share after uncertain HTTP delivery A lost POST response or truncated GET can leave encrypted stream delivery unknown. Retire that polling session and reconnect instead of replaying or discarding its bytes. fixes #92 Assisted-by: gpt-6.1-sol --- crates/notebook/src/live/transport.rs | 75 +++---- crates/notebook/src/live/transport_tests.rs | 205 ++++++++++++++++++++ 2 files changed, 232 insertions(+), 48 deletions(-) create mode 100644 crates/notebook/src/live/transport_tests.rs diff --git a/crates/notebook/src/live/transport.rs b/crates/notebook/src/live/transport.rs index 30088c82cd871762bdb75e9d203922845459dae0..33fc6e21debadd3c452d70c2c54226da91e97e44 100644 --- a/crates/notebook/src/live/transport.rs +++ b/crates/notebook/src/live/transport.rs @@ -383,7 +383,7 @@ fn poll(address: &Address, path: &str) -> Result { session, address, at, - http: Some(http), + http, arrived: Vec::new(), read: 0, }), @@ -404,30 +404,26 @@ fn post(session: &Session, address: &Address, at: &str) { } std::mem::take(&mut *pending) }; - // A connection the relay or the proxy closed meanwhile is opened again, once. - let sent = (0..2).any(|_| { - if http.is_none() { - http = Http::open(address).ok(); - if let (Some(opened), Ok(mut open)) = (&http, session.open.lock()) { - open.extend(opened.tcp.try_clone()); - } - } - let answered = http - .as_mut() - .map(|connection| connection.ask(address, "POST", at, &batch)); - match answered { - Some(Ok((200 | 204, ..))) => true, - Some(Ok((410, ..))) => { + if http.is_none() { + http = Http::open(address).ok(); + if let Some(opened) = &http { + let Ok(tcp) = opened.tcp.try_clone() else { session.close(); - true - } - _ => { - http = None; - false + return; + }; + let mut open = session.open.lock().unwrap(); + if session.closed.load(Ordering::Acquire) { + let _ = tcp.shutdown(Shutdown::Both); + return; } + open.push(tcp); } - }); - if !sent { + } + let answered = http + .as_mut() + .map(|connection| connection.ask(address, "POST", at, &batch)); + // A lost HTTP answer may follow delivery; replaying its bytes corrupts the stream. + if !matches!(answered, Some(Ok((200 | 204, ..)))) { session.close(); return; } @@ -456,7 +452,7 @@ struct PollReader { session: Arc, address: Arc
, at: String, - http: Option, + http: Http, arrived: Vec, read: usize, } @@ -467,23 +463,7 @@ impl Read for PollReader { if self.session.closed.load(Ordering::Acquire) { return Ok(0); } - if self.http.is_none() { - let opened = Http::open(&self.address).map_err(|failure| match failure { - Failure::Trouble(_, error) => error, - Failure::Refused(status, _) => io::Error::other(format!("{status}")), - })?; - self.session - .open - .lock() - .unwrap() - .extend(opened.tcp.try_clone()); - self.http = Some(opened); - } - let answered = - self.http - .as_mut() - .expect("opened above") - .ask(&self.address, "GET", &self.at, &[]); + let answered = self.http.ask(&self.address, "GET", &self.at, &[]); match answered { Ok((200, _, body)) => { self.arrived = body; @@ -494,17 +474,12 @@ impl Read for PollReader { return Ok(0); } Ok((status, ..)) => { + self.session.close(); return Err(io::Error::other(format!("The relay answered {status}"))); } Err(error) => { - // Opened again once; a second failure ends the session. - if self.http.take().is_none() { - return Err(error); - } - self.http = Http::open(&self.address).ok(); - if self.http.is_none() { - return Err(error); - } + self.session.close(); + return Err(error); } } } @@ -606,3 +581,7 @@ impl Write for TlsWriter { Ok(()) } } + +#[cfg(test)] +#[path = "transport_tests.rs"] +mod tests; diff --git a/crates/notebook/src/live/transport_tests.rs b/crates/notebook/src/live/transport_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..c005cb9f6f4f3ed588d6ff67b5c47e1dd5abdd7b --- /dev/null +++ b/crates/notebook/src/live/transport_tests.rs @@ -0,0 +1,205 @@ +use super::*; +use std::{net::TcpListener, thread::JoinHandle, time::Instant}; + +fn request(stream: &mut TcpStream) -> io::Result<(String, Vec)> { + let head = ws::head(stream)?; + let length = ws::header(&head, "Content-Length") + .and_then(|length| length.parse().ok()) + .unwrap_or(0); + let mut body = vec![0; length]; + stream.read_exact(&mut body)?; + Ok((head, body)) +} + +fn answer(stream: &mut TcpStream, status: u16, body: &[u8]) -> io::Result<()> { + write!( + stream, + "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\n\r\n", + body.len() + )?; + stream.write_all(body) +} + +fn polling( + serve: impl FnOnce(TcpListener, TcpStream) -> T + Send + 'static, +) -> (Connection, JoinHandle) { + proxy::use_proxy(Some(None)); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = parse(&format!("ws://{}", listener.local_addr().unwrap())).unwrap(); + let serving = thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let (head, _) = request(&mut stream).unwrap(); + assert!(head.starts_with("GET /v1/room/test?poll=1 ")); + answer(&mut stream, 200, b"session test").unwrap(); + serve(listener, stream) + }); + let connection = match poll(&address, "/v1/room/test") { + Ok(connection) => connection, + Err(_) => panic!("the local polling relay refused the session"), + }; + (connection, serving) +} + +fn accept(listener: &TcpListener) -> Option { + listener.set_nonblocking(true).unwrap(); + let until = Instant::now() + Duration::from_millis(300); + loop { + match listener.accept() { + Ok((stream, _)) => { + stream.set_nonblocking(false).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + return Some(stream); + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock && Instant::now() < until => { + thread::sleep(Duration::from_millis(1)); + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => return None, + Err(error) => panic!("could not accept the polling request: {error}"), + } + } +} + +#[test] +fn an_accepted_post_with_a_lost_answer_retires_its_stream_without_replay() { + let sent = ws::frame(ws::BINARY, b"sealed edit", Some([1, 2, 3, 4])); + let expected = sent.clone(); + let (mut connection, serving) = polling(move |listener, mut reading| { + assert!( + request(&mut reading) + .unwrap() + .0 + .starts_with("GET /v1/poll/test ") + ); + let mut posting = accept(&listener).expect("the edit never reached the relay"); + let (head, body) = request(&mut posting).unwrap(); + assert!(head.starts_with("POST /v1/poll/test ")); + assert_eq!(body, expected); + drop(posting); + let replay = accept(&listener).map(|mut stream| { + let (_, body) = request(&mut stream).unwrap(); + answer(&mut stream, 200, b"").unwrap(); + body + }); + reading + .set_read_timeout(Some(Duration::from_millis(300))) + .unwrap(); + let closed = reading.read(&mut [0]).is_ok_and(|length| length == 0); + let _ = answer(&mut reading, 410, b""); + (replay, closed) + }); + connection.writer.write_all(&sent).unwrap(); + let read = connection.reader.read(&mut [0]); + let writing_closed = connection.writer.write(&sent).is_err(); + (connection.close)(); + let (replay, reading_closed) = serving.join().unwrap(); + assert!( + replay.is_none(), + "the same sealed bytes were replayed after an unknown outcome" + ); + assert!(reading_closed, "the other tracked connection stayed open"); + assert!(read.is_err() || read.is_ok_and(|length| length == 0)); + assert!(writing_closed); +} + +#[test] +fn a_drained_get_with_a_truncated_answer_closes_both_tracked_connections() { + let (mut connection, serving) = polling(|listener, mut reading| { + assert!( + request(&mut reading) + .unwrap() + .0 + .starts_with("GET /v1/poll/test ") + ); + let mut posting = accept(&listener).expect("the outgoing connection never opened"); + assert!( + request(&mut posting) + .unwrap() + .0 + .starts_with("POST /v1/poll/test ") + ); + reading + .write_all(b"HTTP/1.1 200 Test\r\nContent-Length: 5\r\n\r\nxx") + .unwrap(); + drop(reading); + posting + .set_read_timeout(Some(Duration::from_millis(300))) + .unwrap(); + let closed = posting.read(&mut [0]).is_ok_and(|length| length == 0); + let _ = answer(&mut posting, 200, b""); + let retry = accept(&listener).map(|mut stream| { + assert!( + request(&mut stream) + .unwrap() + .0 + .starts_with("GET /v1/poll/test ") + ); + answer(&mut stream, 200, b"later").unwrap(); + }); + (retry.is_some(), closed) + }); + connection + .writer + .write_all(&ws::frame(ws::PING, b"", Some([0; 4]))) + .unwrap(); + let started = Instant::now(); + let read = connection.reader.read(&mut [0; 5]); + let elapsed = started.elapsed(); + let writing_closed = read.is_err() && connection.writer.write(b"x").is_err(); + (connection.close)(); + let (retried, posting_closed) = serving.join().unwrap(); + assert!( + read.is_err(), + "bytes consumed by the failed GET were silently discarded" + ); + assert!( + !retried, + "a failed GET reused the existing encrypted stream" + ); + assert!(posting_closed, "the other tracked connection stayed open"); + assert!(writing_closed); + assert!(elapsed < Duration::from_secs(1)); +} + +#[test] +fn repeated_get_failures_return_without_a_reconnect_loop() { + let (mut connection, serving) = polling(|listener, mut reading| { + assert!( + request(&mut reading) + .unwrap() + .0 + .starts_with("GET /v1/poll/test ") + ); + drop(reading); + let mut attempts = 1; + while let Some(mut stream) = accept(&listener) { + assert!( + request(&mut stream) + .unwrap() + .0 + .starts_with("GET /v1/poll/test ") + ); + attempts += 1; + if attempts == 4 { + answer(&mut stream, 410, b"").unwrap(); + break; + } + } + attempts + }); + let started = Instant::now(); + let read = connection.reader.read(&mut [0]); + let elapsed = started.elapsed(); + (connection.close)(); + let attempts = serving.join().unwrap(); + assert!( + read.is_err(), + "request failures were retried inside the same session" + ); + assert_eq!(attempts, 1); + assert!(elapsed < Duration::from_secs(1)); +} -- 2.54.0