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

General & Crypto

Browse source
This commit is contained in:
Alex Emmet 2026-06-20 15:25:36 +02:00
commit 02f94993c7
27 changed files with 1881 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

93
transport/src/client.rs Normal file
View file

@ -0,0 +1,93 @@
use std::sync::Arc;
use mtp_common::CommunicationError;
use rustls::{ClientConfig as RustlsClientConfig, RootCertStore, pki_types::pem::PemObject};
use wtransport::{ClientConfig, Endpoint};
use crate::{ConnectionHandle, Policy, Receiver, Sender};
pub async fn connect(
url: &str,
server_cert: Option<Vec<u8>>,
policy: Policy,
) -> Result<(Sender, Receiver), CommunicationError> {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let client_config = if let Some(cert_pem) = server_cert {
configure_client_with_cert(cert_pem, &policy)?
} else {
configure_client_system_roots(&policy)?
};
let endpoint = Endpoint::client(client_config)
.map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?;
let connection = endpoint
.connect(url)
.await
.map_err(|e| CommunicationError::ConnectingError(e.to_string()))?;
let handle = Arc::new(ConnectionHandle::new());
let policy = Arc::new(policy);
let sender = Sender::new(connection.clone(), handle.clone(), policy.clone());
let receiver = Receiver::new(connection, handle, policy);
Ok((sender, receiver))
}
fn configure_client_with_cert(
server_cert: Vec<u8>,
policy: &Policy,
) -> Result<ClientConfig, CommunicationError> {
let mut root_store = RootCertStore::empty();
let certs = rustls::pki_types::CertificateDer::pem_slice_iter(&server_cert)
.collect::<Result<Vec<_>, _>>()
.map_err(|_| CommunicationError::CertificateParseFailed)?;
for cert in certs {
root_store
.add(cert)
.map_err(|_| CommunicationError::CertificateParseFailed)?;
}
let mut tls_config = RustlsClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
tls_config.alpn_protocols = vec![b"h3".to_vec()];
Ok(ClientConfig::builder()
.with_bind_default()
.with_custom_tls(tls_config)
.keep_alive_interval(policy.keep_alive_interval)
.max_idle_timeout(policy.max_idle_timeout)
.map_err(|e| CommunicationError::Other(e.to_string()))?
.build())
}
fn configure_client_system_roots(policy: &Policy) -> Result<ClientConfig, CommunicationError> {
let mut root_store = RootCertStore::empty();
// Load native certs
let certs = rustls_native_certs::load_native_certs().certs;
for cert in certs {
root_store.add(cert).ok();
}
let mut tls_config = RustlsClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
tls_config.alpn_protocols = vec![b"h3".to_vec()];
Ok(ClientConfig::builder()
.with_bind_default()
.with_custom_tls(tls_config)
.keep_alive_interval(policy.keep_alive_interval)
.max_idle_timeout(policy.max_idle_timeout)
.map_err(|e| CommunicationError::Other(e.to_string()))?
.build())
}

519
transport/src/connection.rs Normal file
View file

