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

46
client/src/connection.rs
View file

@ -13,6 +13,10 @@ use crate::error::AuthState;
use crate::ping::{PingSession, start_ping_session};
#[cfg(feature = "pipes")]
use crate::pipe::PipeRequest;
#[cfg(feature = "pipes")]
use crate::pipe::is_expired_creation;
#[cfg(feature = "pipes")]
use crate::pipe::{PendingCreation, PendingCreationGuard};
use crate::pipe::{PendingRequest, PipeDispatcher, run_dispatcher};
pub struct MTPConnection {
@ -134,17 +138,33 @@ impl MTPConnection {
description: &str,
) -> Result<crate::pipe::PipeHandle, mtp_common::PipeError> {
let (tx, rx) = tokio::sync::oneshot::channel();
let token = Arc::new(());
let pipe_id = {
let mut pending = self.pipe_dispatcher.pending_creations.lock().await;
let mut pending = self
.pipe_dispatcher
.pending_creations
.lock()
.map_err(|_| mtp_common::PipeError::ConnectionClosed)?;
let pipe_id = loop {
let candidate = rand::random::<u32>();
if candidate != 0 && !pending.contains_key(&candidate) {
if candidate != 0
&& !pending.contains_key(&candidate)
&& !is_expired_creation(&self.pipe_dispatcher, candidate)
{
break candidate;
}
};
pending.insert(pipe_id, tx);
pending.insert(
pipe_id,
PendingCreation {
token: token.clone(),
sender: tx,
},
);
pipe_id
};
let mut creation_guard =
PendingCreationGuard::new(self.pipe_dispatcher.clone(), pipe_id, token.clone());
let request = CommunicationValue::new_with_type_map(
mtp_codec::CommunicationType::PipeRequest,
@ -154,19 +174,17 @@ impl MTPConnection {
.add_typed_default(DataType::Description, DataValue::Str(description.into()));
if let Err(error) = self.sender.send(&request).await {
self.pipe_dispatcher
.pending_creations
.lock()
.await
.remove(&pipe_id);
return Err(mtp_common::PipeError::from(error));
}
creation_guard.disarm();
Ok(crate::pipe::PipeHandle {
pipe_id,
description: description.to_string(),
sender: self.sender.clone(),
response_rx: rx,
dispatcher: self.pipe_dispatcher.clone(),
token,
})
}
@ -207,18 +225,19 @@ pub(crate) async fn connection_from_parts(
#[cfg(feature = "pipes")]
{
let receiver_queue_capacity = config.policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(
config.policy.receiver_queue_capacity,
receiver_queue_capacity,
);
let (pipe_req_tx, pipe_req_rx) =
mpsc::channel::<PipeRequest>(config.policy.receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel::<PipeRequest>(receiver_queue_capacity);
let dispatcher = Arc::new(PipeDispatcher {
pending_requests: Mutex::new(std::collections::HashMap::new()),
expired_requests: Mutex::new(std::collections::HashMap::new()),
#[cfg(feature = "pipes")]
type_map: type_map.clone(),
pending_creations: Mutex::new(std::collections::HashMap::new()),
pending_creations: std::sync::Mutex::new(std::collections::HashMap::new()),
expired_creations: std::sync::Mutex::new(std::collections::HashMap::new()),
pending_pipes: Mutex::new(std::collections::HashMap::new()),
policy: Arc::new(config.policy),
});
@ -255,8 +274,9 @@ pub(crate) async fn connection_from_parts(
#[cfg(not(feature = "pipes"))]
{
let receiver_queue_capacity = config.policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(
config.policy.receiver_queue_capacity,
receiver_queue_capacity,
);
let dispatcher = Arc::new(PipeDispatcher {
pending_requests: Mutex::new(std::collections::HashMap::new()),

21
client/src/lib.rs
View file

@ -222,6 +222,10 @@ impl MTPClient {
sender.set_type_map(&tm).await;
receiver.set_type_map(&tm).await;
let version_str = format!("{}", PROTOCOL_VERSION);
let public_key_bytes = keys
.public_key_bundle()
.try_as_bytes()
.map_err(|error| CommunicationError::ParseError(error.to_string()))?;
let mut ident =
CommunicationValue::new_with_type_map(CommunicationType::Identification, &tm)
@ -233,10 +237,7 @@ impl MTPClient {
// This capability marker lets a non-crypto host reject an
// authentication attempt instead of treating it as a plain
// unauthenticated connection.
.add_typed_default(
DataType::PublicKeys,
DataValue::Bytes(keys.public_key_bundle().as_bytes()),
);
.add_typed_default(DataType::PublicKeys, DataValue::Bytes(public_key_bytes));
if let Some(desc) = &config.description {
ident = ident.add_typed_default(DataType::Description, DataValue::Str(desc.clone()));
}
@ -415,7 +416,9 @@ impl MTPClient {
receiver.set_type_map(&tm).await;
let version_str = format!("{}", PROTOCOL_VERSION);
let pk_bundle = keys.public_key_bundle();
let pk_bytes = pk_bundle.as_bytes();
let pk_bytes = pk_bundle
.try_as_bytes()
.map_err(|error| CommunicationError::ParseError(error.to_string()))?;
let mut register = CommunicationValue::new_with_type_map(CommunicationType::Register, &tm)
.add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
@ -601,7 +604,9 @@ mod tests {
#[cfg(feature = "pipes")]
type_map: mtp_codec::TypeMap::latest(),
#[cfg(feature = "pipes")]
pending_creations: Mutex::new(HashMap::new()),
pending_creations: std::sync::Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
expired_creations: std::sync::Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
pending_pipes: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
@ -639,7 +644,9 @@ mod tests {
#[cfg(feature = "pipes")]
type_map: mtp_codec::TypeMap::latest(),
#[cfg(feature = "pipes")]
pending_creations: Mutex::new(HashMap::new()),
pending_creations: std::sync::Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
expired_creations: std::sync::Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
pending_pipes: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]

8
client/src/ping.rs
View file

@ -66,8 +66,8 @@ pub(crate) async fn start_ping_session(
return None;
}
let (pong_tx, mut pong_rx) = mpsc::unbounded_channel();
receiver.observe_pongs(pong_tx).await;
let (pong_tx, mut pong_rx) = mpsc::channel(1);
receiver.observe_pongs_bounded(pong_tx).await;
let last_ping = Arc::new(Mutex::new(None));
let ping_state = last_ping.clone();
let interval = config.ping_interval;
@ -75,6 +75,7 @@ pub(crate) async fn start_ping_session(
let max_missed_pings = config.max_missed_pings;
let ping_timestamp = config.ping_timestamp;
let type_map = type_map.clone();
let ping_receiver = receiver.clone();
let mut close_rx = receiver.handle().subscribe_close();
let task = tokio::spawn(async move {
@ -91,6 +92,7 @@ pub(crate) async fn start_ping_session(
}
_ = ticker.tick() => {
let missed_pings = tracker.begin_round();
ping_receiver.set_expected_pong_id(None).await;
if max_missed_pings > 0 && missed_pings >= max_missed_pings {
sender.close().await;
break;
@ -121,7 +123,9 @@ pub(crate) async fn start_ping_session(
sender.close().await;
break;
};
ping_receiver.set_expected_pong_id(Some(id)).await;
if sender.send(&ping).await.is_err() {
ping_receiver.set_expected_pong_id(None).await;
sender.close().await;
break;
}

183
client/src/pipe.rs
View file

@ -5,6 +5,8 @@ use mtp_common::CommunicationError;
use mtp_transport::Receiver;
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(feature = "pipes")]
use std::sync::Mutex as StdMutex;
use tokio::sync::{Mutex, mpsc};
use tokio::time::{Duration, Instant};
@ -21,6 +23,8 @@ pub struct PipeHandle {
pub(crate) description: String,
pub(crate) sender: Sender,
pub(crate) response_rx: tokio::sync::oneshot::Receiver<Result<bool, PipeError>>,
pub(crate) dispatcher: Arc<PipeDispatcher>,
pub(crate) token: Arc<()>,
}
#[cfg(feature = "pipes")]
@ -33,9 +37,11 @@ impl PipeHandle {
&self.description
}
pub async fn wait(self) -> Result<Option<mtp_transport::PipeWriter>, PipeError> {
match self.response_rx.await {
Ok(Ok(true)) => {
pub async fn wait(mut self) -> Result<Option<mtp_transport::PipeWriter>, PipeError> {
let response =
tokio::time::timeout(self.dispatcher.policy.read_timeout, &mut self.response_rx).await;
match response {
Ok(Ok(Ok(true))) => {
let writer = self
.sender
.open_pipe(self.pipe_id, &self.description)
@ -43,13 +49,30 @@ impl PipeHandle {
.map_err(PipeError::from)?;
Ok(Some(writer))
}
Ok(Ok(false)) => Ok(None),
Ok(Err(e)) => Err(e),
Err(_) => Err(PipeError::StreamClosed),
Ok(Ok(Ok(false))) => Ok(None),
Ok(Ok(Err(error))) => {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
Err(error)
}
Ok(Err(_)) => {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
Err(PipeError::StreamClosed)
}
Err(_) => {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
Err(PipeError::HandshakeTimeout)
}
}
}
}
#[cfg(feature = "pipes")]
impl Drop for PipeHandle {
fn drop(&mut self) {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
}
}
#[cfg(feature = "pipes")]
pub struct PipeRequest {
pub(crate) pipe_id: u32,
@ -129,14 +152,54 @@ pub(crate) struct PendingRequest {
pub(crate) sender: tokio::sync::oneshot::Sender<Result<CommunicationValue, CommunicationError>>,
}
#[cfg(feature = "pipes")]
pub(crate) struct PendingCreation {
pub(crate) token: Arc<()>,
pub(crate) sender: tokio::sync::oneshot::Sender<Result<bool, PipeError>>,
}
#[cfg(feature = "pipes")]
pub(crate) struct PendingCreationGuard {
dispatcher: Arc<PipeDispatcher>,
pipe_id: u32,
token: Arc<()>,
armed: bool,
}
#[cfg(feature = "pipes")]
impl PendingCreationGuard {
pub(crate) fn new(dispatcher: Arc<PipeDispatcher>, pipe_id: u32, token: Arc<()>) -> Self {
Self {
dispatcher,
pipe_id,
token,
armed: true,
}
}
pub(crate) fn disarm(&mut self) {
self.armed = false;
}
}
#[cfg(feature = "pipes")]
impl Drop for PendingCreationGuard {
fn drop(&mut self) {
if self.armed {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
}
}
}
pub(crate) struct PipeDispatcher {
pub(crate) pending_requests: Mutex<HashMap<u32, PendingRequest>>,
pub(crate) expired_requests: Mutex<HashMap<u32, Instant>>,
#[cfg(feature = "pipes")]
pub(crate) type_map: TypeMap,
#[cfg(feature = "pipes")]
pub(crate) pending_creations:
Mutex<HashMap<u32, tokio::sync::oneshot::Sender<Result<bool, PipeError>>>>,
pub(crate) pending_creations: StdMutex<HashMap<u32, PendingCreation>>,
#[cfg(feature = "pipes")]
pub(crate) expired_creations: StdMutex<HashMap<u32, Instant>>,
#[cfg(feature = "pipes")]
pub(crate) pending_pipes:
Mutex<HashMap<u32, tokio::sync::oneshot::Sender<mtp_transport::PipeReader>>>,
@ -144,6 +207,90 @@ pub(crate) struct PipeDispatcher {
pub(crate) policy: Arc<Policy>,
}
#[cfg(feature = "pipes")]
const EXPIRED_CREATION_TOMBSTONE_TTL: Duration = Duration::from_secs(60);
#[cfg(feature = "pipes")]
const MAX_EXPIRED_CREATION_TOMBSTONES: usize = 1024;
#[cfg(feature = "pipes")]
pub(crate) fn expire_pending_creation(dispatcher: &PipeDispatcher, pipe_id: u32, token: &Arc<()>) {
let removed = dispatcher
.pending_creations
.lock()
.ok()
.and_then(|mut pending| {
if pending
.get(&pipe_id)
.is_some_and(|entry| Arc::ptr_eq(&entry.token, token))
{
pending.remove(&pipe_id);
Some(())
} else {
None
}
});
if removed.is_none() {
return;
}
let Ok(mut expired) = dispatcher.expired_creations.lock() else {
return;
};
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
if expired.len() >= MAX_EXPIRED_CREATION_TOMBSTONES
&& let Some(oldest) = expired
.iter()
.min_by_key(|(_, expires_at)| **expires_at)
.map(|(id, _)| *id)
{
expired.remove(&oldest);
}
expired.insert(pipe_id, now + EXPIRED_CREATION_TOMBSTONE_TTL);
}
#[cfg(feature = "pipes")]
fn consume_expired_creation(dispatcher: &PipeDispatcher, pipe_id: u32) -> bool {
let Ok(mut expired) = dispatcher.expired_creations.lock() else {
return false;
};
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.remove(&pipe_id).is_some()
}
#[cfg(feature = "pipes")]
pub(crate) fn is_expired_creation(dispatcher: &PipeDispatcher, pipe_id: u32) -> bool {
let Ok(mut expired) = dispatcher.expired_creations.lock() else {
return true;
};
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.contains_key(&pipe_id)
}
#[cfg(feature = "pipes")]
pub(crate) fn fail_pending_creations(dispatcher: &PipeDispatcher, error: &CommunicationError) {
let pending = dispatcher
.pending_creations
.lock()
.ok()
.map(|mut pending| std::mem::take(&mut *pending));
if let Some(pending) = pending {
let error = PipeError::from(error.clone());
for (_, pending) in pending {
let _ = pending.sender.send(Err(error.clone()));
}
}
if let Ok(mut expired) = dispatcher.expired_creations.lock() {
expired.clear();
}
}
#[cfg(feature = "pipes")]
pub(crate) async fn fail_pending_pipes(dispatcher: &PipeDispatcher) {
dispatcher.pending_pipes.lock().await.clear();
}
pub(crate) async fn route_message(
msg: CommunicationValue,
app_tx: &mpsc::Sender<Result<CommunicationValue, CommunicationError>>,
@ -283,9 +430,15 @@ pub(crate) async fn run_dispatcher(
continue;
};
let accepted = msg.get_bool(DataType::Accepted).unwrap_or(false);
let mut pending = dispatcher.pending_creations.lock().await;
if let Some(tx) = pending.remove(&pipe_id) {
let _ = tx.send(Ok(accepted));
let pending = dispatcher
.pending_creations
.lock()
.ok()
.and_then(|mut pending| pending.remove(&pipe_id));
if let Some(entry) = pending {
let _ = entry.sender.send(Ok(accepted));
} else {
let _ = consume_expired_creation(&dispatcher, pipe_id);
}
continue;
}
@ -303,6 +456,10 @@ pub(crate) async fn run_dispatcher(
}
Err(e) => {
fail_pending_requests(&dispatcher, e.clone()).await;
#[cfg(feature = "pipes")]
fail_pending_creations(&dispatcher, &e);
#[cfg(feature = "pipes")]
fail_pending_pipes(&dispatcher).await;
let _ = app_tx.send(Err(e)).await;
break;
}
@ -325,6 +482,10 @@ pub(crate) async fn run_dispatcher(
}
Err(e) => {
fail_pending_requests(&dispatcher, e.clone()).await;
#[cfg(feature = "pipes")]
fail_pending_creations(&dispatcher, &e);
#[cfg(feature = "pipes")]
fail_pending_pipes(&dispatcher).await;
let _ = app_tx.send(Err(e)).await;
break;
}