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

[WIP] Security work While on holiday

Browse source
This commit is contained in:
Alex 2026-08-12 22:45:28 +02:00
commit 7f0231e3f1
Signed by: alex
SSH key fingerprint: SHA256:D1+Ub8o0v4K5y1JNivW8IxEOelqLSvPmUzBbDIoZkRQ
109 changed files with 19694 additions and 5210 deletions
Download patch file Download diff file Expand all files Collapse all files

77
client/src/connection.rs
View file

@ -1,4 +1,4 @@
use mtp_codec::{CommunicationValue, Version};
use mtp_codec::{CommunicationValue, Version, registry::VersionedCodec};
#[cfg(feature = "pipes")]
use mtp_codec::{DataType, DataValue};
use mtp_common::CommunicationError;
@ -17,6 +17,7 @@ use crate::pipe::{PendingRequest, PipeDispatcher, run_dispatcher};
pub struct MTPConnection {
pub version: Version,
pub codec: VersionedCodec,
pub sender: mtp_transport::Sender,
pub receiver: mtp_transport::Receiver,
pub description: Option<String>,
@ -45,12 +46,19 @@ impl MTPConnection {
request: &CommunicationValue,
expected_response: Option<mtp_codec::CommunicationType>,
) -> Result<CommunicationValue, CommunicationError> {
let request_id = request.get_id();
let request_id = request
.id()
.ok_or_else(|| CommunicationError::Other("request frame must contain an id".into()))?;
if request_id == 0 {
return Err(CommunicationError::Other(
"request frame must have a non-zero id".into(),
));
}
if crate::pipe::is_expired_request(&self.pipe_dispatcher, request_id).await {
return Err(CommunicationError::Other(format!(
"request id {request_id} recently timed out; use a new request id"
)));
}
let (sender, receiver) = tokio::sync::oneshot::channel();
let token = Arc::new(());
@ -86,7 +94,7 @@ impl MTPConnection {
result?
}
Err(_) => {
crate::pipe::remove_pending_request(&self.pipe_dispatcher, request_id, &token)
crate::pipe::expire_pending_request(&self.pipe_dispatcher, request_id, &token)
.await;
return Err(CommunicationError::Other(format!(
"request {request_id} timed out after {:?}",
@ -96,7 +104,7 @@ impl MTPConnection {
};
if let Some(expected) = expected_response {
let expected_type = expected.try_to_id(&mtp_codec::TypeMap::latest());
let expected_type = expected.try_to_id(self.codec.type_map());
if Some(response.get_type()) != expected_type {
return Err(CommunicationError::Other(format!(
"unexpected response type: expected {:?}, got {:?}; parsed {}",
@ -125,23 +133,35 @@ impl MTPConnection {
&self,
description: &str,
) -> Result<crate::pipe::PipeHandle, mtp_common::PipeError> {
let pipe_id = rand::random::<u32>();
let (tx, rx) = tokio::sync::oneshot::channel();
{
let pipe_id = {
let mut pending = self.pipe_dispatcher.pending_creations.lock().await;
let pipe_id = loop {
let candidate = rand::random::<u32>();
if candidate != 0 && !pending.contains_key(&candidate) {
break candidate;
}
};
pending.insert(pipe_id, tx);
pipe_id
};
let request = CommunicationValue::new_with_type_map(
mtp_codec::CommunicationType::PipeRequest,
self.codec.type_map(),
)
.with_id(pipe_id)
.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));
}
let request = CommunicationValue::new(mtp_codec::CommunicationType::PipeRequest)
.with_id(pipe_id)
.add_typed_default(DataType::Description, DataValue::Str(description.into()));
self.sender
.send(&request)
.await
.map_err(mtp_common::PipeError::from)?;
Ok(crate::pipe::PipeHandle {
pipe_id,
description: description.to_string(),
@ -164,11 +184,26 @@ pub(crate) async fn connection_from_parts(
sender: mtp_transport::Sender,
receiver: mtp_transport::Receiver,
version: Version,
codec: VersionedCodec,
#[cfg(feature = "crypto")] auth_state: AuthState,
#[cfg(feature = "crypto")] client_id: u64,
) -> MTPConnection {
#[cfg(feature = "pipes")]
let type_map = codec.type_map().clone();
receiver.set_type_map(codec.type_map()).await;
let remote_addr = sender.handle().remote_addr();
let ping = start_ping_session(&config, sender.clone(), &receiver).await;
#[cfg(feature = "crypto")]
let ping_client_id = client_id;
#[cfg(not(feature = "crypto"))]
let ping_client_id = config.client_id;
let ping = start_ping_session(
&config,
sender.clone(),
&receiver,
codec.type_map(),
ping_client_id,
)
.await;
#[cfg(feature = "pipes")]
{
@ -180,6 +215,9 @@ pub(crate) async fn connection_from_parts(
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_pipes: Mutex::new(std::collections::HashMap::new()),
policy: Arc::new(config.policy),
@ -197,6 +235,7 @@ pub(crate) async fn connection_from_parts(
MTPConnection {
version,
codec,
sender,
receiver,
app_rx: Mutex::new(app_rx),
@ -221,11 +260,15 @@ pub(crate) async fn connection_from_parts(
);
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,
});
let task = tokio::spawn(run_dispatcher(receiver.clone(), app_tx, dispatcher.clone()));
MTPConnection {
version,
codec,
sender,
receiver,
app_rx: Mutex::new(app_rx),

36
client/src/crypto.rs
View file

@ -1,4 +1,4 @@
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, Version};
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap, Version};
use mtp_common::CommunicationError;
pub(crate) fn unexpected_response_type_error(
@ -24,7 +24,7 @@ pub(crate) async fn verify_host_challenge(
use mtp_crypto::{auth, verify_ed25519};
let sig = match challenge.get_data(DataType::Signature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => {
return Err(CommunicationError::AuthenticationFailed(
"Missing host challenge signature".into(),
@ -32,11 +32,11 @@ pub(crate) async fn verify_host_challenge(
}
};
let pq_sig = match challenge.get_data(DataType::PqSignature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => vec![],
};
let host_requires_pq = challenge.get_data(DataType::RequirePq) == &DataValue::BoolTrue;
let host_requires_pq = challenge.get_data(DataType::RequirePq) == Some(&DataValue::BoolTrue);
if host_requires_pq && host_pk.sig_pq_public_key.as_bytes().is_empty() {
return Err(CommunicationError::AuthenticationFailed(
"Host requires post-quantum authentication but its PQ public key is absent".into(),
@ -80,7 +80,7 @@ pub(crate) async fn verify_host_final(
use mtp_crypto::{auth, verify_ed25519};
match response.get_data(DataType::ClientNonce) {
DataValue::UnsignedNumber(n) if *n == client_nonce => {}
Some(DataValue::UnsignedNumber(n)) if *n == client_nonce => {}
_ => {
return Err(CommunicationError::AuthenticationFailed(
"Nonce mismatch".into(),
@ -89,7 +89,7 @@ pub(crate) async fn verify_host_final(
}
let sig = match response.get_data(DataType::Signature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => {
return Err(CommunicationError::AuthenticationFailed(
"Missing signature".into(),
@ -97,7 +97,7 @@ pub(crate) async fn verify_host_final(
}
};
let pq_sig = match response.get_data(DataType::PqSignature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => vec![],
};
if require_pq && pq_sig.is_empty() {
@ -130,8 +130,8 @@ pub(crate) fn check_connected(
reject_msg: &str,
) -> Result<(), CommunicationError> {
match response.get_data(DataType::Connected) {
DataValue::BoolTrue => Ok(()),
DataValue::BoolFalse => Err(CommunicationError::AuthenticationFailed(
Some(DataValue::BoolTrue) => Ok(()),
Some(DataValue::BoolFalse) => Err(CommunicationError::AuthenticationFailed(
response
.get_str(DataType::ErrorMessage)
.unwrap_or(reject_msg)
@ -147,7 +147,7 @@ pub(crate) fn negotiated_version(
response: &CommunicationValue,
) -> Result<Version, CommunicationError> {
match response.get_data(DataType::Version) {
DataValue::Str(version) => Version::parse(version).ok_or_else(|| {
Some(DataValue::Str(version)) => Version::parse(version).ok_or_else(|| {
CommunicationError::AuthenticationFailed(
"Host returned an invalid negotiated protocol version".into(),
)
@ -162,16 +162,18 @@ pub(crate) async fn signed_challenge_response(
keys: &mtp_crypto::Keyring,
proof_payload: Vec<u8>,
client_nonce: u128,
type_map: &TypeMap,
) -> Result<CommunicationValue, CommunicationError> {
use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme};
let signer = Ed25519Signer::new(&keys.sig_cl_secret_key)
.map_err(|e| CommunicationError::Other(e.to_string()))?;
let mut proof = CommunicationValue::new(CommunicationType::ChallengeResponse)
.add_typed_default(
DataType::ClientNonce,
DataValue::UnsignedNumber(client_nonce),
);
let mut proof =
CommunicationValue::new_with_type_map(CommunicationType::ChallengeResponse, type_map)
.add_typed_default(
DataType::ClientNonce,
DataValue::UnsignedNumber(client_nonce),
);
if keys.sig_pq_secret_key.as_bytes().is_empty() {
let signature = signer
@ -213,14 +215,14 @@ pub(crate) async fn receive_verified_challenge(
}
let server_challenge = match challenge.get_data(DataType::ServerNonce) {
DataValue::UnsignedNumber(n) => *n,
Some(DataValue::UnsignedNumber(n)) => *n,
_ => {
return Err(CommunicationError::AuthenticationFailed(
"Missing server challenge".into(),
));
}
};
if challenge.get_data(DataType::RequirePq) == &DataValue::BoolTrue && !client_has_pq_key {
if challenge.get_data(DataType::RequirePq) == Some(&DataValue::BoolTrue) && !client_has_pq_key {
return Err(CommunicationError::AuthenticationFailed(
"Host requires post-quantum authentication but the client PQ key is absent".into(),
));

238
client/src/lib.rs
View file

@ -31,20 +31,22 @@ mod error {
}
}
use mtp_codec::{CommunicationValue, DataType, DataValue, PROTOCOL_VERSION, Version};
use mtp_codec::{
CommunicationValue, DataType, DataValue, PROTOCOL_VERSION, Version,
registry::{Registry, VersionedCodec},
};
use mtp_common::{CommunicationError, HandshakeOutcome, RejectionReason};
use connection::connection_from_parts;
fn parse_handshake_response(
response: &CommunicationValue,
type_map: &mtp_codec::TypeMap,
) -> Result<HandshakeOutcome, CommunicationError> {
let tm = mtp_codec::TypeMap::latest();
let bad_version = mtp_codec::CommunicationType::ErrorBadVersion.try_to_id(&tm);
let bad_version = mtp_codec::CommunicationType::ErrorBadVersion.try_to_id(type_map);
if Some(response.get_type()) == bad_version {
let supported_versions = match response.get_data(DataType::Version) {
DataValue::Str(v) if !v.is_empty() => v.split(',').map(String::from).collect(),
Some(DataValue::Str(v)) if !v.is_empty() => v.split(',').map(String::from).collect(),
_ => vec![],
};
return Ok(HandshakeOutcome::Rejected {
@ -53,7 +55,7 @@ fn parse_handshake_response(
}
let expected = mtp_codec::CommunicationType::IdentificationResponse
.try_to_id(&tm)
.try_to_id(type_map)
.ok_or_else(|| {
CommunicationError::Other("IdentificationResponse is absent from the type map".into())
})?;
@ -68,9 +70,9 @@ fn parse_handshake_response(
}
match response.get_data(DataType::Connected) {
DataValue::BoolTrue => {
Some(DataValue::BoolTrue) => {
let version = match response.get_data(DataType::Version) {
DataValue::Str(v) => v.clone(),
Some(DataValue::Str(v)) => v.clone(),
_ => {
return Err(CommunicationError::Other(
"host omitted the negotiated version".into(),
@ -78,15 +80,21 @@ fn parse_handshake_response(
}
};
let assigned_id = match response.get_data(DataType::Id) {
DataValue::UnsignedNumber(n) => *n as u64,
_ => 0,
Some(DataValue::UnsignedNumber(n)) => u64::try_from(*n).map_err(|_| {
CommunicationError::Other("host returned an out-of-range client id".into())
})?,
_ => {
return Err(CommunicationError::Other(
"host omitted the assigned client id".into(),
));
}
};
Ok(HandshakeOutcome::Accepted {
version,
assigned_id,
})
}
DataValue::BoolFalse => {
Some(DataValue::BoolFalse) => {
let detail = response
.get_str(DataType::ErrorMessage)
.unwrap_or("host rejected the connection")
@ -99,20 +107,34 @@ fn parse_handshake_response(
}
}
fn codec_for_version(version: &Version) -> Result<VersionedCodec, CommunicationError> {
VersionedCodec::for_version(Registry::builtin(), version.clone()).ok_or_else(|| {
CommunicationError::Other(format!(
"host returned unsupported protocol version {version}"
))
})
}
pub struct MTPClient;
impl MTPClient {
pub async fn connect(config: ClientConfig) -> Result<MTPConnection, CommunicationError> {
let (sender, receiver) =
mtp_transport::connect(&config.url, config.server_cert(), config.policy).await?;
let opening_codec = codec_for_version(&PROTOCOL_VERSION)?;
sender.set_type_map(opening_codec.type_map()).await;
receiver.set_type_map(opening_codec.type_map()).await;
let version_str = format!("{}", PROTOCOL_VERSION);
let mut ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification)
.add_typed_default(DataType::Version, DataValue::Str(version_str))
.add_typed_default(
DataType::Id,
DataValue::UnsignedNumber(config.client_id.into()),
);
let mut ident = CommunicationValue::new_with_type_map(
mtp_codec::CommunicationType::Identification,
opening_codec.type_map(),
)
.add_typed_default(DataType::Version, DataValue::Str(version_str))
.add_typed_default(
DataType::Id,
DataValue::UnsignedNumber(config.client_id.into()),
);
if let Some(desc) = &config.description {
ident = ident.add_typed_default(DataType::Description, DataValue::Str(desc.clone()));
}
@ -120,32 +142,47 @@ impl MTPClient {
sender.send(&ident).await?;
let response = receiver.receive().await?;
let outcome = parse_handshake_response(&response)?;
let negotiated = match outcome {
mtp_common::HandshakeOutcome::Accepted { version, .. } => Version::parse(&version)
.ok_or_else(|| {
let outcome = parse_handshake_response(&response, opening_codec.type_map())?;
let (negotiated, assigned_id) = match outcome {
mtp_common::HandshakeOutcome::Accepted {
version,
assigned_id,
} => (
Version::parse(&version).ok_or_else(|| {
CommunicationError::Other("host returned an invalid negotiated version".into())
})?,
assigned_id,
),
mtp_common::HandshakeOutcome::Rejected { reason } => {
sender.close().await;
return Err(CommunicationError::Other(reason.to_string()));
}
};
#[cfg(not(feature = "crypto"))]
let _ = assigned_id;
if negotiated != PROTOCOL_VERSION {
sender.close().await;
return Err(CommunicationError::Other(
"host selected a protocol version the client did not offer".into(),
));
}
let codec = codec_for_version(&negotiated)?;
#[cfg(feature = "crypto")]
let client_id = config.client_id;
let client_id = assigned_id;
#[cfg(feature = "crypto")]
return Ok(connection_from_parts(
config,
sender,
receiver,
negotiated,
codec,
error::AuthState::Unauthenticated,
client_id,
)
.await);
#[cfg(not(feature = "crypto"))]
Ok(connection_from_parts(config, sender, receiver, negotiated).await)
Ok(connection_from_parts(config, sender, receiver, negotiated, codec).await)
}
}
@ -180,15 +217,26 @@ impl MTPClient {
let (sender, receiver) =
mtp_transport::connect(&config.url, config.server_cert(), config.policy).await?;
let tm = mtp_codec::TypeMap::latest();
let handshake_codec = codec_for_version(&PROTOCOL_VERSION)?;
let tm = handshake_codec.type_map().clone();
sender.set_type_map(&tm).await;
receiver.set_type_map(&tm).await;
let version_str = format!("{}", PROTOCOL_VERSION);
let mut ident = CommunicationValue::new(CommunicationType::Identification)
.add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
.add_typed_default(
DataType::Id,
DataValue::UnsignedNumber(config.client_id as u128),
);
let mut ident =
CommunicationValue::new_with_type_map(CommunicationType::Identification, &tm)
.add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
.add_typed_default(
DataType::Id,
DataValue::UnsignedNumber(config.client_id as u128),
)
// 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()),
);
if let Some(desc) = &config.description {
ident = ident.add_typed_default(DataType::Description, DataValue::Str(desc.clone()));
}
@ -223,14 +271,14 @@ impl MTPClient {
client_nonce,
);
let proof = match crypto::signed_challenge_response(keys, proof_payload, client_nonce).await
{
Ok(p) => p,
Err(e) => {
sender.close().await;
return Err(e);
}
};
let proof =
match crypto::signed_challenge_response(keys, proof_payload, client_nonce, &tm).await {
Ok(p) => p,
Err(e) => {
sender.close().await;
return Err(e);
}
};
if let Err(e) = sender.send(&proof).await {
sender.close().await;
return Err(e);
@ -276,12 +324,41 @@ impl MTPClient {
return Err(e);
}
let assigned_id = match response.get_data(DataType::Id) {
Some(DataValue::UnsignedNumber(id)) => u64::try_from(*id).map_err(|_| {
CommunicationError::AuthenticationFailed(
"host returned an out-of-range client id".into(),
)
})?,
_ => {
sender.close().await;
return Err(CommunicationError::AuthenticationFailed(
"host omitted the authenticated client id".into(),
));
}
};
if assigned_id != config.client_id {
sender.close().await;
return Err(CommunicationError::AuthenticationFailed(
"host returned a different authenticated client id".into(),
));
}
let negotiated = crypto::negotiated_version(&response)?;
if negotiated != PROTOCOL_VERSION {
sender.close().await;
return Err(CommunicationError::AuthenticationFailed(
"host selected a protocol version the client did not offer".into(),
));
}
let codec = codec_for_version(&negotiated)?;
let client_id = config.client_id;
Ok(connection_from_parts(
config,
sender,
receiver,
crypto::negotiated_version(&response)?,
negotiated,
codec,
error::AuthState::Authenticated,
client_id,
)
@ -332,12 +409,15 @@ impl MTPClient {
let (sender, receiver) =
mtp_transport::connect(&config.url, config.server_cert(), config.policy).await?;
let tm = mtp_codec::TypeMap::latest();
let handshake_codec = codec_for_version(&PROTOCOL_VERSION)?;
let tm = handshake_codec.type_map().clone();
sender.set_type_map(&tm).await;
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 mut register = CommunicationValue::new(CommunicationType::Register)
let mut register = CommunicationValue::new_with_type_map(CommunicationType::Register, &tm)
.add_typed_default(DataType::Version, DataValue::Str(version_str.clone()))
.add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes.clone()));
if let Some(desc) = &config.description {
@ -375,14 +455,14 @@ impl MTPClient {
client_nonce,
);
let proof = match crypto::signed_challenge_response(keys, proof_payload, client_nonce).await
{
Ok(p) => p,
Err(e) => {
sender.close().await;
return Err(e);
}
};
let proof =
match crypto::signed_challenge_response(keys, proof_payload, client_nonce, &tm).await {
Ok(p) => p,
Err(e) => {
sender.close().await;
return Err(e);
}
};
if let Err(e) = sender.send(&proof).await {
sender.close().await;
return Err(e);
@ -413,7 +493,11 @@ impl MTPClient {
return Err(e);
}
let assigned_id = match response.get_data(DataType::Id) {
DataValue::UnsignedNumber(n) => *n as u64,
Some(DataValue::UnsignedNumber(n)) => u64::try_from(*n).map_err(|_| {
CommunicationError::AuthenticationFailed(
"host returned an out-of-range client id".into(),
)
})?,
_ => {
sender.close().await;
return Err(CommunicationError::AuthenticationFailed(
@ -435,11 +519,20 @@ impl MTPClient {
return Err(e);
}
let negotiated = crypto::negotiated_version(&response)?;
if negotiated != PROTOCOL_VERSION {
sender.close().await;
return Err(CommunicationError::AuthenticationFailed(
"host selected a protocol version the client did not offer".into(),
));
}
let codec = codec_for_version(&negotiated)?;
Ok(connection_from_parts(
config,
sender,
receiver,
crypto::negotiated_version(&response)?,
negotiated,
codec,
error::AuthState::Authenticated,
assigned_id,
)
@ -504,6 +597,9 @@ mod tests {
let dispatcher = pipe::PipeDispatcher {
pending_requests: Mutex::new(HashMap::new()),
expired_requests: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
type_map: mtp_codec::TypeMap::latest(),
#[cfg(feature = "pipes")]
pending_creations: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
@ -523,11 +619,47 @@ mod tests {
let unrelated = CommunicationValue::new(mtp_codec::CommunicationType::Ping).with_id(8);
assert!(pipe::route_message(unrelated, &app_tx, &dispatcher).await);
assert_eq!(app_rx.recv().await.unwrap().unwrap().get_id(), 8);
assert_eq!(app_rx.recv().await.unwrap().unwrap().id(), Some(8));
let response = CommunicationValue::new(mtp_codec::CommunicationType::Pong).with_id(7);
assert!(pipe::route_message(response, &app_tx, &dispatcher).await);
assert_eq!(response_rx.await.unwrap().unwrap().get_id(), 7);
assert_eq!(response_rx.await.unwrap().unwrap().id(), Some(7));
assert!(app_rx.try_recv().is_err());
}
#[tokio::test]
async fn test_expired_request_response_is_consumed() {
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
let dispatcher = pipe::PipeDispatcher {
pending_requests: Mutex::new(HashMap::new()),
expired_requests: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
type_map: mtp_codec::TypeMap::latest(),
#[cfg(feature = "pipes")]
pending_creations: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
pending_pipes: Mutex::new(HashMap::new()),
#[cfg(feature = "pipes")]
policy: Arc::new(Policy::default()),
};
let token = Arc::new(());
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
dispatcher.pending_requests.lock().await.insert(
9,
pipe::PendingRequest {
token: token.clone(),
sender: response_tx,
},
);
pipe::expire_pending_request(&dispatcher, 9, &token).await;
let (app_tx, mut app_rx) = tokio::sync::mpsc::channel(1);
let response = CommunicationValue::new(mtp_codec::CommunicationType::Pong).with_id(9);
assert!(pipe::route_message(response, &app_tx, &dispatcher).await);
assert!(response_rx.await.is_err());
assert!(app_rx.try_recv().is_err());
}

20
client/src/ping.rs
View file

@ -3,7 +3,7 @@ use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use tokio::time::{Duration, Instant};
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
use mtp_transport::{Receiver, Sender};
pub(crate) struct PingSession {
@ -59,6 +59,8 @@ pub(crate) async fn start_ping_session(
config: &crate::config::ClientConfig,
sender: Sender,
receiver: &Receiver,
type_map: &TypeMap,
client_id: u64,
) -> Option<PingSession> {
if config.ping_interval.is_zero() {
return None;
@ -72,6 +74,7 @@ pub(crate) async fn start_ping_session(
let ping_jitter = config.ping_jitter;
let max_missed_pings = config.max_missed_pings;
let ping_timestamp = config.ping_timestamp;
let type_map = type_map.clone();
let mut close_rx = receiver.handle().subscribe_close();
let task = tokio::spawn(async move {
@ -99,7 +102,11 @@ pub(crate) async fn start_ping_session(
tokio::time::sleep(Duration::from_millis(extra)).await;
}
let mut ping = CommunicationValue::new(CommunicationType::Ping);
let mut ping = CommunicationValue::new_with_type_map(
CommunicationType::Ping,
&type_map,
)
.with_sender(client_id);
if ping_timestamp {
let sent_at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
@ -110,7 +117,10 @@ pub(crate) async fn start_ping_session(
DataValue::UnsignedNumber(sent_at),
);
}
let id = ping.get_id();
let Some(id) = ping.id() else {
sender.close().await;
break;
};
if sender.send(&ping).await.is_err() {
sender.close().await;
break;
@ -119,7 +129,9 @@ pub(crate) async fn start_ping_session(
}
pong = pong_rx.recv() => match pong {
Some(pong) => {
if let Some(ping) = tracker.received(pong.get_id()) {
if let Some(id) = pong.id()
&& let Some(ping) = tracker.received(id)
{
let mut last_ping = ping_state.lock().await;
*last_ping = Some(ping);
}

157
client/src/pipe.rs
View file

@ -1,9 +1,12 @@
use mtp_codec::CommunicationValue;
#[cfg(feature = "pipes")]
use mtp_codec::TypeMap;
use mtp_common::CommunicationError;
use mtp_transport::Receiver;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use tokio::time::{Duration, Instant};
#[cfg(feature = "pipes")]
use mtp_codec::{CommunicationType, DataType, DataValue};
@ -72,22 +75,50 @@ impl PipeRequest {
pending.insert(self.pipe_id, pipe_tx);
}
let resp = CommunicationValue::new(CommunicationType::PipeResponse)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolTrue);
self.sender.send(&resp).await.map_err(PipeError::from)?;
let resp = CommunicationValue::new_with_type_map(
CommunicationType::PipeResponse,
&self.dispatcher.type_map,
)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolTrue);
if let Err(error) = self.sender.send(&resp).await {
self.dispatcher
.pending_pipes
.lock()
.await
.remove(&self.pipe_id);
return Err(PipeError::from(error));
}
let timeout = self.dispatcher.policy.read_timeout;
tokio::time::timeout(timeout, pipe_rx)
.await
.map_err(|_| PipeError::HandshakeTimeout)?
.map_err(|_| PipeError::StreamClosed)
match tokio::time::timeout(timeout, pipe_rx).await {
Ok(Ok(reader)) => Ok(reader),
Ok(Err(_)) => {
self.dispatcher
.pending_pipes
.lock()
.await
.remove(&self.pipe_id);
Err(PipeError::StreamClosed)
}
Err(_) => {
self.dispatcher
.pending_pipes
.lock()
.await
.remove(&self.pipe_id);
Err(PipeError::HandshakeTimeout)
}
}
}
pub async fn deny(self) -> Result<(), PipeError> {
let resp = CommunicationValue::new(CommunicationType::PipeResponse)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolFalse);
let resp = CommunicationValue::new_with_type_map(
CommunicationType::PipeResponse,
&self.dispatcher.type_map,
)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolFalse);
self.sender.send(&resp).await.map_err(PipeError::from)?;
Ok(())
}
@ -100,6 +131,9 @@ pub(crate) struct PendingRequest {
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>>>>,
@ -115,14 +149,27 @@ pub(crate) async fn route_message(
app_tx: &mpsc::Sender<Result<CommunicationValue, CommunicationError>>,
dispatcher: &PipeDispatcher,
) -> bool {
let pending = dispatcher
.pending_requests
.lock()
.await
.remove(&msg.get_id());
if let Some(tx) = pending {
let _ = tx.sender.send(Ok(msg));
return true;
if !matches!(msg.id(), Some(id) if id != 0)
&& msg
.get_type_name()
.is_some_and(|name| name.ends_with("Response"))
{
return app_tx
.send(Err(CommunicationError::Other(
"response frame must contain a non-zero id".into(),
)))
.await
.is_ok();
}
if let Some(id) = msg.id() {
let pending = dispatcher.pending_requests.lock().await.remove(&id);
if let Some(tx) = pending {
let _ = tx.sender.send(Ok(msg));
return true;
}
if consume_expired_request(dispatcher, id).await {
return true;
}
}
app_tx.send(Ok(msg)).await.is_ok()
@ -135,6 +182,44 @@ pub(crate) async fn fail_pending_requests(dispatcher: &PipeDispatcher, error: Co
}
}
const EXPIRED_REQUEST_TOMBSTONE_TTL: Duration = Duration::from_secs(60);
const MAX_EXPIRED_REQUEST_TOMBSTONES: usize = 1024;
pub(crate) async fn expire_pending_request(
dispatcher: &PipeDispatcher,
request_id: u32,
token: &Arc<()>,
) {
let mut pending = dispatcher.pending_requests.lock().await;
if pending
.get(&request_id)
.is_some_and(|entry| Arc::ptr_eq(&entry.token, token))
{
pending.remove(&request_id);
drop(pending);
let mut expired = dispatcher.expired_requests.lock().await;
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
if expired.len() >= MAX_EXPIRED_REQUEST_TOMBSTONES {
if let Some(oldest) = expired
.iter()
.min_by_key(|(_, expires_at)| **expires_at)
.map(|(id, _)| *id)
{
expired.remove(&oldest);
}
}
expired.insert(request_id, now + EXPIRED_REQUEST_TOMBSTONE_TTL);
}
}
pub(crate) async fn is_expired_request(dispatcher: &PipeDispatcher, request_id: u32) -> bool {
let mut expired = dispatcher.expired_requests.lock().await;
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.contains_key(&request_id)
}
pub(crate) async fn remove_pending_request(
dispatcher: &PipeDispatcher,
request_id: u32,
@ -149,6 +234,13 @@ pub(crate) async fn remove_pending_request(
}
}
async fn consume_expired_request(dispatcher: &PipeDispatcher, request_id: u32) -> bool {
let mut expired = dispatcher.expired_requests.lock().await;
let now = Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.remove(&request_id).is_some()
}
#[cfg(feature = "pipes")]
pub(crate) async fn run_dispatcher(
receiver: Receiver,
@ -157,14 +249,19 @@ pub(crate) async fn run_dispatcher(
pipe_req_tx: mpsc::Sender<PipeRequest>,
dispatcher: Arc<PipeDispatcher>,
) {
let pipe_req_type = CommunicationType::PipeRequest.try_to_id(&mtp_codec::TypeMap::latest());
let pipe_resp_type = CommunicationType::PipeResponse.try_to_id(&mtp_codec::TypeMap::latest());
loop {
match receiver.receive_event().await {
Ok(mtp_transport::TransportEvent::Message(msg)) => {
if Some(msg.get_type()) == pipe_req_type {
let pipe_id = msg.get_id();
if msg.is_type(CommunicationType::PipeRequest) {
let Some(pipe_id) = msg.id().filter(|id| *id != 0) else {
let error = CommunicationError::Other(
"PipeRequest frame must contain a non-zero id".into(),
);
if app_tx.send(Err(error)).await.is_err() {
break;
}
continue;
};
let description = msg.get_str(DataType::Description).unwrap_or("").to_string();
let req = PipeRequest {
pipe_id,
@ -176,8 +273,16 @@ pub(crate) async fn run_dispatcher(
continue;
}
if Some(msg.get_type()) == pipe_resp_type {
let pipe_id = msg.get_id();
if msg.is_type(CommunicationType::PipeResponse) {
let Some(pipe_id) = msg.id().filter(|id| *id != 0) else {
let error = CommunicationError::Other(
"PipeResponse frame must contain a non-zero id".into(),
);
if app_tx.send(Err(error)).await.is_err() {
break;
}
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) {