@ -0,0 +1,519 @@
use crate::ConnectionHandle;
use mtp_codec::CommunicationValue;
use mtp_common::CommunicationError;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use tokio::time::{Duration, sleep, timeout};
use wtransport::Connection;
const APPLICATION_CLOSE_REASON: &str = "mtp-close";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SendMode {
PersistentStream,
SingleStreamPerMessage,
}
#[derive(Debug, Clone)]
pub struct Policy {
pub send_mode: SendMode,
pub max_message_size: u64,
pub close_frame_len: u32,
pub application_close_code: u32,
pub open_stream_timeout: Duration,
pub write_timeout: Duration,
pub accept_stream_timeout: Duration,
pub read_timeout: Duration,
pub keep_alive_interval: Option<Duration>,
pub max_idle_timeout: Option<Duration>,
pub force_close_delay: Duration,
pub max_transient_recv_errors: usize,
pub transient_recv_backoff: Duration,
pub receiver_queue_capacity: usize,
}
impl Default for Policy {
fn default() -> Self {
Self {
send_mode: SendMode::PersistentStream,
max_message_size: 1_000_000_000,
close_frame_len: u32::MAX,
application_close_code: 0,
open_stream_timeout: Duration::from_millis(2_000),
write_timeout: Duration::from_millis(2_000),
accept_stream_timeout: Duration::from_millis(10_000),
read_timeout: Duration::from_millis(30_000),
keep_alive_interval: Some(Duration::from_secs(3)),
max_idle_timeout: Some(Duration::from_secs(30)),
force_close_delay: Duration::from_millis(300),
max_transient_recv_errors: 20,
transient_recv_backoff: Duration::from_millis(100),
receiver_queue_capacity: 1000,
}
}
}
#[allow(unused)]
enum ReceivedFrame {
Message(CommunicationValue),
ClosedByPeer,
Idle,
}
pub struct Sender {
send_guard: Mutex<()>,
stream_guard: Mutex<Option<wtransport::SendStream>>,
handle: Arc<ConnectionHandle>,
connection: Connection,
policy: Arc<Policy>,
}
impl Sender {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
Self {
send_guard: Mutex::new(()),
stream_guard: Mutex::new(None),
handle,
connection,
policy,
}
}
async fn write_frame(
stream: &mut wtransport::SendStream,
data: &CommunicationValue,
policy: &Policy,
) -> Result<(), CommunicationError> {
let bytes = data.to_bytes();
if bytes.len() as u64 > policy.max_message_size
|| bytes.len() as u64 >= policy.close_frame_len as u64
{
return Err(CommunicationError::MessageTooLarge);
}
use tokio::io::AsyncWriteExt;
timeout(policy.write_timeout, stream.write_u32(bytes.len() as u32))
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(|_| CommunicationError::StreamError)?;
timeout(policy.write_timeout, stream.write_all(&bytes))
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(CommunicationError::from)?;
Ok(())
}
fn normalize_send_error(error: CommunicationError) -> CommunicationError {
match error {
CommunicationError::ConnectionError(_)
| CommunicationError::ReadExactError(_)
| CommunicationError::ClosedError(_)
| CommunicationError::StreamReadExactError(_)
| CommunicationError::StreamError => CommunicationError::StreamClosed,
other => other,
}
}
async fn open_uni_stream(
conn: &Connection,
policy: &Policy,
) -> Result<wtransport::SendStream, CommunicationError> {
let opening = timeout(policy.open_stream_timeout, conn.open_uni())
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(CommunicationError::ConnectionError)?;
let stream = timeout(policy.open_stream_timeout, opening)
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(|_| CommunicationError::StreamError)?;
Ok(stream)
}
async fn ensure_stream<'a>(
conn: &Connection,
stream_opt: &'a mut Option<wtransport::SendStream>,
policy: &Policy,
) -> Result<&'a mut wtransport::SendStream, CommunicationError> {
if stream_opt.is_none() {
*stream_opt = Some(Self::open_uni_stream(conn, policy).await?);
}
match stream_opt.as_mut() {
Some(stream) => Ok(stream),
_ => Err(CommunicationError::StreamError),
}
}
async fn send_on_persistent_stream(
conn: &Connection,
stream_opt: &mut Option<wtransport::SendStream>,
data: &CommunicationValue,
policy: &Policy,
) -> Result<(), CommunicationError> {
let mut tries = 0usize;
loop {
if conn.quic_connection().close_reason().is_some() {
return Err(CommunicationError::StreamClosed);
}
let res = {
let stream = Self::ensure_stream(conn, stream_opt, policy).await?;
Self::write_frame(stream, data, policy).await
};
if res.is_ok() {
return Ok(());
}
*stream_opt = None;
tries += 1;
if tries >= 4 {
let stream = Self::ensure_stream(conn, stream_opt, policy).await?;
return Self::write_frame(stream, data, policy).await;
}
tokio::time::sleep(std::time::Duration::from_millis(20 * tries as u64)).await;
}
}
async fn send_on_single_stream(
conn: &Connection,
data: &CommunicationValue,
policy: &Policy,
) -> Result<(), CommunicationError> {
let mut stream = Self::open_uni_stream(conn, policy).await?;
Self::write_frame(&mut stream, data, policy).await?;
timeout(policy.write_timeout, stream.finish())
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(|_| CommunicationError::StreamError)?;
Ok(())
}
async fn send_close_frame(
conn: &Connection,
policy: &Policy,
) -> Result<(), CommunicationError> {
let mut stream = Self::open_uni_stream(conn, policy).await?;
use tokio::io::AsyncWriteExt;
timeout(
policy.write_timeout,
stream.write_u32(policy.close_frame_len),
)
.await
.map_err(|_| CommunicationError::StreamError)?
.map_err(|_| CommunicationError::StreamError)?;
if let Err(e) = timeout(policy.write_timeout, stream.finish())
.await
.map_err(|_| CommunicationError::StreamError)?
{
println!("[Sender] close frame finish failed: {e}");
}
Ok(())
}
pub async fn send(&self, data: &CommunicationValue) -> Result<(), CommunicationError> {
if self.handle.is_closed() {
return Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::UseAfterClosed));
}
let _send_lock = self.send_guard.lock().await;
if self.connection.quic_connection().close_reason().is_some() {
let reason = self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed);
self.handle.close(Some(reason.clone()));
return Err(reason);
}
let res = match self.policy.send_mode {
SendMode::PersistentStream => {
let mut stream_opt = self.stream_guard.lock().await;
let r = Self::send_on_persistent_stream(
&self.connection,
&mut stream_opt,
data,
&self.policy,
)
.await;
if r.is_err() {
*stream_opt = None;
}
r
}
SendMode::SingleStreamPerMessage => {
Self::send_on_single_stream(&self.connection, data, &self.policy).await
}
};
match res {
Ok(()) => Ok(()),
Err(e) => {
let normalized = Self::normalize_send_error(e);
if self.connection.quic_connection().close_reason().is_some()
|| matches!(normalized, CommunicationError::StreamClosed)
{
self.handle.close(Some(normalized.clone()));
}
Err(normalized)
}
}
}
pub fn handle(&self) -> &Arc<ConnectionHandle> {
&self.handle
}
pub fn close(&self) {
let connection = self.connection.clone();
let handle = self.handle.clone();
let policy = self.policy.clone();
tokio::spawn(async move {
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
handle.close(Some(CommunicationError::StreamClosed));
return;
}
let _ = Self::send_close_frame(&connection, &policy).await;
handle.close(Some(CommunicationError::StreamClosed));
sleep(policy.force_close_delay).await;
if connection.quic_connection().close_reason().is_none() {
connection.quic_connection().close(
policy.application_close_code.into(),
APPLICATION_CLOSE_REASON.as_bytes(),
);
}
});
}
pub fn is_open(&self) -> bool {
self.handle.is_open()
}
pub fn is_closed(&self) -> bool {
self.handle.is_closed()
}
pub fn close_reason(&self) -> Option<CommunicationError> {
self.handle.close_reason()
}
}
pub struct Receiver {
rx: Mutex<mpsc::Receiver<Result<CommunicationValue, CommunicationError>>>,
_accept_task: tokio::task::JoinHandle<()>,
handle: Arc<ConnectionHandle>,
}
impl Receiver {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
let (tx, rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(
policy.receiver_queue_capacity,
);
let conn_handle = handle.clone();
let accept_connection = connection.clone();
let accept_policy = policy.clone();
let accept_task = tokio::spawn(async move {
let mut close_rx = conn_handle.subscribe_close();
loop {
tokio::select! {
_ = close_rx.changed() => {
if close_rx.borrow().is_some() {
break;
}
}
accepted = timeout(
accept_policy.accept_stream_timeout,
accept_connection.accept_uni()
) => {
match accepted {
Ok(Ok(stream)) => {
let tx_stream = tx.clone();
let stream_handle = conn_handle.clone();
let stream_policy = accept_policy.clone();
tokio::spawn(async move {
let mut s = stream;
loop {
match Self::read_one_frame(&mut s, &stream_policy).await {
Ok(ReceivedFrame::Message(msg)) => {
if tx_stream.send(Ok(msg)).await.is_err() {
break;
}
}
Ok(ReceivedFrame::ClosedByPeer) => {
let close_error = CommunicationError::StreamClosed;
let _ = tx_stream.send(Err(close_error.clone())).await;
stream_handle.close(Some(close_error));
break;
}
Ok(ReceivedFrame::Idle) => {
break;
}
Err(e) => {
let close_error = match e {
CommunicationError::ConnectionError(_)
| CommunicationError::ReadExactError(_)
| CommunicationError::ClosedError(_)
| CommunicationError::StreamReadExactError(_)
| CommunicationError::StreamError => CommunicationError::StreamClosed,
other => other,
};
let _ = tx_stream.send(Err(close_error.clone())).await;
stream_handle.close(Some(close_error));
break;
}
}
}
});
}
Ok(Err(_e)) => {
// A connection error from accept_uni means the connection is permanently closed.
let close_error = CommunicationError::StreamClosed;
let _ = tx.send(Err(close_error.clone())).await;
conn_handle.close(Some(close_error));
break;
}
Err(_) => {
if accept_connection.quic_connection().close_reason().is_some() {
let close_error = CommunicationError::StreamClosed;
let _ = tx.send(Err(close_error.clone())).await;
conn_handle.close(Some(close_error));
break;
}
}
}
}
}
if conn_handle.is_closed() {
break;
}
}
});
Self {
rx: Mutex::new(rx),
_accept_task: accept_task,
handle,
}
}
async fn read_one_frame(
stream: &mut wtransport::RecvStream,
policy: &Policy,
) -> Result<ReceivedFrame, CommunicationError> {
use std::io::ErrorKind;
use tokio::io::AsyncReadExt;
let mut attempts = 0;
let len = loop {
match stream.read_u32().await {
Ok(len) => break len,
Err(e) => {
if e.kind() == ErrorKind::Interrupted && attempts < 3 {
attempts += 1;
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
continue;
}
if e.kind() == ErrorKind::UnexpectedEof {
return Ok(ReceivedFrame::Idle);
}
println!("[Receiver] read_u32 failed: {e}");
return Err(CommunicationError::StreamError);
}
}
};
if len == policy.close_frame_len {
return Ok(ReceivedFrame::ClosedByPeer);
}
let len = len as usize;
if len as u64 > policy.max_message_size {
return Err(CommunicationError::MessageTooLarge);
}
let mut buf = vec![0u8; len];
match timeout(policy.read_timeout, stream.read_exact(&mut buf)).await {
Ok(Ok(())) => {}
Ok(Err(e)) => match timeout(policy.read_timeout, stream.read_exact(&mut buf)).await {
Ok(Ok(())) => {}
_ => return Err(e.into()),
},
Err(_) => {
println!("[Receiver] read_exact timed out (len={})", len);
return Err(CommunicationError::StreamError);
}
}
let message = CommunicationValue::from_bytes(&buf)
.ok_or(CommunicationError::ParseCommunicationValue)?;
Ok(ReceivedFrame::Message(message))
}
pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
if self.handle.is_closed() {
return Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed));
}
let mut rx = self.rx.lock().await;
match rx.recv().await {
Some(result) => result,
_ => Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed)),
}
}
pub fn handle(&self) -> &Arc<ConnectionHandle> {
&self.handle
}
pub fn close(&self) {
self.handle.close(None);
}
pub fn is_open(&self) -> bool {
self.handle.is_open()
}
pub fn is_closed(&self) -> bool {
self.handle.is_closed()
}
pub fn close_reason(&self) -> Option<CommunicationError> {
self.handle.close_reason()
}
}

