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

[Fix] Harden MTP codec, transport, and SDK security

Browse source
This commit is contained in:
Alex Emmet 2026-08-18 20:57:45 +02:00
commit a7e804c603
No known key found for this signature in database
73 changed files with 11892 additions and 5756 deletions
Download patch file Download diff file Expand all files Collapse all files

427
transport/src/connection.rs
View file

@ -1,12 +1,14 @@
use crate::ConnectionHandle;
use crate::framing::RetryClassifier;
#[cfg(feature = "pipes")]
use crate::pipe::PipeReader;
use mtp_codec::{CommunicationValue, DecodeLimits, TypeMap};
use mtp_codec::{CommunicationValue, DecodeError, DecodeLimits, EncodeLimits, TypeMap};
use mtp_common::CommunicationError;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{Mutex, Notify, RwLock, Semaphore, mpsc};
use tokio::time::{Duration, sleep, timeout};
use tokio::time::{Duration, Instant, sleep, timeout, timeout_at};
use tracing::{debug, info, instrument, trace, warn};
use wtransport::Connection;
@ -19,6 +21,63 @@ pub enum TransportEvent<R = wtransport::RecvStream> {
const APPLICATION_CLOSE_REASON: &str = "mtp-close";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DecodeRejectionClass {
Malformed,
ResourceLimit,
DuplicateField,
}
pub fn classify_decode_error(error: &DecodeError) -> DecodeRejectionClass {
match error {
DecodeError::MalformedEncoding => DecodeRejectionClass::Malformed,
DecodeError::DepthLimit
| DecodeError::ValueCountLimit
| DecodeError::BlobLimit
| DecodeError::AllocationLimit
| DecodeError::RecipientLimit => DecodeRejectionClass::ResourceLimit,
DecodeError::DuplicateField => DecodeRejectionClass::DuplicateField,
}
}
#[derive(Debug, Default)]
pub struct DecodeRejectionCounters {
malformed: AtomicU64,
resource_limit: AtomicU64,
duplicate_field: AtomicU64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct DecodeRejectionCounts {
pub malformed: u64,
pub resource_limit: u64,
pub duplicate_field: u64,
}
impl DecodeRejectionCounters {
pub(crate) fn record(&self, error: &DecodeError) {
match classify_decode_error(error) {
DecodeRejectionClass::Malformed => {
self.malformed.fetch_add(1, Ordering::Relaxed);
}
DecodeRejectionClass::ResourceLimit => {
self.resource_limit.fetch_add(1, Ordering::Relaxed);
}
DecodeRejectionClass::DuplicateField => {
self.duplicate_field.fetch_add(1, Ordering::Relaxed);
}
}
}
pub(crate) fn snapshot(&self) -> DecodeRejectionCounts {
DecodeRejectionCounts {
malformed: self.malformed.load(Ordering::Relaxed),
resource_limit: self.resource_limit.load(Ordering::Relaxed),
duplicate_field: self.duplicate_field.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SendMode {
PersistentStream,
@ -135,27 +194,62 @@ impl Policy {
}
}
/// A validated policy snapshot used after a public [`Policy`] crosses into a
/// transport implementation. `Policy` intentionally remains a plain public
/// struct for source compatibility, so callers can construct it directly and
/// bypass builder methods. Every transport constructor takes this snapshot
/// before creating channels or semaphores.
#[derive(Debug, Clone, Copy)]
pub(crate) struct RuntimePolicy(Policy);
impl RuntimePolicy {
pub(crate) fn from_public(policy: &Policy) -> Self {
let mut policy = *policy;
policy.receiver_queue_capacity = policy.receiver_queue_capacity.max(1);
policy.max_concurrent_stream_tasks = policy.max_concurrent_stream_tasks.max(1);
Self(policy)
}
}
impl Deref for RuntimePolicy {
type Target = Policy;
fn deref(&self) -> &Self::Target {
&self.0
}
}
enum ReceivedFrame {
Message(CommunicationValue),
ClosedByPeer,
Idle,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SenderState {
Open,
Closing,
Closed,
}
#[derive(Clone)]
pub struct Sender {
send_guard: Arc<Mutex<()>>,
stream_guard: Arc<Mutex<Option<wtransport::SendStream>>>,
state: Arc<Mutex<SenderState>>,
handle: Arc<ConnectionHandle>,
connection: Connection,
policy: Arc<Policy>,
policy: Arc<RuntimePolicy>,
type_map: Arc<RwLock<TypeMap>>,
}
impl Sender {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
let policy = Arc::new(RuntimePolicy::from_public(&policy));
Self {
send_guard: Arc::new(Mutex::new(())),
stream_guard: Arc::new(Mutex::new(None)),
state: Arc::new(Mutex::new(SenderState::Open)),
handle,
connection,
policy,
@ -174,7 +268,11 @@ impl Sender {
data: &CommunicationValue,
policy: &Policy,
) -> Result<(), CommunicationError> {
let bytes = data.to_bytes().map_err(|_| CommunicationError::Encode)?;
let bytes = data
.to_bytes_with_limits(EncodeLimits::for_transport_message_size(
policy.max_message_size,
))
.map_err(|_| CommunicationError::Encode)?;
if bytes.len() as u64 > policy.max_message_size
|| bytes.len() as u64 >= policy.close_frame_len as u64
{
@ -269,11 +367,11 @@ impl Sender {
return Ok(());
}
let err = res.err().unwrap_or(CommunicationError::StreamError);
if !matches!(
err,
CommunicationError::StreamError | CommunicationError::StreamClosed
) {
let err = match res {
Ok(()) => return Ok(()),
Err(error) => error,
};
if !RetryClassifier::retry_persistent_stream(&err) {
return Err(err);
}
*stream_opt = None;
@ -356,20 +454,24 @@ impl Sender {
#[instrument(skip(self, data), level = "trace")]
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;
{
let state = self.state.lock().await;
if *state != SenderState::Open {
return Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed));
}
}
if self.connection.quic_connection().close_reason().is_some() {
let reason = self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed);
*self.state.lock().await = SenderState::Closed;
self.handle.close(Some(reason.clone()));
return Err(reason);
}
@ -402,6 +504,7 @@ impl Sender {
if self.connection.quic_connection().close_reason().is_some()
|| matches!(normalized, CommunicationError::StreamClosed)
{
*self.state.lock().await = SenderState::Closed;
self.handle.close(Some(normalized.clone()));
}
@ -413,6 +516,12 @@ impl Sender {
#[instrument(skip(self), level = "trace")]
pub async fn finish_stream(&self) -> Result<(), CommunicationError> {
let _send_lock = self.send_guard.lock().await;
if *self.state.lock().await != SenderState::Open {
return Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::StreamClosed));
}
let mut stream_opt = self.stream_guard.lock().await;
if let Some(mut stream) = stream_opt.take() {
match timeout(self.policy.write_timeout, stream.finish()).await {
@ -445,11 +554,12 @@ impl Sender {
pipe_id: u32,
description: &str,
) -> Result<crate::pipe::PipeWriter, CommunicationError> {
if self.handle.is_closed() {
let _send_lock = self.send_guard.lock().await;
if *self.state.lock().await != SenderState::Open {
return Err(self
.handle
.close_reason()
.unwrap_or(CommunicationError::UseAfterClosed));
.unwrap_or(CommunicationError::StreamClosed));
}
if self.connection.quic_connection().close_reason().is_some() {
@ -486,31 +596,47 @@ impl Sender {
let handle = self.handle.clone();
let policy = self.policy.clone();
let stream_guard = self.stream_guard.clone();
let send_guard = self.send_guard.clone();
let state = self.state.clone();
tokio::spawn(async move {
let _send_lock = send_guard.lock().await;
{
let mut sender_state = state.lock().await;
if *sender_state != SenderState::Open {
return;
}
*sender_state = SenderState::Closing;
}
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
*state.lock().await = SenderState::Closed;
handle.close(Some(CommunicationError::StreamClosed));
return;
}
if let Some(mut stream) = stream_guard.lock().await.take() {
match timeout(policy.write_timeout, stream.finish()).await {
Ok(Ok(())) => {}
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => warn!(
"[Sender] persistent stream finish failed: peer sent STOP_SENDING (error code {code})"
),
Ok(Err(e)) => {
warn!("[Sender] persistent stream finish failed: {e}")
{
if let Some(mut stream) = stream_guard.lock().await.take() {
match timeout(policy.write_timeout, stream.finish()).await {
Ok(Ok(())) => {}
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => warn!(
"[Sender] persistent stream finish failed: peer sent STOP_SENDING (error code {code})"
),
Ok(Err(e)) => {
warn!("[Sender] persistent stream finish failed: {e}")
}
Err(_) => warn!("[Sender] persistent stream finish timed out"),
}
Err(_) => warn!("[Sender] persistent stream finish timed out"),
}
}
let _ = Self::send_close_frame(&connection, &policy).await;
*state.lock().await = SenderState::Closed;
handle.close(Some(CommunicationError::StreamClosed));
info!(target = "mtp.transport", "connection closed");
drop(_send_lock);
sleep(policy.force_close_delay).await;
if connection.quic_connection().close_reason().is_none() {
connection.quic_connection().close(
@ -528,35 +654,48 @@ impl Sender {
let connection = self.connection.clone();
let handle = self.handle.clone();
let policy = self.policy.clone();
let mut stream_opt = self.stream_guard.lock().await;
let _send_lock = self.send_guard.lock().await;
{
let mut state = self.state.lock().await;
if *state != SenderState::Open {
return;
}
*state = SenderState::Closing;
}
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
*self.state.lock().await = SenderState::Closed;
handle.close(Some(CommunicationError::StreamClosed));
return;
}
if let Some(mut stream) = stream_opt.take() {
let close_bytes = policy.close_frame_len.to_be_bytes();
let close_write = async {
stream.write_all(&close_bytes).await?;
stream.finish().await
};
{
let mut stream_opt = self.stream_guard.lock().await;
if let Some(mut stream) = stream_opt.take() {
let close_bytes = policy.close_frame_len.to_be_bytes();
let close_write = async {
stream.write_all(&close_bytes).await?;
stream.finish().await
};
match timeout(policy.write_timeout, close_write).await {
Ok(Ok(())) => {}
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => {
warn!("[Sender] close failed: peer sent STOP_SENDING (error code {code})")
match timeout(policy.write_timeout, close_write).await {
Ok(Ok(())) => {}
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => {
warn!("[Sender] close failed: peer sent STOP_SENDING (error code {code})")
}
Ok(Err(e)) => warn!("[Sender] close failed: {e}"),
Err(_) => warn!("[Sender] close timed out"),
}
Ok(Err(e)) => warn!("[Sender] close failed: {e}"),
Err(_) => warn!("[Sender] close timed out"),
} else {
let _ = Self::send_close_frame(&connection, &policy).await;
}
} else {
let _ = Self::send_close_frame(&connection, &policy).await;
}
*self.state.lock().await = SenderState::Closed;
handle.close(Some(CommunicationError::StreamClosed));
info!(target = "mtp.transport", "connection closed");
drop(_send_lock);
sleep(policy.force_close_delay).await;
if connection.quic_connection().close_reason().is_none() {
connection.quic_connection().close(
@ -605,6 +744,7 @@ struct ReceiverInner {
queue_notify: Arc<Notify>,
max_message_size: Arc<AtomicU64>,
type_map: Arc<RwLock<TypeMap>>,
decode_rejections: Arc<DecodeRejectionCounters>,
}
impl Clone for Receiver {
@ -626,12 +766,26 @@ impl Drop for Receiver {
#[derive(Clone, Default)]
struct PingControl {
pong_sender: Option<Sender>,
pong_observer: Option<mpsc::UnboundedSender<CommunicationValue>>,
pong_observer: Option<mpsc::Sender<CommunicationValue>>,
expected_pong_id: Option<u32>,
}
impl PingControl {
fn accepts_pong(&mut self, id: Option<u32>) -> bool {
if self.expected_pong_id == id && id.is_some() {
self.expected_pong_id = None;
true
} else {
false
}
}
}
impl Receiver {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
Self::new_with_max_message_size(connection, handle, policy.clone(), policy.max_message_size)
let policy = Arc::new(RuntimePolicy::from_public(&policy));
let max_message_size = policy.max_message_size;
Self::new_with_max_message_size(connection, handle, policy, max_message_size)
}
#[cfg(feature = "host")]
@ -640,6 +794,7 @@ impl Receiver {
handle: Arc<ConnectionHandle>,
policy: Arc<Policy>,
) -> Self {
let policy = Arc::new(RuntimePolicy::from_public(&policy));
let initial_max = policy
.handshake_max_message_size
.min(policy.max_message_size);
@ -649,7 +804,7 @@ impl Receiver {
fn new_with_max_message_size(
connection: Connection,
handle: Arc<ConnectionHandle>,
policy: Arc<Policy>,
policy: Arc<RuntimePolicy>,
initial_max_message_size: u64,
) -> Self {
#[cfg(feature = "pipes")]
@ -674,6 +829,8 @@ impl Receiver {
let accept_max_message_size = max_message_size.clone();
let type_map = Arc::new(RwLock::new(TypeMap::latest()));
let accept_type_map = type_map.clone();
let decode_rejections = Arc::new(DecodeRejectionCounters::default());
let accept_decode_rejections = decode_rejections.clone();
let stream_limit = Arc::new(Semaphore::new(policy.max_concurrent_stream_tasks.max(1)));
let accept_stream_limit = stream_limit.clone();
debug!(
@ -742,6 +899,7 @@ impl Receiver {
let stream_ping_control = accept_ping_control.clone();
let stream_max_message_size = accept_max_message_size.clone();
let stream_type_map = accept_type_map.clone();
let stream_decode_rejections = accept_decode_rejections.clone();
tokio::spawn(async move {
let _permit = permit;
@ -761,7 +919,14 @@ impl Receiver {
}
let frame_limit = stream_max_message_size.load(Ordering::Relaxed);
match Self::read_one_frame(&mut s, &stream_policy, frame_limit).await {
match Self::read_one_frame(
&mut s,
&stream_policy,
frame_limit,
&stream_decode_rejections,
)
.await
{
Ok(ReceivedFrame::Message(mut msg)) => {
let negotiated_type_map =
stream_type_map.read().await.clone();
@ -805,24 +970,17 @@ impl Receiver {
}
}
let control = {
let control = stream_ping_control.read().await;
if msg.is_type(mtp_codec::CommunicationType::Ping) {
control
.pong_sender
.clone()
.map(|sender| (Some(sender), None))
} else if msg.is_type(mtp_codec::CommunicationType::Pong) {
control
.pong_observer
.clone()
.map(|observer| (None, Some(observer)))
} else {
None
}
let pong_sender = if msg.is_type(mtp_codec::CommunicationType::Ping) {
stream_ping_control
.read()
.await
.pong_sender
.clone()
} else {
None
};
if let Some((Some(sender), _)) = control {
if let Some(sender) = pong_sender {
let mut pong = CommunicationValue::new_with_type_map(
mtp_codec::CommunicationType::Pong,
&negotiated_type_map,
@ -844,8 +1002,18 @@ impl Receiver {
continue;
}
if let Some((_, Some(observer))) = control {
let _ = observer.send(msg);
if msg.is_type(mtp_codec::CommunicationType::Pong) {
let observer = {
let mut control = stream_ping_control.write().await;
if control.accepts_pong(msg.id()) {
control.pong_observer.clone()
} else {
None
}
};
if let Some(observer) = observer {
let _ = observer.try_send(msg);
}
continue;
}
@ -942,6 +1110,7 @@ impl Receiver {
queue_notify,
max_message_size,
type_map,
decode_rejections,
}),
}
}
@ -958,6 +1127,14 @@ impl Receiver {
*self.inner.type_map.write().await = type_map.clone();
}
/// Return local counts for frames rejected by the structured decoder.
///
/// These counters are intentionally local-only; peers continue to receive
/// the generic protocol parse failure.
pub fn decode_rejection_counts(&self) -> DecodeRejectionCounts {
self.inner.decode_rejections.snapshot()
}
/* Respond to reserved Ping frames without exposing them to application I/O. */
pub fn respond_to_pings(&self, sender: Sender) {
if let Ok(mut control) = self.inner.ping_control.try_write() {
@ -967,16 +1144,21 @@ impl Receiver {
}
}
/* Route reserved Pong frames to a connection-level observer. */
pub async fn observe_pongs(&self, observer: mpsc::UnboundedSender<CommunicationValue>) {
/* Route only the currently expected reserved Pong through a bounded observer. */
pub async fn observe_pongs_bounded(&self, observer: mpsc::Sender<CommunicationValue>) {
self.inner.ping_control.write().await.pong_observer = Some(observer);
}
#[instrument(skip(stream, policy), level = "trace")]
pub async fn set_expected_pong_id(&self, expected_pong_id: Option<u32>) {
self.inner.ping_control.write().await.expected_pong_id = expected_pong_id;
}
#[instrument(skip(stream, policy, decode_rejections), level = "trace")]
async fn read_one_frame(
stream: &mut wtransport::RecvStream,
policy: &Policy,
policy: &RuntimePolicy,
max_message_size: u64,
decode_rejections: &DecodeRejectionCounters,
) -> Result<ReceivedFrame, CommunicationError> {
use wtransport::error::{StreamReadError, StreamReadExactError};
@ -1011,6 +1193,7 @@ impl Receiver {
if len == policy.close_frame_len {
return Ok(ReceivedFrame::ClosedByPeer);
}
let deadline = Instant::now() + policy.read_timeout;
let body_len = len as usize;
let frame_len = body_len
@ -1020,29 +1203,29 @@ impl Receiver {
return Err(CommunicationError::MessageTooLarge);
}
// Grow in bounded chunks instead of trusting the peer's length prefix
// enough to allocate the complete frame up front.
let mut buf = Vec::new();
buf.try_reserve(body_len.min(16 * 1024))
// The length has already been checked against the admitted frame
// limit, so reserve one bounded framing buffer and decode it without a
// second prefix-plus-body allocation/copy.
let mut frame = Vec::new();
frame
.try_reserve_exact(frame_len)
.map_err(|_| CommunicationError::MessageTooLarge)?;
while buf.len() < body_len {
let chunk_len = (body_len - buf.len()).min(16 * 1024);
let mut chunk = [0u8; 16 * 1024];
match timeout(
policy.read_timeout,
stream.read_exact(&mut chunk[..chunk_len]),
frame.extend_from_slice(&len_buf);
frame.resize(frame_len, 0);
let mut body_offset = 4usize;
while body_offset < frame_len {
let chunk_len = (frame_len - body_offset).min(16 * 1024);
match timeout_at(
deadline,
stream.read_exact(&mut frame[body_offset..body_offset + chunk_len]),
)
.await
{
Ok(Ok(())) => {
buf.try_reserve(chunk_len)
.map_err(|_| CommunicationError::MessageTooLarge)?;
buf.extend_from_slice(&chunk[..chunk_len]);
}
Ok(Ok(())) => body_offset += chunk_len,
Ok(Err(StreamReadExactError::FinishedEarly(n))) => {
warn!(
"[Receiver] body read ended early ({}/{body_len} bytes): stream closed by peer",
buf.len() + n
body_offset.saturating_sub(4) + n
);
return Err(CommunicationError::StreamError);
}
@ -1063,14 +1246,21 @@ impl Receiver {
}
}
let mut frame = Vec::with_capacity(frame_len);
frame.extend_from_slice(&len_buf);
frame.extend_from_slice(&buf);
let message = CommunicationValue::from_bytes_with_limits(
let message = match CommunicationValue::try_from_bytes_with_limits(
&frame,
DecodeLimits::for_transport_message_size(max_message_size),
)
.map_err(|_| CommunicationError::ParseCommunicationValue)?;
) {
Ok(message) => message,
Err(error) => {
decode_rejections.record(&error);
warn!(
?error,
class = ?classify_decode_error(&error),
"[Receiver] rejected frame during bounded decode"
);
return Err(CommunicationError::ParseCommunicationValue);
}
};
Ok(ReceivedFrame::Message(message))
}
@ -1267,4 +1457,61 @@ mod tests {
let debug_str = format!("{:?}", p);
assert!(debug_str.contains("Policy"));
}
#[test]
fn runtime_policy_normalizes_zero_channel_and_task_limits() {
let mut policy = Policy::default();
policy.receiver_queue_capacity = 0;
policy.max_concurrent_stream_tasks = 0;
let runtime = RuntimePolicy::from_public(&policy);
assert_eq!(runtime.receiver_queue_capacity, 1);
assert_eq!(runtime.max_concurrent_stream_tasks, 1);
assert_eq!(policy.receiver_queue_capacity, 0);
assert_eq!(policy.max_concurrent_stream_tasks, 0);
}
#[test]
fn ping_control_accepts_only_the_current_expected_id() {
let mut control = PingControl {
expected_pong_id: Some(7),
..PingControl::default()
};
assert!(!control.accepts_pong(Some(6)));
assert_eq!(control.expected_pong_id, Some(7));
assert!(control.accepts_pong(Some(7)));
assert_eq!(control.expected_pong_id, None);
assert!(!control.accepts_pong(Some(7)));
}
#[test]
fn decode_rejection_classes_are_stable_and_counted() {
assert_eq!(
classify_decode_error(&DecodeError::MalformedEncoding),
DecodeRejectionClass::Malformed
);
assert_eq!(
classify_decode_error(&DecodeError::AllocationLimit),
DecodeRejectionClass::ResourceLimit
);
assert_eq!(
classify_decode_error(&DecodeError::DuplicateField),
DecodeRejectionClass::DuplicateField
);
let counters = DecodeRejectionCounters::default();
counters.record(&DecodeError::MalformedEncoding);
counters.record(&DecodeError::DepthLimit);
counters.record(&DecodeError::DuplicateField);
assert_eq!(
counters.snapshot(),
DecodeRejectionCounts {
malformed: 1,
resource_limit: 1,
duplicate_field: 1,
}
);
}
}

11
transport/src/connection_handle.rs
View file

@ -2,22 +2,26 @@ use mtp_common::CommunicationError;
use std::net::SocketAddr;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
atomic::{AtomicBool, AtomicU64, Ordering},
};
use tokio::sync::watch;
#[derive(Debug)]
pub struct ConnectionHandle {
connection_id: u64,
closed: AtomicBool,
close_tx: watch::Sender<Option<CommunicationError>>,
close_rx: watch::Receiver<Option<CommunicationError>>,
remote_addr: Option<SocketAddr>,
}
static NEXT_CONNECTION_ID: AtomicU64 = AtomicU64::new(1);
impl ConnectionHandle {
pub fn new() -> Self {
let (close_tx, close_rx) = watch::channel(None);
Self {
connection_id: NEXT_CONNECTION_ID.fetch_add(1, Ordering::Relaxed).max(1),
closed: AtomicBool::new(false),
close_tx,
close_rx,
@ -35,6 +39,11 @@ impl ConnectionHandle {
self.remote_addr
}
/// Stable process-local identifier for authentication-rate-limit scopes.
pub fn connection_id(&self) -> u64 {
self.connection_id
}
pub fn is_open(&self) -> bool {
!self.closed.load(Ordering::SeqCst)
}

22
transport/src/framing.rs
View file

@ -1,7 +1,21 @@
use crate::{Policy, TransportSendStream};
use mtp_codec::CommunicationValue;
use mtp_codec::{CommunicationValue, EncodeLimits};
use mtp_common::CommunicationError;
/// Classifies failures that may be recovered by replacing a persistent
/// application stream. Encoding and frame-size failures are deterministic and
/// must reach the caller without opening more streams.
pub(crate) struct RetryClassifier;
impl RetryClassifier {
pub(crate) fn retry_persistent_stream(error: &CommunicationError) -> bool {
matches!(
error,
CommunicationError::StreamError | CommunicationError::StreamClosed
)
}
}
/// Writes the canonical self-framed MTP value used by every transport.
///
/// `CommunicationValue` already begins with the four-byte body length. The
@ -12,7 +26,11 @@ pub(crate) async fn write_frame<S: TransportSendStream>(
value: &CommunicationValue,
policy: &Policy,
) -> Result<(), CommunicationError> {
let bytes = value.to_bytes().map_err(|_| CommunicationError::Encode)?;
let bytes = value
.to_bytes_with_limits(EncodeLimits::for_transport_message_size(
policy.max_message_size,
))
.map_err(|_| CommunicationError::Encode)?;
if bytes.len() as u64 > policy.max_message_size
|| bytes.len() as u64 >= policy.close_frame_len as u64
{

148
transport/src/generic_connection.rs
View file

@ -5,21 +5,23 @@
//! wrappers while the framing implementation below is shared by adapters.
use crate::{
Policy, TransportConnection, TransportRecvStream, TransportSendStream, framing::write_frame,
Policy, TransportConnection, TransportRecvStream, TransportSendStream,
connection::{DecodeRejectionCounters, RuntimePolicy, classify_decode_error},
framing::{RetryClassifier, write_frame},
};
use mtp_codec::{CommunicationValue, DecodeLimits, TypeMap};
use mtp_common::CommunicationError;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc};
use tokio::time::timeout;
use tokio::sync::{Mutex, Notify, RwLock, Semaphore, mpsc};
use tokio::time::{Instant, timeout, timeout_at};
#[cfg(feature = "pipes")]
use crate::pipe::{PipeReader, PipeWriter};
pub struct GenericSender<C: TransportConnection> {
connection: C,
policy: Arc<Policy>,
policy: Arc<RuntimePolicy>,
persistent: Arc<Mutex<Option<C::SendStream>>>,
send_lock: Arc<Mutex<()>>,
type_map: Arc<RwLock<TypeMap>>,
@ -39,6 +41,7 @@ impl<C: TransportConnection> Clone for GenericSender<C> {
impl<C: TransportConnection> GenericSender<C> {
pub fn new(connection: C, policy: Arc<Policy>) -> Self {
let policy = Arc::new(RuntimePolicy::from_public(&policy));
Self {
connection,
policy,
@ -84,20 +87,30 @@ impl<C: TransportConnection> GenericSender<C> {
if stream.is_none() {
*stream = Some(self.open().await?);
}
let result = timeout(
self.policy.write_timeout,
write_frame(stream.as_mut().unwrap(), value, &self.policy),
)
.await
.map_err(|_| CommunicationError::StreamError)
.and_then(|r| r);
let result = match stream.as_mut() {
Some(stream) => timeout(
self.policy.write_timeout,
write_frame(stream, value, &self.policy),
)
.await
.map_err(|_| CommunicationError::StreamError)
.and_then(|result| result),
None => Err(CommunicationError::StreamError),
};
if result.is_ok() {
return result;
return Ok(());
}
let error = match result {
Ok(()) => return Ok(()),
Err(error) => error,
};
if !RetryClassifier::retry_persistent_stream(&error) {
return Err(error);
}
*stream = None;
attempts += 1;
if attempts > self.policy.persistent_stream_max_retries {
return result;
return Err(error);
}
tokio::time::sleep(
self.policy.persistent_stream_retry_backoff * attempts as u32,
@ -114,6 +127,7 @@ impl<C: TransportConnection> GenericSender<C> {
pipe_id: u32,
description: &str,
) -> Result<PipeWriter<C::SendStream>, CommunicationError> {
let _send_lock = self.send_lock.lock().await;
if self.connection.close_reason().is_some() {
return Err(CommunicationError::StreamClosed);
}
@ -179,6 +193,8 @@ pub struct GenericReceiver<C: TransportConnection> {
ping_sender: Arc<RwLock<Option<GenericSender<C>>>>,
max_message_size: Arc<AtomicU64>,
type_map: Arc<RwLock<TypeMap>>,
queue_notify: Arc<Notify>,
decode_rejections: Arc<DecodeRejectionCounters>,
_accept_task: Arc<tokio::task::JoinHandle<()>>,
}
@ -192,6 +208,8 @@ impl<C: TransportConnection> Clone for GenericReceiver<C> {
ping_sender: self.ping_sender.clone(),
max_message_size: self.max_message_size.clone(),
type_map: self.type_map.clone(),
queue_notify: self.queue_notify.clone(),
decode_rejections: self.decode_rejections.clone(),
_accept_task: self._accept_task.clone(),
}
}
@ -207,6 +225,7 @@ impl<C: TransportConnection> Drop for GenericReceiver<C> {
impl<C: TransportConnection> GenericReceiver<C> {
pub fn new(connection: C, policy: Arc<Policy>) -> Self {
let policy = Arc::new(RuntimePolicy::from_public(&policy));
let (tx, rx) = mpsc::channel(policy.receiver_queue_capacity);
#[cfg(feature = "pipes")]
let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity);
@ -222,6 +241,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
let task_max_message_size = max_message_size.clone();
let type_map = Arc::new(RwLock::new(TypeMap::latest()));
let task_type_map = type_map.clone();
let queue_notify = Arc::new(Notify::new());
let task_queue_notify = queue_notify.clone();
let decode_rejections = Arc::new(DecodeRejectionCounters::default());
let task_decode_rejections = decode_rejections.clone();
let task_accept_task_tx = tx.clone();
#[cfg(feature = "pipes")]
let task_accept_task_pipe_tx = pipe_tx.clone();
@ -237,8 +260,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
#[cfg(not(feature = "pipes"))]
let cap_full = task_accept_task_tx.capacity() == 0;
let notified = task_queue_notify.notified();
tokio::pin!(notified);
if cap_full {
tokio::time::sleep(std::time::Duration::from_millis(1)).await;
notified.await;
continue;
}
@ -277,11 +302,12 @@ impl<C: TransportConnection> GenericReceiver<C> {
let ping_sender = task_ping_sender.clone();
let connection = task_connection.clone();
let type_map = task_type_map.clone();
let decode_rejections = task_decode_rejections.clone();
tokio::spawn(async move {
let _permit = permit;
let mut stream = stream;
let mut frames = 0usize;
loop {
'stream: loop {
if policy
.max_frames_per_stream
.is_some_and(|max| frames >= max)
@ -304,7 +330,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
policy.application_close_code,
b"frame header read error",
);
break;
break 'stream;
}
Err(_) => {
break;
@ -314,6 +340,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
if len == policy.close_frame_len {
break;
}
let deadline = Instant::now() + policy.read_timeout;
let frame_limit = max_message_size.load(Ordering::Relaxed);
let body_len = len as usize;
let frame_len = match body_len.checked_add(4) {
@ -330,29 +357,25 @@ impl<C: TransportConnection> GenericReceiver<C> {
connection.close(policy.application_close_code, b"frame too large");
break;
}
let target_len = body_len;
let mut body = Vec::new();
if body.try_reserve(target_len.min(16 * 1024)).is_err() {
tracing::warn!(
target_len,
"MTP receive stream could not reserve frame body"
);
let mut frame = Vec::new();
if frame.try_reserve_exact(frame_len).is_err() {
tracing::warn!(frame_len, "MTP receive stream could not reserve frame");
let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await;
connection
.close(policy.application_close_code, b"frame allocation failed");
break;
}
while body.len() < target_len {
let chunk_len = (target_len - body.len()).min(16 * 1024);
let mut chunk = [0u8; 16 * 1024];
let body_read = tokio::time::timeout(
policy.read_timeout,
stream.read_exact(&mut chunk[..chunk_len]),
frame.extend_from_slice(&len.to_be_bytes());
frame.resize(frame_len, 0);
let mut body_offset = 4usize;
while body_offset < frame_len {
let chunk_len = (frame_len - body_offset).min(16 * 1024);
let body_read = timeout_at(
deadline,
stream.read_exact(&mut frame[body_offset..body_offset + chunk_len]),
)
.await;
if !matches!(&body_read, Ok(Ok(())))
|| body.try_reserve(chunk_len).is_err()
{
if !matches!(&body_read, Ok(Ok(()))) {
tracing::warn!(
pipe_chunk_len = chunk_len,
?body_read,
@ -361,24 +384,23 @@ impl<C: TransportConnection> GenericReceiver<C> {
let _ = tx.send(Err(CommunicationError::StreamError)).await;
connection
.close(policy.application_close_code, b"frame body read error");
break;
break 'stream;
}
body.extend_from_slice(&chunk[..chunk_len]);
}
if body.len() != target_len {
break;
body_offset += chunk_len;
}
frames += 1;
let mut frame = Vec::with_capacity(frame_len);
frame.extend_from_slice(&len.to_be_bytes());
frame.extend_from_slice(&body);
let mut message = match CommunicationValue::from_bytes_with_limits(
let mut message = match CommunicationValue::try_from_bytes_with_limits(
&frame,
DecodeLimits::for_transport_message_size(frame_limit),
) {
Ok(message) => message,
Err(_) => {
tracing::warn!("MTP receive stream contained an invalid frame");
Err(error) => {
tracing::warn!(
?error,
class = ?classify_decode_error(&error),
"MTP receive stream rejected by bounded decode"
);
decode_rejections.record(&error);
let _ = tx
.send(Err(CommunicationError::ParseCommunicationValue))
.await;
@ -463,6 +485,8 @@ impl<C: TransportConnection> GenericReceiver<C> {
ping_sender,
max_message_size,
type_map,
queue_notify,
decode_rejections,
_accept_task: Arc::new(accept_task),
}
}
@ -480,13 +504,23 @@ impl<C: TransportConnection> GenericReceiver<C> {
pub async fn set_type_map(&self, type_map: &TypeMap) {
*self.type_map.write().await = type_map.clone();
}
/// Return local counts for frames rejected by the structured decoder.
pub fn decode_rejection_counts(&self) -> crate::DecodeRejectionCounts {
self.decode_rejections.snapshot()
}
pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
self.incoming
let result = self
.incoming
.lock()
.await
.recv()
.await
.unwrap_or(Err(CommunicationError::StreamClosed))
.unwrap_or(Err(CommunicationError::StreamClosed));
if result.is_ok() {
self.queue_notify.notify_one();
}
result
}
#[cfg(feature = "pipes")]
@ -498,7 +532,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
tokio::select! {
msg = incoming.recv() => {
match msg {
Some(Ok(val)) => Ok(crate::TransportEvent::Message(val)),
Some(Ok(val)) => {
self.queue_notify.notify_one();
Ok(crate::TransportEvent::Message(val))
}
Some(Err(e)) => Err(e),
None => Err(self
.connection
@ -508,7 +545,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
}
pipe = pipes.recv() => {
match pipe {
Some(reader) => Ok(crate::TransportEvent::Pipe(reader)),
Some(reader) => {
self.queue_notify.notify_one();
Ok(crate::TransportEvent::Pipe(reader))
}
None => Err(self
.connection
.close_reason()
@ -520,12 +560,17 @@ impl<C: TransportConnection> GenericReceiver<C> {
#[cfg(feature = "pipes")]
pub async fn receive_pipe(&self) -> Result<PipeReader<C::RecvStream>, CommunicationError> {
self.pipes
let result = self
.pipes
.lock()
.await
.recv()
.await
.ok_or(CommunicationError::StreamClosed)
.ok_or(CommunicationError::StreamClosed);
if result.is_ok() {
self.queue_notify.notify_one();
}
result
}
#[cfg(feature = "pipes")]
@ -534,7 +579,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
) -> Result<Option<PipeReader<C::RecvStream>>, CommunicationError> {
match self.pipes.try_lock() {
Ok(mut rx) => match rx.try_recv() {
Ok(reader) => Ok(Some(reader)),
Ok(reader) => {
self.queue_notify.notify_one();
Ok(Some(reader))
}
Err(mpsc::error::TryRecvError::Empty) => Ok(None),
Err(mpsc::error::TryRecvError::Disconnected) => {
Err(CommunicationError::StreamClosed)

5
transport/src/lib.rs
View file

@ -11,7 +11,10 @@ pub mod encrypted_pipe;
#[cfg(feature = "pipes")]
pub mod pipe;
pub use connection::{Policy, Receiver, SendMode, Sender};
pub use connection::{
DecodeRejectionClass, DecodeRejectionCounters, DecodeRejectionCounts, Policy, Receiver,
SendMode, Sender, classify_decode_error,
};
pub use generic_connection::{GenericReceiver, GenericSender};
#[cfg(feature = "pipes")]

4
transport/tests/generic_pipe.rs
View file

@ -265,7 +265,7 @@ async fn test_regular_messages_still_work_alongside_pipes() -> Result<(), Box<dy
let sender = GenericSender::new(conn_a, policy.clone());
let receiver = GenericReceiver::new(conn_b, policy);
let msg = CommunicationValue::new(mtp_codec::CommunicationType::Pong);
let msg = CommunicationValue::new(mtp_codec::CommunicationType::BadRequest);
sender.send(&msg).await?;
let _pipe_writer = sender.open_pipe(1, "mixed-pipe").await?;
@ -273,7 +273,7 @@ async fn test_regular_messages_still_work_alongside_pipes() -> Result<(), Box<dy
let received = receiver.receive().await?;
assert_eq!(
received.get_type(),
mtp_codec::CommunicationType::Pong
mtp_codec::CommunicationType::BadRequest
.try_to_id(&mtp_codec::TypeMap::latest())
.unwrap()
);

16
transport/tests/integration.rs
View file

@ -127,12 +127,12 @@ async fn test_send_receive_roundtrip() -> Result<(), Box<dyn std::error::Error>>
assert_numbered_message(&received, CommunicationType::Ping, 42, &tm);
// Host sends a response
let resp = numbered_message(CommunicationType::Pong, 99, &tm);
let resp = numbered_message(CommunicationType::BadRequest, 99, &tm);
host_tx.send(&resp).await?;
// Client receives it
let client_received = client_rx.receive().await?;
assert_numbered_message(&client_received, CommunicationType::Pong, 99, &tm);
assert_numbered_message(&client_received, CommunicationType::BadRequest, 99, &tm);
// Close both sides
client_tx.close().await;
@ -179,13 +179,13 @@ async fn test_concurrent_messages() -> Result<(), Box<dyn std::error::Error>> {
// Send 3 responses back
for i in 0..3u128 {
let msg = numbered_message(CommunicationType::Pong, i * 10, &tm);
let msg = numbered_message(CommunicationType::BadRequest, i * 10, &tm);
client_tx.send(&msg).await?;
}
for i in 0..3u128 {
let received = host_rx.receive().await?;
assert_numbered_message(&received, CommunicationType::Pong, i * 10, &tm);
assert_numbered_message(&received, CommunicationType::BadRequest, i * 10, &tm);
}
client_tx.close().await;
@ -260,11 +260,11 @@ async fn test_drop_receiver_keeps_sender_alive() -> Result<(), Box<dyn std::erro
let tm = TypeMap::latest();
let resp = numbered_message(CommunicationType::Pong, 7, &tm);
let resp = numbered_message(CommunicationType::BadRequest, 7, &tm);
host_tx.send(&resp).await?;
let got = client_rx.receive().await?;
assert_numbered_message(&got, CommunicationType::Pong, 7, &tm);
assert_numbered_message(&got, CommunicationType::BadRequest, 7, &tm);
client_tx.close().await;
host_tx.close().await;
@ -284,10 +284,10 @@ async fn test_persistent_stream_reopens_after_local_finish()
client_tx.finish_stream().await?;
let msg2 = numbered_message(CommunicationType::Pong, 22, &tm);
let msg2 = numbered_message(CommunicationType::BadRequest, 22, &tm);
client_tx.send(&msg2).await?;
let received2 = host_rx.receive().await?;
assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm);
assert_numbered_message(&received2, CommunicationType::BadRequest, 22, &tm);
client_tx.close().await;
host_tx.close().await;