use super::*;
use crate::messages::encode_raw_length;
use crate::transport::raw_capture::test_support;
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
#[tokio::test]
async fn read_framed_message_taps_raw_bytes_including_rejected_prefixes() {
let dir = tempfile::TempDir::new().unwrap();
let tap = RawFrameTap::capturing_to(dir.path());
let good = encode_raw_length(&[0, 0, 0, 9, 42]);
let bad_prefix = u32::MAX.to_be_bytes();
let mut stream = std::io::Cursor::new([good.clone(), bad_prefix.to_vec()].concat());
assert_eq!(read_framed_message(&mut stream, &tap).await.unwrap(), [0, 0, 0, 9, 42]);
let err = read_framed_message(&mut stream, &tap).await.expect_err("prefix must be rejected");
assert!(matches!(err, Error::InvalidFrame(_)), "got {err:?}");
let mut expected = good;
expected.extend_from_slice(&bad_prefix);
assert_eq!(test_support::frames(dir.path()), expected);
}
const LARGE: usize = 32 * 1024 * 1024;
async fn socket_pair() -> (Arc<AsyncTcpSocket>, TcpStream, tokio::net::TcpListener) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap().to_string();
let (socket, peer) = tokio::join!(AsyncTcpSocket::connect(&address, true), listener.accept());
(Arc::new(socket.unwrap()), peer.unwrap().0, listener)
}
async fn abandon_write(socket: &AsyncTcpSocket, bytes: &[u8]) {
let written = tokio::time::timeout(Duration::from_millis(100), socket.write_all(bytes)).await;
assert!(written.is_err(), "write completed; the peer should not be reading yet");
}
async fn until_writer_held(socket: &AsyncTcpSocket) {
tokio::time::timeout(Duration::from_secs(1), async {
while socket.writer.try_lock().is_ok() {
tokio::task::yield_now().await;
}
})
.await
.expect("large write never took the writer");
}
pub(in crate::transport::r#async) async fn assert_peer_sees_eof(mut peer: TcpStream) {
let eof = tokio::time::timeout(Duration::from_secs(1), peer.read_u8())
.await
.expect("TCP socket retained");
assert_eq!(eof.unwrap_err().kind(), std::io::ErrorKind::UnexpectedEof);
}
fn read_to_end(mut peer: TcpStream) -> tokio::task::JoinHandle<Vec<u8>> {
tokio::spawn(async move {
let mut received = Vec::new();
peer.read_to_end(&mut received).await.unwrap();
received
})
}
#[tokio::test]
async fn write_cut_off_mid_frame_breaks_the_connection() {
let (socket, peer, _listener) = socket_pair().await;
let large = vec![0xAA; LARGE];
let (read, ()) = tokio::join!(
tokio::time::timeout(Duration::from_secs(1), socket.read_message()),
abandon_write(&socket, &large),
);
let read = read.expect("pending read not woken");
assert!(matches!(read, Err(Error::ConnectionReset)), "read: {read:?}");
let written = socket.write_all(&[0xCC; 3]).await;
assert!(matches!(written, Err(Error::ConnectionReset)), "write: {written:?}");
let reader = read_to_end(peer);
drop(socket);
let received = reader.await.unwrap();
assert!(!received.is_empty() && received.len() < LARGE, "no fragment: {} bytes", received.len());
assert!(received.iter().all(|&b| b == 0xAA), "bytes followed the fragment");
}
#[tokio::test]
async fn write_abandoned_before_it_starts_sends_nothing() {
let (socket, peer, _listener) = socket_pair().await;
let large = vec![0xAA; LARGE];
let stalled = tokio::spawn({
let socket = socket.clone();
let large = large.clone();
async move { socket.write_all(&large).await }
});
until_writer_held(&socket).await;
abandon_write(&socket, &[0xBB; 3]).await;
let reader = read_to_end(peer);
stalled.await.unwrap().unwrap();
socket.write_all(&[0xCC; 3]).await.unwrap();
drop(socket);
let received = reader.await.unwrap();
assert!(received == [large, vec![0xCC; 3]].concat(), "received {} bytes", received.len());
}
#[tokio::test]
async fn reconnect_clears_a_broken_connection() {
let (socket, _peer, listener) = socket_pair().await;
abandon_write(&socket, &vec![0xAA; LARGE]).await;
assert!(matches!(socket.write_all(&[0xCC; 3]).await, Err(Error::ConnectionReset)));
let (reconnected, accepted) = tokio::join!(socket.reconnect(), listener.accept());
reconnected.unwrap();
let mut peer = accepted.unwrap().0;
socket.write_all(&[0xCC; 3]).await.unwrap();
let mut received = [0u8; 3];
peer.read_exact(&mut received).await.unwrap();
assert_eq!(received, [0xCC; 3]);
peer.write_all(&encode_raw_length(&[0, 0, 0, 9, 42])).await.unwrap();
assert_eq!(socket.read_message().await.unwrap(), [0, 0, 0, 9, 42]);
}
#[tokio::test]
async fn physical_close_wakes_read_and_is_terminal() {
let (socket, peer, _listener) = socket_pair().await;
let reading = socket.clone();
let read = tokio::spawn(async move { reading.read_message().await });
tokio::task::yield_now().await;
tokio::time::timeout(Duration::from_secs(1), socket.close())
.await
.expect("close blocked on read");
let read = read.await.unwrap();
assert!(matches!(read, Err(Error::Shutdown)), "read: {read:?}");
assert_peer_sees_eof(peer).await;
socket.close().await;
assert!(matches!(socket.write_all(&[1]).await, Err(Error::Shutdown)));
assert!(matches!(socket.read_message().await, Err(Error::Shutdown)));
assert!(matches!(socket.reconnect().await, Err(Error::Shutdown)));
}
#[tokio::test]
async fn physical_close_overrides_a_broken_connection() {
let (socket, _peer, _listener) = socket_pair().await;
abandon_write(&socket, &vec![0xAA; LARGE]).await;
assert!(matches!(socket.write_all(&[0xCC; 3]).await, Err(Error::ConnectionReset)));
socket.close().await;
let read = socket.read_message().await;
assert!(matches!(read, Err(Error::Shutdown)), "read: {read:?}");
let written = socket.write_all(&[0xCC; 3]).await;
assert!(matches!(written, Err(Error::Shutdown)), "write: {written:?}");
assert!(matches!(socket.reconnect().await, Err(Error::Shutdown)));
}
#[tokio::test]
async fn physical_close_prevents_a_queued_write() {
let (socket, peer, _listener) = socket_pair().await;
let writer = socket.writer.lock().await;
let sending = socket.clone();
let write = tokio::spawn(async move { sending.write_all(&[0xAA; 4]).await });
tokio::task::yield_now().await;
let closing = socket.clone();
let close = tokio::spawn(async move { closing.close().await });
tokio::time::timeout(Duration::from_secs(1), async {
while !socket.is_closed() {
tokio::task::yield_now().await;
}
})
.await
.expect("close never latched");
drop(writer);
tokio::time::timeout(Duration::from_secs(1), close).await.unwrap().unwrap();
assert!(matches!(write.await.unwrap(), Err(Error::Shutdown)));
assert_peer_sees_eof(peer).await;
}
#[tokio::test]
async fn physical_close_wakes_a_backpressured_write() {
let (socket, _peer, _listener) = socket_pair().await;
let write = tokio::spawn({
let socket = socket.clone();
async move { socket.write_all(&vec![0xAA; LARGE]).await }
});
until_writer_held(&socket).await;
tokio::time::timeout(Duration::from_secs(1), socket.close())
.await
.expect("close blocked on write");
let written = write.await.unwrap();
assert!(matches!(written, Err(Error::Shutdown)), "write: {written:?}");
assert!(matches!(socket.reconnect().await, Err(Error::Shutdown)));
}