65
transport/src/connection_handle.rs Normal file
View file

@ -0,0 +1,65 @@
use mtp_common::CommunicationError;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use tokio::sync::watch;
#[derive(Debug)]
pub struct ConnectionHandle {
closed: AtomicBool,
close_tx: watch::Sender<Option<CommunicationError>>,
close_rx: watch::Receiver<Option<CommunicationError>>,
}
impl ConnectionHandle {
pub fn new() -> Self {
let (close_tx, close_rx) = watch::channel(None);
Self {
closed: AtomicBool::new(false),
close_tx,
close_rx,
}
}
pub fn is_open(&self) -> bool {
!self.closed.load(Ordering::SeqCst)
}
pub fn is_closed(&self) -> bool {
self.closed.load(Ordering::SeqCst)
}
pub fn close(&self, reason: Option<CommunicationError>) {
if !self.closed.swap(true, Ordering::SeqCst) {
let _ = self.close_tx.send(reason);
}
}
pub fn close_reason(&self) -> Option<CommunicationError> {
self.close_rx.borrow().clone()
}
pub fn subscribe_close(&self) -> watch::Receiver<Option<CommunicationError>> {
self.close_rx.clone()
}
pub fn close_with_error(&self, error: CommunicationError) {
self.close(Some(error));
}
pub async fn wait_closed(self: Arc<Self>) -> Option<CommunicationError> {
let mut rx = self.subscribe_close();
if self.is_closed() {
return rx.borrow().clone();
}
let _ = rx.changed().await.ok()?;
rx.borrow().clone()
}
}
impl Default for ConnectionHandle {
fn default() -> Self {
Self::new()
}
}

