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

1
host/Cargo.toml
View file

@ -9,6 +9,7 @@ mtp-codec = { version = "0.3.0", path = "../codec", features = ["registry"] }
mtp-transport = { version = "0.3.0", path = "../transport", features = ["host"] }
mtp-crypto = { version = "0.3.0", path = "../crypto", optional = true }
rand = "0.10"
thiserror = "2"
tokio = { version = "1", features = ["macros", "rt", "time", "sync"] }
tracing = "0.1"
wtransport = "0.7"

230
host/src/config.rs
View file

@ -5,10 +5,14 @@ use std::collections::HashMap;
#[cfg(feature = "crypto")]
use std::collections::HashSet;
#[cfg(feature = "crypto")]
use std::collections::VecDeque;
#[cfg(feature = "crypto")]
use std::pin::Pin;
#[cfg(feature = "crypto")]
use std::sync::{Arc, Mutex};
#[cfg(feature = "crypto")]
use std::time::{Duration as StdDuration, Instant};
#[cfg(feature = "crypto")]
use tokio::time::Duration;
pub use mtp_transport::Policy;
@ -82,6 +86,158 @@ pub enum AuthenticationPolicy {
Unauthenticated,
}
/// Transport-supplied identity used to scope authentication attempt limits.
/// Concrete hosts should populate these fields from the accepted connection;
/// the zero/empty defaults exist only for transport-neutral callers.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AuthenticationContext {
pub peer_network_identity: Option<String>,
pub connection_id: u64,
}
#[cfg(feature = "crypto")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AuthenticationAttempt {
pub peer_network_identity: Option<String>,
pub connection_id: u64,
pub claimed_client_id: Option<u64>,
pub registration: bool,
}
#[cfg(feature = "crypto")]
#[derive(Debug, thiserror::Error)]
pub enum AuthenticationLimitError {
#[error("authentication limiter storage is unavailable")]
Store,
}
#[cfg(feature = "crypto")]
pub trait AuthenticationAttemptLimiter: Send + Sync {
fn allow(&self, context: &AuthenticationAttempt) -> Result<bool, AuthenticationLimitError>;
}
#[cfg(feature = "crypto")]
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
enum AuthenticationLimitKey {
Peer(String),
Connection(u64),
Client(u64),
Registration,
Global,
}
#[cfg(feature = "crypto")]
#[derive(Debug)]
pub struct InMemoryAuthenticationAttemptLimiter {
max_attempts: usize,
window: StdDuration,
max_keys: usize,
by_peer: bool,
by_connection: bool,
by_client: bool,
by_registration: bool,
attempts: Mutex<HashMap<AuthenticationLimitKey, VecDeque<Instant>>>,
}
#[cfg(feature = "crypto")]
impl InMemoryAuthenticationAttemptLimiter {
pub fn new(max_attempts: usize, window: StdDuration) -> Self {
Self {
max_attempts,
window,
max_keys: 100_000,
by_peer: true,
by_connection: true,
by_client: true,
by_registration: true,
attempts: Mutex::new(HashMap::new()),
}
}
pub fn with_keys(
mut self,
by_peer: bool,
by_connection: bool,
by_client: bool,
by_registration: bool,
) -> Self {
self.by_peer = by_peer;
self.by_connection = by_connection;
self.by_client = by_client;
self.by_registration = by_registration;
self
}
pub fn with_max_keys(mut self, max_keys: usize) -> Self {
self.max_keys = max_keys.max(1);
self
}
fn keys(&self, context: &AuthenticationAttempt) -> Vec<AuthenticationLimitKey> {
let mut keys = Vec::with_capacity(5);
if self.by_peer
&& let Some(peer) = context.peer_network_identity.as_ref()
{
keys.push(AuthenticationLimitKey::Peer(peer.clone()));
}
if self.by_connection && context.connection_id != 0 {
keys.push(AuthenticationLimitKey::Connection(context.connection_id));
}
if self.by_client
&& let Some(client_id) = context.claimed_client_id
{
keys.push(AuthenticationLimitKey::Client(client_id));
}
if self.by_registration && context.registration {
keys.push(AuthenticationLimitKey::Registration);
}
// Keep one global bucket as a backstop when an attacker varies the
// claimed client ID or presents no peer/connection identity.
keys.push(AuthenticationLimitKey::Global);
keys
}
}
#[cfg(feature = "crypto")]
impl AuthenticationAttemptLimiter for InMemoryAuthenticationAttemptLimiter {
fn allow(&self, context: &AuthenticationAttempt) -> Result<bool, AuthenticationLimitError> {
if self.max_attempts == 0 {
return Ok(false);
}
let now = Instant::now();
let cutoff = now.checked_sub(self.window);
let keys = self.keys(context);
let mut attempts = self
.attempts
.lock()
.map_err(|_| AuthenticationLimitError::Store)?;
for key in &keys {
if let Some(history) = attempts.get_mut(key) {
while history
.front()
.is_some_and(|timestamp| cutoff.is_some_and(|cutoff| *timestamp <= cutoff))
{
history.pop_front();
}
if history.len() >= self.max_attempts {
return Ok(false);
}
}
}
for key in keys {
if !attempts.contains_key(&key) && attempts.len() >= self.max_keys {
if let Some(oldest) = attempts.keys().next().cloned() {
attempts.remove(&oldest);
}
}
attempts.entry(key).or_default().push_back(now);
}
Ok(true)
}
}
pub struct HostConfig {
pub ip: IpAddr,
pub port: u16,
@ -113,6 +269,10 @@ pub struct HostConfig {
pub complete_register: CompleteRegister,
#[cfg(feature = "crypto")]
pub find_registered_client: Option<FindRegisteredClient>,
#[cfg(feature = "crypto")]
pub auth_limiter: Arc<dyn AuthenticationAttemptLimiter>,
#[cfg(feature = "crypto")]
pub conceal_authentication_identities: bool,
}
impl HostConfig {
@ -153,6 +313,13 @@ impl HostConfig {
complete_register: Box::new(|_, _| Box::pin(async { 0 })),
#[cfg(feature = "crypto")]
find_registered_client: None,
#[cfg(feature = "crypto")]
auth_limiter: Arc::new(InMemoryAuthenticationAttemptLimiter::new(
32,
StdDuration::from_secs(60),
)),
#[cfg(feature = "crypto")]
conceal_authentication_identities: true,
}
}
@ -210,4 +377,67 @@ impl HostConfig {
self.find_registered_client = Some(lookup);
self
}
#[cfg(feature = "crypto")]
pub fn with_authentication_limiter(
mut self,
limiter: Arc<dyn AuthenticationAttemptLimiter>,
) -> Self {
self.auth_limiter = limiter;
self
}
#[cfg(feature = "crypto")]
pub fn with_authentication_identity_concealment(mut self, conceal: bool) -> Self {
self.conceal_authentication_identities = conceal;
self
}
}
#[cfg(all(test, feature = "crypto"))]
mod tests {
use super::*;
#[test]
fn authentication_attempt_limiter_rejects_repeated_attempts() {
let limiter = InMemoryAuthenticationAttemptLimiter::new(1, StdDuration::from_secs(60))
.with_keys(false, true, false, false);
let attempt = AuthenticationAttempt {
peer_network_identity: None,
connection_id: 9,
claimed_client_id: Some(42),
registration: false,
};
assert!(limiter.allow(&attempt).expect("first attempt decision"));
assert!(!limiter.allow(&attempt).expect("second attempt decision"));
}
#[test]
fn authentication_attempt_limiter_can_scope_registration_separately() {
let limiter = InMemoryAuthenticationAttemptLimiter::new(2, StdDuration::from_secs(60))
.with_keys(false, false, false, true);
let login = AuthenticationAttempt {
peer_network_identity: None,
connection_id: 1,
claimed_client_id: None,
registration: false,
};
let registration = AuthenticationAttempt {
registration: true,
..login.clone()
};
assert!(limiter.allow(&login).expect("login attempt decision"));
assert!(
limiter
.allow(&registration)
.expect("registration attempt decision")
);
assert!(
!limiter
.allow(&registration)
.expect("repeated registration decision")
);
}
}

