methanium/mtp
Watch
3
0
Fork
You've already forked mtp
 0
Code Issues Pull requests 3 Projects Releases 22 Packages Wiki Activity Actions

Merge

Browse source 
Crypto
WASM
TESTS
This commit is contained in:
Alex Emmet 2026-06-25 22:08:44 +02:00
commit 687e6f9642
49 changed files with 6272 additions and 366 deletions
Download patch file Download diff file Expand all files Collapse all files

37
transport/src/connection.rs
View file

@ -61,7 +61,7 @@ enum ReceivedFrame {
pub struct Sender {
send_guard: Mutex<()>,
stream_guard: Mutex<Option<wtransport::SendStream>>,
stream_guard: Arc<Mutex<Option<wtransport::SendStream>>>,
handle: Arc<ConnectionHandle>,
connection: Connection,
policy: Arc<Policy>,
@ -71,7 +71,7 @@ impl Sender {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
Self {
send_guard: Mutex::new(()),
stream_guard: Mutex::new(None),
stream_guard: Arc::new(Mutex::new(None)),
handle,
connection,
policy,
@ -275,6 +275,18 @@ impl Sender {
}
}
pub async fn finish_stream(&self) -> Result<(), CommunicationError> {
let _send_lock = self.send_guard.lock().await;
let mut stream_opt = self.stream_guard.lock().await;
if let Some(mut stream) = stream_opt.take() {
timeout(self.policy.write_timeout, stream.finish())
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(|_| CommunicationError::StreamError)?;
}
Ok(())
}
pub fn handle(&self) -> &Arc<ConnectionHandle> {
&self.handle
}
@ -283,6 +295,7 @@ impl Sender {
let connection = self.connection.clone();
let handle = self.handle.clone();
let policy = self.policy.clone();
let stream_guard = self.stream_guard.clone();
tokio::spawn(async move {
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
@ -290,6 +303,14 @@ impl Sender {
return;
}
if let Some(mut stream) = stream_guard.lock().await.take() {
match timeout(policy.write_timeout, stream.finish()).await {
Ok(Ok(())) => {}
Ok(Err(e)) => log::warn!("[Sender] persistent stream finish failed: {e}"),
Err(_) => log::warn!("[Sender] persistent stream finish timed out"),
}
}
let _ = Self::send_close_frame(&connection, &policy).await;
handle.close(Some(CommunicationError::StreamClosed));
@ -323,6 +344,18 @@ pub struct Receiver {
handle: Arc<ConnectionHandle>,
}
impl Drop for Receiver {
fn drop(&mut self) {
// The accept loop holds clones of the connection and the shared
// ConnectionHandle. Without this, dropping a Receiver without first
// closing the connection would leave that task running forever. Abort
// it directly rather than closing the shared handle, so a still-live
// Sender on the same connection is unaffected. abort() is a no-op if
// the task already finished (e.g. the connection was closed).
self._accept_task.abort();
}
}
impl Receiver {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
let (tx, rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(

2
transport/src/connection_handle.rs
View file

@ -53,7 +53,7 @@ impl ConnectionHandle {
if self.is_closed() {
return rx.borrow().clone();
}
let _ = rx.changed().await.ok()?;
rx.changed().await.ok()?;
rx.borrow().clone()
}
}

16
transport/src/host.rs
View file

@ -26,6 +26,22 @@ impl Host {
pub fn local_addr(&self) -> std::net::SocketAddr {
self.local_addr
}
/// Stop accepting new connections. Already-accepted connections run on their
/// own spawned tasks and are not affected.
pub fn shutdown(&mut self) {
self._task.abort();
}
}
impl Drop for Host {
fn drop(&mut self) {
// The accept loop runs forever on its own task; dropping the Host
// JoinHandle would only detach it. Abort it so dropping the Host
// actually stops accepting new connections. Per-connection handler
// tasks are spawned independently and keep running.
self._task.abort();
}
}
pub async fn host(

142
transport/tests/integration.rs
View file

@ -1,6 +1,6 @@
use std::net::{IpAddr, Ipv4Addr};
use mtp_transport::{Policy, host, connect};
use mtp_transport::{Policy, connect, host};
fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
let key_pair = rcgen::KeyPair::generate().unwrap();
@ -44,18 +44,18 @@ async fn test_send_receive_roundtrip() {
let addr = h.local_addr();
let url = format!("https://127.0.0.1:{}", addr.port());
let (client_tx, client_rx) =
connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
let (client_tx, client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await
.unwrap();
// Accept on host side
let (host_tx, host_rx) = h.next().await.unwrap();
// Client sends a simple message
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping)
.add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(42),
);
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping).add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(42),
);
client_tx.send(&msg).await.unwrap();
// Host receives it
@ -65,16 +65,18 @@ async fn test_send_receive_roundtrip() {
assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(42));
// Host sends a response
let resp = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong)
.add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(99),
);
let resp = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong).add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(99),
);
host_tx.send(&resp).await.unwrap();
// Client receives it
let client_received = client_rx.receive().await.unwrap();
assert_eq!(client_received.get_type(), mtp_codec::CommunicationTypeId(20)); // Pong
assert_eq!(
client_received.get_type(),
mtp_codec::CommunicationTypeId(20)
); // Pong
let client_val = client_received.get_data(mtp_codec::DataTypeId(6)).clone();
assert_eq!(client_val, mtp_codec::DataValue::UnsignedNumber(99));
@ -98,18 +100,18 @@ async fn test_concurrent_messages() {
let addr = h.local_addr();
let url = format!("https://127.0.0.1:{}", addr.port());
let (client_tx, _client_rx) =
connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await
.unwrap();
let (_host_tx, host_rx) = h.next().await.unwrap();
// Send 5 messages in sequence
for i in 0..5u128 {
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping)
.add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(i),
);
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping).add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(i),
);
client_tx.send(&msg).await.unwrap();
}
@ -122,11 +124,10 @@ async fn test_concurrent_messages() {
// Send 3 responses back
for i in 0..3u128 {
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong)
.add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(i * 10),
);
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong).add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(i * 10),
);
client_tx.send(&msg).await.unwrap();
}
@ -154,8 +155,9 @@ async fn test_close_detection() {
let addr = h.local_addr();
let url = format!("https://127.0.0.1:{}", addr.port());
let (client_tx, _client_rx) =
connect(&url, Some(cert_pem), Policy::default()).await.unwrap();
let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await
.unwrap();
let (_host_tx, host_rx) = h.next().await.unwrap();
@ -172,3 +174,87 @@ async fn test_close_detection() {
let result = host_rx.receive().await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_host_shutdown_stops_accepting() {
let (cert_pem, key_pem) = generate_self_signed_cert();
let mut h = host(
IpAddr::V4(Ipv4Addr::LOCALHOST),
0,
cert_pem.clone(),
key_pem,
Policy::default(),
)
.await
.unwrap();
let addr = h.local_addr();
let url = format!("https://127.0.0.1:{}", addr.port());
// A connection succeeds while the host is accepting.
let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default())
.await
.unwrap();
let _accepted = h.next().await.unwrap();
// After shutdown the accept task is aborted and its endpoint is dropped, so
// new connections no longer succeed. Guard with a timeout so a hung connect
// still fails the assertion rather than blocking the test.
h.shutdown();
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
connect(&url, Some(cert_pem), Policy::default()),
)
.await;
assert!(
matches!(result, Err(_) | Ok(Err(_))),
"connect should not succeed after host shutdown"
);
}
#[tokio::test]
async fn test_drop_receiver_keeps_sender_alive() {
let (cert_pem, key_pem) = generate_self_signed_cert();
let mut h = host(
IpAddr::V4(Ipv4Addr::LOCALHOST),
0,
cert_pem.clone(),
key_pem,
Policy::default(),
)
.await
.unwrap();
let addr = h.local_addr();
let url = format!("https://127.0.0.1:{}", addr.port());
let (client_tx, client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await
.unwrap();
let (host_tx, host_rx) = h.next().await.unwrap();
// Client sends a message the host receives.
let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping);
client_tx.send(&msg).await.unwrap();
let _ = host_rx.receive().await.unwrap();
// Dropping the host Receiver aborts only its accept task; the Sender shares
// the same connection and must keep working.
drop(host_rx);
let resp = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong).add_data(
mtp_codec::DataTypeId(6),
mtp_codec::DataValue::UnsignedNumber(7),
);
host_tx.send(&resp).await.unwrap();
let got = client_rx.receive().await.unwrap();
assert_eq!(got.get_type(), mtp_codec::CommunicationTypeId(20)); // Pong
assert_eq!(
got.get_data(mtp_codec::DataTypeId(6)).clone(),
mtp_codec::DataValue::UnsignedNumber(7)
);
client_tx.close();
host_tx.close();
}