123
transport/src/host.rs Normal file
View file

@ -0,0 +1,123 @@
use crate::{ConnectionHandle, Policy, Receiver, Sender};
use mtp_common::CommunicationError;
use rustls::pki_types::{PrivateKeyDer, pem::PemObject};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use wtransport::{Connection as WTConnection, Endpoint, ServerConfig};
pub struct Host {
incoming: tokio::sync::mpsc::Receiver<(Sender, Receiver)>,
local_addr: std::net::SocketAddr,
_task: tokio::task::JoinHandle<()>,
}
impl Host {
pub async fn next(&mut self) -> Option<(Sender, Receiver)> {
self.incoming.recv().await
}
pub fn local_addr(&self) -> std::net::SocketAddr {
self.local_addr
}
}
pub async fn host(
port: u16,
cert_pem: Vec<u8>,
key_pem: Vec<u8>,
policy: Policy,
) -> Result<Host, CommunicationError> {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let server_config = configure_server(port, cert_pem, key_pem, &policy).await?;
let endpoint = Endpoint::server(server_config)
.map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?;
let local_addr = endpoint
.local_addr()
.map_err(|e| CommunicationError::Other(e.to_string()))?;
let (incoming_tx, incoming_rx) = tokio::sync::mpsc::channel(16);
let policy = Arc::new(policy);
let task = tokio::spawn(async move {
loop {
let incoming_session = endpoint.accept().await;
let request = match incoming_session.await {
Ok(req) => req,
Err(_) => {
continue;
}
};
let connection = match request.accept().await {
Ok(conn) => conn,
Err(_) => {
continue;
}
};
let incoming_tx = incoming_tx.clone();
let policy = policy.clone();
tokio::spawn(handle_connection(connection, incoming_tx, policy));
}
});
Ok(Host {
incoming: incoming_rx,
local_addr,
_task: task,
})
}
async fn handle_connection(
connection: WTConnection,
tx: tokio::sync::mpsc::Sender<(Sender, Receiver)>,
policy: Arc<Policy>,
) {
let handle = Arc::new(ConnectionHandle::new());
let sender = Sender::new(connection.clone(), handle.clone(), policy.clone());
let receiver = Receiver::new(connection, handle, policy);
let _ = tx.send((sender, receiver)).await;
}
async fn configure_server(
port: u16,
cert_pem: Vec<u8>,
key_pem: Vec<u8>,
policy: &Policy,
) -> Result<ServerConfig, CommunicationError> {
let cert_chain = rustls::pki_types::CertificateDer::pem_slice_iter(&cert_pem)
.collect::<Result<Vec<_>, _>>()
.map_err(|_| CommunicationError::CertificateLoadFailed)?;
let key = PrivateKeyDer::from_pem_slice(&key_pem)
.map_err(|_| CommunicationError::CertificateParseFailed)?;
let mut tls_config = rustls::ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key)
.map_err(|_| CommunicationError::CertificateLoadFailed)?;
tls_config.alpn_protocols = vec![b"h3".to_vec()];
let bind_ip = std::env::var("mtp_BIND")
.ok()
.and_then(|s| if s.is_empty() { None } else { Some(s) })
.unwrap_or_else(|| "::".to_string())
.parse::<IpAddr>()?;
let bind_addr = SocketAddr::new(bind_ip, port);
let server_config = ServerConfig::builder()
.with_bind_address(bind_addr)
.with_custom_tls(tls_config)
.keep_alive_interval(policy.keep_alive_interval)
.max_idle_timeout(policy.max_idle_timeout)
.map_err(|e| CommunicationError::Other(e.to_string()))?
.build();
Ok(server_config)
}

22
transport/src/lib.rs
View file

@ -1,14 +1,12 @@
pub fn add(left: u64, right: u64) -> u64 {
left + right
}
pub mod client;
pub mod connection;
pub mod connection_handle;
#[cfg(test)]
mod tests {
use super::*;
pub use client::connect;
pub use connection::{Policy, Receiver, SendMode, Sender};
pub use connection_handle::ConnectionHandle;
#[test]
fn it_works() {
let result = add(2, 2);
assert_eq!(result, 4);
}
}
#[cfg(feature = "host")]
pub mod host;
#[cfg(feature = "host")]
pub use host::{Host, host};