57
host/src/connection.rs
View file

@ -11,7 +11,10 @@ use tokio::sync::{Mutex, mpsc};
#[cfg(feature = "crypto")]
use crate::error::random_client_id;
#[cfg(feature = "pipes")]
use crate::pipe::{PipeDispatcher, PipeReceiver, PipeRequest, PipeSender, run_dispatcher};
use crate::pipe::{
PendingCreationGuard, PipeDispatcher, PipeReceiver, PipeRequest, PipeSender,
is_expired_creation, run_dispatcher,
};
#[cfg(feature = "pipes")]
use mtp_transport::Policy;
@ -145,10 +148,12 @@ where
remote_addr: Option<SocketAddr>,
) -> Self {
let policy = Arc::new(Policy::default());
let (app_tx, app_rx) = mpsc::channel(policy.receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(policy.receiver_queue_capacity);
let receiver_queue_capacity = policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel(receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(receiver_queue_capacity);
let dispatcher = Arc::new(PipeDispatcher {
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,
type_map: codec.type_map().clone(),
@ -198,10 +203,12 @@ where
remote_addr: Option<SocketAddr>,
policy: Arc<Policy>,
) -> Self {
let (app_tx, app_rx) = mpsc::channel(policy.receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(policy.receiver_queue_capacity);
let receiver_queue_capacity = policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel(receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(receiver_queue_capacity);
let dispatcher = Arc::new(PipeDispatcher {
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,
type_map: codec.type_map().clone(),
@ -321,19 +328,37 @@ where
pub async fn create_pipe(
&self,
description: &str,
) -> Result<crate::pipe::PipeHandle<S>, mtp_common::PipeError> {
) -> Result<crate::pipe::PipeHandle<S, P>, mtp_common::PipeError> {
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
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, response_tx);
pipe_id
let token = Arc::new(());
pending.insert(
pipe_id,
crate::pipe::PendingCreation {
token: token.clone(),
sender: response_tx,
},
);
drop(pending);
(pipe_id, token)
};
let (pipe_id, token) = pipe_id;
let mut creation_guard =
PendingCreationGuard::new(self.pipe_dispatcher.clone(), pipe_id, token.clone());
let request = CommunicationValue::new_with_type_map(
CommunicationType::PipeRequest,
@ -342,19 +367,17 @@ where
.with_id(pipe_id)
.add_typed_default(DataType::Description, DataValue::Str(description.into()));
if let Err(error) = self.sender.send_pipe_message(&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_owned(),
sender: self.sender.clone(),
response_rx,
dispatcher: self.pipe_dispatcher.clone(),
token,
})
}

159
host/src/engine.rs
View file

@ -4,7 +4,7 @@
//! and the web server's `MTPWebServer` to perform the MTP opening handshake,
//! version negotiation, authentication, and guest assignment.
use crate::config::HostConfig;
use crate::config::{AuthenticationContext, HostConfig};
use crate::error::AcceptError;
use mtp_codec::{
CommunicationType, CommunicationValue, DataType, DataValue, TypeMap, Version,
@ -127,19 +127,32 @@ impl HandshakeEngine {
&self,
sender: &S,
receiver: &R,
) -> Result<HandshakeResult, AcceptError> {
self.accept_with_context(sender, receiver, AuthenticationContext::default())
.await
}
/// Run the opening handshake with transport-provided authentication
/// scoping information.
pub async fn accept_with_context<S: HandshakeSender, R: HandshakeReceiver>(
&self,
sender: &S,
receiver: &R,
context: AuthenticationContext,
) -> Result<HandshakeResult, AcceptError> {
#[cfg(feature = "crypto")]
{
self.accept_until(
self.accept_until_with_context(
sender,
receiver,
tokio::time::Instant::now() + self.config.auth_timeout,
context,
)
.await
}
#[cfg(not(feature = "crypto"))]
{
let result = self.accept_inner(sender, receiver).await;
let result = self.accept_inner(sender, receiver, &context).await;
if result.is_err() {
sender.close();
}
@ -159,7 +172,22 @@ impl HandshakeEngine {
receiver: &R,
deadline: tokio::time::Instant,
) -> Result<HandshakeResult, AcceptError> {
match tokio::time::timeout_at(deadline, self.accept_inner(sender, receiver)).await {
self.accept_until_with_context(sender, receiver, deadline, AuthenticationContext::default())
.await
}
/// Run the crypto handshake until a deadline with transport-provided
/// authentication scoping information.
#[cfg(feature = "crypto")]
pub async fn accept_until_with_context<S: HandshakeSender, R: HandshakeReceiver>(
&self,
sender: &S,
receiver: &R,
deadline: tokio::time::Instant,
context: AuthenticationContext,
) -> Result<HandshakeResult, AcceptError> {
match tokio::time::timeout_at(deadline, self.accept_inner(sender, receiver, &context)).await
{
Ok(result) => {
if result.is_err() {
sender.close();
@ -186,6 +214,7 @@ impl HandshakeEngine {
&self,
sender: &S,
receiver: &R,
_authentication_context: &AuthenticationContext,
) -> Result<HandshakeResult, AcceptError> {
let mut first_msg = receiver.receive().await.map_err(AcceptError::Receive)?;
@ -257,6 +286,42 @@ impl HandshakeEngine {
#[cfg(feature = "crypto")]
{
let claimed_client_id = match first_msg.get_data(DataType::Id) {
Some(DataValue::UnsignedNumber(value)) => u64::try_from(*value).ok(),
_ => None,
};
let registration = Some(first_msg.get_type())
== CommunicationType::Register.try_to_id(codec.type_map());
let authentication_requested = matches!(
self.config.authentication_policy,
crate::config::AuthenticationPolicy::ForceAuthentication
) || registration
|| first_msg.get_data(DataType::PublicKeys).is_some()
|| claimed_client_id.is_some_and(|client_id| client_id != 0);
if authentication_requested {
let attempt = crate::config::AuthenticationAttempt {
peer_network_identity: _authentication_context.peer_network_identity.clone(),
connection_id: _authentication_context.connection_id,
claimed_client_id,
registration,
};
match self.config.auth_limiter.allow(&attempt) {
Ok(true) => {}
Ok(false) | Err(_) => {
let error =
AcceptError::AuthenticationFailed("authentication rejected".into());
send_rejection_generic(
sender,
RejectionReason::RateLimited,
Some(codec.type_map()),
)
.await;
sender.close();
return Err(error);
}
}
}
match self.config.authentication_policy {
crate::config::AuthenticationPolicy::ForceAuthentication => {
self.force_auth_handshake(
@ -408,7 +473,14 @@ impl HandshakeEngine {
return Err(error);
}
};
let pk_bytes = bundle.as_bytes();
let pk_bytes = match bundle.try_as_bytes() {
Ok(bytes) => bytes,
Err(error) => {
let error = AcceptError::AuthenticationFailed(error.to_string());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
return self
.complete_auth_handshake(
sender,
@ -453,6 +525,28 @@ impl HandshakeEngine {
// Unknown or zero ID: an Identification carrying PublicKeys is an
// explicit authentication attempt, not a guest connection.
if first_msg.get_data(DataType::PublicKeys).is_some() {
if self.config.conceal_authentication_identities && cid > 0 {
/* Keep an unknown authenticated ID on the same
challenge/proof path as a known ID. The fixed host
identity makes the eventual proof fail without
disclosing whether the lookup succeeded. */
return self
.complete_auth_handshake(
sender,
receiver,
Flow::Login {
id: cid,
bundle: self.config.host_keyring.public_key_bundle(),
},
CommunicationType::IdentificationResponse,
&negotiated,
&codec,
description,
version_str,
client_version,
)
.await;
}
let error = AcceptError::AuthenticationFailed(
"unknown authenticated client identity".into(),
);
@ -525,20 +619,29 @@ impl HandshakeEngine {
let bundle = match (self.config.get_existing_client)(cid, description.clone()).await {
Some(b) => b,
None => {
let rejection = CommunicationValue::new_with_type_map(
CommunicationType::IdentificationResponse,
tm,
)
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
.add_typed_default(
DataType::ErrorMessage,
DataValue::Str("unknown client id".into()),
);
let _ = sender.send(&rejection).await;
sender.close();
return Err(AcceptError::AuthenticationFailed(
"unknown client id".into(),
));
if self.config.conceal_authentication_identities {
// Use a valid fixed-cost dummy identity so an unknown
// client follows the same challenge/proof sequence as
// a registered client. The host public bundle is
// already public and the peer cannot produce its
// private-key proof.
self.config.host_keyring.public_key_bundle()
} else {
let rejection = CommunicationValue::new_with_type_map(
CommunicationType::IdentificationResponse,
tm,
)
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
.add_typed_default(
DataType::ErrorMessage,
DataValue::Str("unknown client id".into()),
);
let _ = sender.send(&rejection).await;
sender.close();
return Err(AcceptError::AuthenticationFailed(
"unknown client id".into(),
));
}
}
};
(
@ -553,7 +656,14 @@ impl HandshakeEngine {
return Err(error);
}
};
let pk_bytes = bundle.as_bytes();
let pk_bytes = match bundle.try_as_bytes() {
Ok(bytes) => bytes,
Err(error) => {
let error = AcceptError::AuthenticationFailed(error.to_string());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
(
Flow::Register { bundle, pk_bytes },
CommunicationType::RegisterResponse,
@ -787,7 +897,14 @@ impl HandshakeEngine {
Flow::Login { id, bundle } => (id, bundle),
Flow::Register { bundle, .. } => {
let _registration_guard = self.config.registration_lock.lock().await;
let identity = bundle.as_bytes();
let identity = match bundle.try_as_bytes() {
Ok(identity) => identity,
Err(error) => {
let error = AcceptError::AuthenticationFailed(error.to_string());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let cached_id = self
.config
.registration_ids

31
host/src/handshake.rs
View file

@ -11,7 +11,7 @@ use std::time::Instant;
#[cfg(feature = "pipes")]
use tokio::sync::mpsc;
use crate::config::HostConfig;
use crate::config::{AuthenticationContext, HostConfig};
use crate::connection::MTPConnection;
use crate::engine::HandshakeEngine;
use crate::error::AcceptError;
@ -135,7 +135,16 @@ impl HandshakeContext {
receiver: Receiver,
) -> Result<Option<MTPConnection>, AcceptError> {
let engine = HandshakeEngine::new(self.registry.clone(), self.config.clone());
let result = engine.accept(&sender, &receiver).await?;
let authentication_context = AuthenticationContext {
peer_network_identity: sender
.handle()
.remote_addr()
.map(|address| address.to_string()),
connection_id: sender.handle().connection_id(),
};
let result = engine
.accept_with_context(&sender, &receiver, authentication_context)
.await?;
#[cfg(feature = "crypto")]
{
Ok(Some(self.connection_from_handshake_result(
@ -171,12 +180,13 @@ impl HandshakeContext {
receiver.respond_to_pings(sender.clone());
}
let (app_tx, app_rx) = mpsc::channel(self.config.policy.receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) =
mpsc::channel(self.config.policy.receiver_queue_capacity);
let receiver_queue_capacity = self.config.policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel(receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(receiver_queue_capacity);
let dispatcher = Arc::new(PipeDispatcher {
pending_creations: tokio::sync::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: tokio::sync::Mutex::new(std::collections::HashMap::new()),
policy: Arc::new(self.config.policy),
type_map: type_map.clone(),
@ -260,12 +270,13 @@ impl HandshakeContext {
receiver.respond_to_pings(sender.clone());
}
let (app_tx, app_rx) = mpsc::channel(self.config.policy.receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) =
mpsc::channel(self.config.policy.receiver_queue_capacity);
let receiver_queue_capacity = self.config.policy.receiver_queue_capacity.max(1);
let (app_tx, app_rx) = mpsc::channel(receiver_queue_capacity);
let (pipe_req_tx, pipe_req_rx) = mpsc::channel(receiver_queue_capacity);
let dispatcher = Arc::new(PipeDispatcher {
pending_creations: tokio::sync::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: tokio::sync::Mutex::new(std::collections::HashMap::new()),
policy: Arc::new(self.config.policy),
type_map,

5
host/src/lib.rs
View file

@ -29,8 +29,9 @@ pub use mtp_codec::registry::Registry;
#[cfg(feature = "crypto")]
pub use config::{
AuthenticationPolicy, CompleteRegister, FindRegisteredClient, GetExistingClient,
GuestIdGenerator,
AuthenticationAttempt, AuthenticationAttemptLimiter, AuthenticationContext,
AuthenticationLimitError, AuthenticationPolicy, CompleteRegister, FindRegisteredClient,
GetExistingClient, GuestIdGenerator, InMemoryAuthenticationAttemptLimiter,
};
#[cfg(feature = "crypto")]
pub use error::AuthState;

185
host/src/pipe.rs
View file

@ -3,6 +3,7 @@ use mtp_common::{CommunicationError, PipeError};
use mtp_transport::{PipeReader, PipeWriter, Policy, TransportEvent};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use tokio::sync::{Mutex, mpsc};
/// The sender operations needed by the transport-independent pipe protocol.
@ -93,14 +94,20 @@ where
}
}
pub struct PipeHandle<S: PipeSender> {
pub struct PipeHandle<S: PipeSender, P = wtransport::RecvStream> {
pub(crate) pipe_id: u32,
pub(crate) description: String,
pub(crate) sender: S,
pub(crate) response_rx: tokio::sync::oneshot::Receiver<Result<bool, PipeError>>,
pub(crate) dispatcher: Arc<PipeDispatcher<P>>,
pub(crate) token: Arc<()>,
}
impl<S: PipeSender> PipeHandle<S> {
impl<S, P> PipeHandle<S, P>
where
S: PipeSender,
P: tokio::io::AsyncRead + Send + Unpin + 'static,
{
pub fn pipe_id(&self) -> u32 {
self.pipe_id
}
@ -109,21 +116,42 @@ impl<S: PipeSender> PipeHandle<S> {
&self.description
}
pub async fn wait(self) -> Result<Option<PipeWriter<S::Writer>>, PipeError> {
match self.response_rx.await {
Ok(Ok(true)) => self
pub async fn wait(mut self) -> Result<Option<PipeWriter<S::Writer>>, PipeError> {
let response =
tokio::time::timeout(self.dispatcher.policy.read_timeout, &mut self.response_rx).await;
match response {
Ok(Ok(Ok(true))) => self
.sender
.open_pipe_stream(self.pipe_id, &self.description)
.await
.map(Some)
.map_err(PipeError::from),
Ok(Ok(false)) => Ok(None),
Ok(Err(error)) => Err(error),
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)
}
}
}
}
impl<S, P> Drop for PipeHandle<S, P>
where
S: PipeSender,
{
fn drop(&mut self) {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
}
}
pub struct PipeRequest<S, P> {
pub(crate) pipe_id: u32,
pub(crate) description: String,
@ -203,13 +231,132 @@ where
}
pub(crate) struct PipeDispatcher<P> {
pub(crate) pending_creations:
Mutex<HashMap<u32, tokio::sync::oneshot::Sender<Result<bool, PipeError>>>>,
pub(crate) pending_creations: StdMutex<HashMap<u32, PendingCreation>>,
pub(crate) expired_creations: StdMutex<HashMap<u32, tokio::time::Instant>>,
pub(crate) pending_pipes: Mutex<HashMap<u32, tokio::sync::oneshot::Sender<PipeReader<P>>>>,
pub(crate) policy: Arc<Policy>,
pub(crate) type_map: TypeMap,
}
pub(crate) struct PendingCreation {
pub(crate) token: Arc<()>,
pub(crate) sender: tokio::sync::oneshot::Sender<Result<bool, PipeError>>,
}
pub(crate) struct PendingCreationGuard<P> {
dispatcher: Arc<PipeDispatcher<P>>,
pipe_id: u32,
token: Arc<()>,
armed: bool,
}
impl<P> PendingCreationGuard<P> {
pub(crate) fn new(dispatcher: Arc<PipeDispatcher<P>>, pipe_id: u32, token: Arc<()>) -> Self {
Self {
dispatcher,
pipe_id,
token,
armed: true,
}
}
pub(crate) fn disarm(&mut self) {
self.armed = false;
}
}
impl<P> Drop for PendingCreationGuard<P> {
fn drop(&mut self) {
if self.armed {
expire_pending_creation(&self.dispatcher, self.pipe_id, &self.token);
}
}
}
const EXPIRED_CREATION_TOMBSTONE_TTL: tokio::time::Duration = tokio::time::Duration::from_secs(60);
const MAX_EXPIRED_CREATION_TOMBSTONES: usize = 1024;
pub(crate) fn expire_pending_creation<P>(
dispatcher: &PipeDispatcher<P>,
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 = tokio::time::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);
}
fn consume_expired_creation<P>(dispatcher: &PipeDispatcher<P>, pipe_id: u32) -> bool {
let Ok(mut expired) = dispatcher.expired_creations.lock() else {
return false;
};
let now = tokio::time::Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.remove(&pipe_id).is_some()
}
pub(crate) fn is_expired_creation<P>(dispatcher: &PipeDispatcher<P>, pipe_id: u32) -> bool {
let Ok(mut expired) = dispatcher.expired_creations.lock() else {
return true;
};
let now = tokio::time::Instant::now();
expired.retain(|_, expires_at| *expires_at > now);
expired.contains_key(&pipe_id)
}
pub(crate) fn fail_pending_creations<P>(
dispatcher: &PipeDispatcher<P>,
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();
}
}
pub(crate) async fn fail_pending_pipes<P>(dispatcher: &PipeDispatcher<P>) {
dispatcher.pending_pipes.lock().await.clear();
}
pub(crate) async fn run_dispatcher<S, R, P>(
receiver: R,
sender: S,
@ -256,10 +403,17 @@ pub(crate) async fn run_dispatcher<S, R, P>(
}
continue;
};
let mut pending = dispatcher.pending_creations.lock().await;
if let Some(reply) = pending.remove(&pipe_id) {
let _ =
reply.send(Ok(message.get_bool(DataType::Accepted).unwrap_or(false)));
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(message.get_bool(DataType::Accepted).unwrap_or(false)));
} else if consume_expired_creation(&dispatcher, pipe_id) {
tracing::debug!(pipe_id, "ignored late pipe creation response");
}
continue;
}
@ -297,9 +451,12 @@ pub(crate) async fn run_dispatcher<S, R, P>(
let _ = pipe_req_tx.send(request).await;
}
Err(error) => {
fail_pending_creations(&dispatcher, &error);
fail_pending_pipes(&dispatcher).await;
if app_tx.send(Err(error)).await.is_err() {
break;
}
break;
}
}
}