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

60
host/src/config.rs
View file

@ -1,8 +1,14 @@
use std::net::IpAddr;
#[cfg(feature = "crypto")]
use std::collections::HashMap;
#[cfg(feature = "crypto")]
use std::collections::HashSet;
#[cfg(feature = "crypto")]
use std::pin::Pin;
#[cfg(feature = "crypto")]
use std::sync::{Arc, Mutex};
#[cfg(feature = "crypto")]
use tokio::time::Duration;
pub use mtp_transport::Policy;
@ -26,18 +32,23 @@ pub type GetExistingClient = Box<
/// Callback that assigns a guest (unauthenticated) client ID.
///
/// Return `Some(id)` to accept the guest with the given ID, or `None` to reject
/// the connection. The returned ID must fit in 48 bits
/// (`id <= mtp_codec::MAX_WIRE_ID`); values outside that range are rejected
/// automatically.
/// Return `Some(id)` to accept the guest with the given full-width `u64` ID, or
/// `None` to reject the connection.
///
/// When set to `None` on `HostConfig`, the built-in generator produces a random
/// 48-bit ID that avoids collisions with registered clients.
/// full-width non-zero ID that avoids collisions with registered clients and
/// currently connected guests.
#[cfg(feature = "crypto")]
pub type GuestIdGenerator =
Box<dyn Fn() -> Pin<Box<dyn std::future::Future<Output = Option<u64>> + Send>> + Send + Sync>;
#[cfg(feature = "crypto")]
/// Callback that commits a new registration and returns its non-zero ID.
///
/// The host serializes registration commits and remembers successful identity
/// assignments for the lifetime of the host. Applications that need retry
/// recovery across a host restart should also configure [`FindRegisteredClient`]
/// to look up the public identity in persistent storage.
pub type CompleteRegister = Box<
dyn Fn(
mtp_crypto::PublicKeyBundle,
@ -47,6 +58,22 @@ pub type CompleteRegister = Box<
+ Sync,
>;
/// Callback that recovers an existing registration by its public identity.
///
/// Returning an ID makes a registration retry idempotent: the host can send
/// the same final response when the original response was lost after the
/// application committed the registration. Returning `None` asks the host to
/// invoke [`CompleteRegister`] for a new registration.
#[cfg(feature = "crypto")]
pub type FindRegisteredClient = Box<
dyn Fn(
mtp_crypto::PublicKeyBundle,
Option<String>,
) -> Pin<Box<dyn std::future::Future<Output = Option<u64>> + Send>>
+ Send
+ Sync,
>;
#[cfg(feature = "crypto")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthenticationPolicy {
@ -75,9 +102,17 @@ pub struct HostConfig {
#[cfg(feature = "crypto")]
pub get_existing_client: GetExistingClient,
#[cfg(feature = "crypto")]
pub(crate) active_guest_ids: Arc<Mutex<HashSet<u64>>>,
#[cfg(feature = "crypto")]
pub(crate) registration_ids: Arc<Mutex<HashMap<Vec<u8>, u64>>>,
#[cfg(feature = "crypto")]
pub(crate) registration_lock: Arc<tokio::sync::Mutex<()>>,
#[cfg(feature = "crypto")]
pub guest_id_generator: Option<GuestIdGenerator>,
#[cfg(feature = "crypto")]
pub complete_register: CompleteRegister,
#[cfg(feature = "crypto")]
pub find_registered_client: Option<FindRegisteredClient>,
}
impl HostConfig {
@ -107,9 +142,17 @@ impl HostConfig {
#[cfg(feature = "crypto")]
get_existing_client: Box::new(|_, _| Box::pin(async { None })),
#[cfg(feature = "crypto")]
active_guest_ids: Arc::new(Mutex::new(HashSet::new())),
#[cfg(feature = "crypto")]
registration_ids: Arc::new(Mutex::new(HashMap::new())),
#[cfg(feature = "crypto")]
registration_lock: Arc::new(tokio::sync::Mutex::new(())),
#[cfg(feature = "crypto")]
guest_id_generator: None,
#[cfg(feature = "crypto")]
complete_register: Box::new(|_, _| Box::pin(async { 0 })),
#[cfg(feature = "crypto")]
find_registered_client: None,
}
}
@ -160,4 +203,11 @@ impl HostConfig {
self.guest_id_generator = Some(generator);
self
}
/// Configure the lookup used to make registration retries idempotent.
#[cfg(feature = "crypto")]
pub fn with_registration_lookup(mut self, lookup: FindRegisteredClient) -> Self {
self.find_registered_client = Some(lookup);
self
}
}

67
host/src/connection.rs
View file

@ -77,12 +77,30 @@ pub struct MTPConnection<
pub(crate) _pipe_stream: std::marker::PhantomData<P>,
pub description: Option<String>,
pub(crate) _dispatcher_task: tokio::task::JoinHandle<()>,
/// Keeps an outer server admission permit alive for this MTP session.
/// Native hosts leave it empty; WebTransport hosts use it to make the
/// configured connection limit cover the session lifetime.
pub(crate) _connection_guard: Option<tokio::sync::OwnedSemaphorePermit>,
#[cfg(feature = "crypto")]
pub auth_state: crate::error::AuthState,
#[cfg(feature = "crypto")]
pub client_id: u64,
#[cfg(feature = "crypto")]
pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
#[cfg(feature = "crypto")]
pub(crate) guest_id_lease: Option<crate::engine::GuestIdLease>,
}
impl<S, R, P> MTPConnection<S, R, P> {
/// Keep an outer server admission permit until this connection is dropped.
pub fn set_connection_guard(&mut self, guard: tokio::sync::OwnedSemaphorePermit) {
self._connection_guard = Some(guard);
}
#[cfg(feature = "crypto")]
pub fn set_guest_id_lease(&mut self, lease: Option<crate::engine::GuestIdLease>) {
self.guest_id_lease = lease;
}
}
#[cfg(feature = "pipes")]
@ -133,6 +151,7 @@ where
pending_creations: Mutex::new(std::collections::HashMap::new()),
pending_pipes: Mutex::new(std::collections::HashMap::new()),
policy,
type_map: codec.type_map().clone(),
});
let task = tokio::spawn(run_dispatcher(
receiver.clone(),
@ -153,12 +172,15 @@ where
pipe_dispatcher: dispatcher,
description,
_dispatcher_task: task,
_connection_guard: None,
#[cfg(feature = "crypto")]
auth_state: crate::error::AuthState::Unauthenticated,
#[cfg(feature = "crypto")]
client_id: random_client_id(),
#[cfg(feature = "crypto")]
client_public_key: None,
#[cfg(feature = "crypto")]
guest_id_lease: None,
}
}
@ -182,6 +204,7 @@ where
pending_creations: Mutex::new(std::collections::HashMap::new()),
pending_pipes: Mutex::new(std::collections::HashMap::new()),
policy,
type_map: codec.type_map().clone(),
});
let task = tokio::spawn(run_dispatcher(
receiver.clone(),
@ -202,12 +225,15 @@ where
pipe_dispatcher: dispatcher,
description,
_dispatcher_task: task,
_connection_guard: None,
#[cfg(feature = "crypto")]
auth_state: crate::error::AuthState::Unauthenticated,
#[cfg(feature = "crypto")]
client_id: random_client_id(),
#[cfg(feature = "crypto")]
client_public_key: None,
#[cfg(feature = "crypto")]
guest_id_lease: None,
}
}
}
@ -252,12 +278,15 @@ impl<S, R, P> MTPConnection<S, R, P> {
description,
_pipe_stream: std::marker::PhantomData,
_dispatcher_task: tokio::spawn(async {}),
_connection_guard: None,
#[cfg(feature = "crypto")]
auth_state: crate::error::AuthState::Unauthenticated,
#[cfg(feature = "crypto")]
client_id: random_client_id(),
#[cfg(feature = "crypto")]
client_public_key: None,
#[cfg(feature = "crypto")]
guest_id_lease: None,
}
}
}
@ -293,21 +322,33 @@ where
&self,
description: &str,
) -> Result<crate::pipe::PipeHandle<S>, mtp_common::PipeError> {
let pipe_id = rand::random::<u32>();
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
self.pipe_dispatcher
.pending_creations
.lock()
.await
.insert(pipe_id, response_tx);
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, response_tx);
pipe_id
};
let request = CommunicationValue::new(CommunicationType::PipeRequest)
.with_id(pipe_id)
.add_typed_default(DataType::Description, DataValue::Str(description.into()));
self.sender
.send_pipe_message(&request)
.await
.map_err(mtp_common::PipeError::from)?;
let request = CommunicationValue::new_with_type_map(
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_pipe_message(&request).await {
self.pipe_dispatcher
.pending_creations
.lock()
.await
.remove(&pipe_id);
return Err(mtp_common::PipeError::from(error));
}
Ok(crate::pipe::PipeHandle {
pipe_id,

519
host/src/engine.rs
View file

@ -7,11 +7,15 @@
use crate::config::HostConfig;
use crate::error::AcceptError;
use mtp_codec::{
CommunicationType, CommunicationValue, DataType, DataValue, Version,
CommunicationType, CommunicationValue, DataType, DataValue, TypeMap, Version,
registry::{Registry, VersionedCodec},
};
use mtp_common::{CommunicationError, RejectionReason};
#[cfg(feature = "crypto")]
use std::collections::HashSet;
use std::sync::Arc;
#[cfg(feature = "crypto")]
use std::sync::Mutex;
/// Trait for sending handshake messages during the opening exchange.
///
@ -24,6 +28,9 @@ pub trait HandshakeSender: Send + Sync {
fn finish_stream(
&self,
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send;
fn set_type_map(&self, _type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async {}
}
fn close(&self);
}
@ -34,6 +41,37 @@ pub trait HandshakeReceiver: Send + Sync {
fn receive(
&self,
) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>> + Send;
/// Bind subsequently decoded frames to the negotiated type map.
///
/// The opening frame must be decoded with the transport's bootstrap map so
/// that it can reveal the version. Once negotiation succeeds, all later
/// frames—including the remainder of the authentication exchange—must use
/// the negotiated map rather than whichever map happens to be latest at
/// compile time.
fn set_type_map(&self, _type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async {}
}
}
/// A non-zero guest ID reserved for the lifetime of a connected session.
///
/// The lease is moved into the resulting `MTPConnection`, so dropping that
/// connection releases the ID for a later guest session.
#[cfg(feature = "crypto")]
#[derive(Debug)]
pub struct GuestIdLease {
active_ids: Arc<Mutex<HashSet<u64>>>,
id: u64,
}
#[cfg(feature = "crypto")]
impl Drop for GuestIdLease {
fn drop(&mut self) {
if let Ok(mut active_ids) = self.active_ids.lock() {
active_ids.remove(&self.id);
}
}
}
/// The result of a successful handshake, containing everything needed to
@ -49,6 +87,8 @@ pub struct HandshakeResult {
pub client_id: u64,
#[cfg(feature = "crypto")]
pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
#[cfg(feature = "crypto")]
pub guest_id_lease: Option<GuestIdLease>,
}
/// Transport-independent handshake state machine.
@ -90,13 +130,56 @@ impl HandshakeEngine {
) -> Result<HandshakeResult, AcceptError> {
#[cfg(feature = "crypto")]
{
let timeout = self.config.auth_timeout;
tokio::time::timeout(timeout, self.accept_inner(sender, receiver))
.await
.unwrap_or(Err(AcceptError::AuthenticationTimedOut))
self.accept_until(
sender,
receiver,
tokio::time::Instant::now() + self.config.auth_timeout,
)
.await
}
#[cfg(not(feature = "crypto"))]
self.accept_inner(sender, receiver).await
{
let result = self.accept_inner(sender, receiver).await;
if result.is_err() {
sender.close();
}
result
}
}
/// Run the crypto handshake until an absolute deadline.
///
/// WebTransport authentication may wait for a shared semaphore before it
/// reaches this engine. Passing the deadline through keeps that queueing
/// time from silently granting the handshake another full timeout.
#[cfg(feature = "crypto")]
pub async fn accept_until<S: HandshakeSender, R: HandshakeReceiver>(
&self,
sender: &S,
receiver: &R,
deadline: tokio::time::Instant,
) -> Result<HandshakeResult, AcceptError> {
match tokio::time::timeout_at(deadline, self.accept_inner(sender, receiver)).await {
Ok(result) => {
if result.is_err() {
sender.close();
}
result
}
Err(_) => {
let error = AcceptError::AuthenticationTimedOut;
send_rejection_generic(
sender,
RejectionReason::AuthenticationFailed {
detail: error.to_string(),
},
None,
)
.await;
sender.close();
Err(error)
}
}
}
async fn accept_inner<S: HandshakeSender, R: HandshakeReceiver>(
@ -104,16 +187,17 @@ impl HandshakeEngine {
sender: &S,
receiver: &R,
) -> Result<HandshakeResult, AcceptError> {
let first_msg = receiver.receive().await.map_err(AcceptError::Receive)?;
let mut first_msg = receiver.receive().await.map_err(AcceptError::Receive)?;
let version_str = match first_msg.get_data(DataType::Version) {
DataValue::Str(s) => s.clone(),
Some(DataValue::Str(s)) => s.clone(),
_ => {
send_rejection_generic(
sender,
RejectionReason::AuthenticationFailed {
detail: "opening message omitted a valid protocol version".into(),
},
None,
)
.await;
sender.close();
@ -128,6 +212,7 @@ impl HandshakeEngine {
RejectionReason::AuthenticationFailed {
detail: "opening message omitted a valid protocol version".into(),
},
None,
)
.await;
sender.close();
@ -150,6 +235,7 @@ impl HandshakeEngine {
.map(|v| v.to_string())
.collect(),
},
None,
)
.await;
sender.close();
@ -160,8 +246,12 @@ impl HandshakeEngine {
let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone())
.ok_or_else(|| AcceptError::UnsupportedVersion(negotiated.clone()))?;
sender.set_type_map(codec.type_map()).await;
receiver.set_type_map(codec.type_map()).await;
first_msg.set_type_map(codec.type_map());
let description = match first_msg.get_data(DataType::Description) {
DataValue::Str(s) => Some(s.clone()),
Some(DataValue::Str(s)) => Some(s.clone()),
_ => None,
};
@ -209,9 +299,28 @@ impl HandshakeEngine {
#[cfg(not(feature = "crypto"))]
{
let _ = sender;
let authentication_requested = Some(first_msg.get_type())
== CommunicationType::Register.try_to_id(codec.type_map())
|| first_msg.get_data(DataType::PublicKeys).is_some();
if authentication_requested {
let error = AcceptError::AuthenticationFailed(
"authentication is unavailable on this non-crypto host".into(),
);
send_rejection_generic(
sender,
RejectionReason::AuthenticationFailed {
detail: error.to_string(),
},
Some(codec.type_map()),
)
.await;
sender.close();
return Err(error);
}
let _ = receiver;
let _ = first_msg;
send_accepted_generic(sender, &negotiated, codec.type_map(), Some(0))
.await
.map_err(AcceptError::Send)?;
Ok(HandshakeResult {
negotiated_version: negotiated,
codec,
@ -229,15 +338,21 @@ impl HandshakeEngine {
codec: VersionedCodec,
description: Option<String>,
) -> Result<HandshakeResult, AcceptError> {
let tm = mtp_codec::TypeMap::latest();
let tm = codec.type_map();
// Reject Register frames on unauthenticated hosts
if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) {
// Reject explicit authentication attempts on unauthenticated hosts.
// Authenticated clients include PublicKeys in Identification as an
// intent marker; this avoids acknowledging the opening as a guest
// connection and leaving the client waiting for a Challenge.
if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm)
|| first_msg.get_data(DataType::PublicKeys).is_some()
{
send_rejection_generic(
sender,
RejectionReason::AuthenticationFailed {
detail: "authentication not allowed on this host".into(),
},
Some(tm),
)
.await;
sender.close();
@ -246,8 +361,15 @@ impl HandshakeEngine {
));
}
let guest_id = self.assign_guest_id().await?;
send_accepted_generic(sender, &negotiated, Some(guest_id))
let guest_id_lease = match self.assign_guest_id().await {
Ok(lease) => lease,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let guest_id = guest_id_lease.id;
send_accepted_generic(sender, &negotiated, tm, Some(guest_id))
.await
.map_err(AcceptError::Send)?;
@ -258,6 +380,7 @@ impl HandshakeEngine {
auth_state: crate::error::AuthState::Unauthenticated,
client_id: guest_id,
client_public_key: None,
guest_id_lease: Some(guest_id_lease),
})
}
@ -274,11 +397,17 @@ impl HandshakeEngine {
version_str: &str,
client_version: Version,
) -> Result<HandshakeResult, AcceptError> {
let tm = mtp_codec::TypeMap::latest();
let tm = codec.type_map();
// Register frames always go through full authentication
if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) {
let bundle = extract_register_bundle(&first_msg)?;
let bundle = match extract_register_bundle(&first_msg) {
Ok(bundle) => bundle,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let pk_bytes = bundle.as_bytes();
return self
.complete_auth_handshake(
@ -298,7 +427,7 @@ impl HandshakeEngine {
// Identification: try lookup, fall back to guest
if Some(first_msg.get_type()) == CommunicationType::Identification.try_to_id(&tm) {
let cid = match first_msg.get_data(DataType::Id) {
DataValue::UnsignedNumber(n) => *n as u64,
Some(DataValue::UnsignedNumber(n)) => u64::try_from(*n).unwrap_or(0),
_ => 0,
};
@ -321,9 +450,26 @@ impl HandshakeEngine {
.await;
}
// 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() {
let error = AcceptError::AuthenticationFailed(
"unknown authenticated client identity".into(),
);
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
// Unknown or zero ID: fall back to guest
let guest_id = self.assign_guest_id().await?;
send_accepted_generic(sender, &negotiated, Some(guest_id))
let guest_id_lease = match self.assign_guest_id().await {
Ok(lease) => lease,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let guest_id = guest_id_lease.id;
send_accepted_generic(sender, &negotiated, tm, Some(guest_id))
.await
.map_err(AcceptError::Send)?;
return Ok(HandshakeResult {
@ -333,13 +479,13 @@ impl HandshakeEngine {
auth_state: crate::error::AuthState::Unauthenticated,
client_id: guest_id,
client_public_key: None,
guest_id_lease: Some(guest_id_lease),
});
}
sender.close();
Err(AcceptError::AuthenticationFailed(
"unexpected message type".into(),
))
let error = AcceptError::AuthenticationFailed("unexpected message type".into());
reject_error_generic(sender, &error, tm).await;
Err(error)
}
#[cfg(feature = "crypto")]
@ -355,30 +501,39 @@ impl HandshakeEngine {
version_str: &str,
client_version: Version,
) -> Result<HandshakeResult, AcceptError> {
let tm = mtp_codec::TypeMap::latest();
let tm = codec.type_map();
let (flow, response_type) = if Some(first_msg.get_type())
== CommunicationType::Identification.try_to_id(&tm)
{
let cid = match first_msg.get_data(DataType::Id) {
DataValue::UnsignedNumber(n) => *n as u64,
Some(DataValue::UnsignedNumber(n)) => match u64::try_from(*n) {
Ok(id) => id,
Err(_) => {
let error =
AcceptError::AuthenticationFailed("client id is out of range".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
},
_ => {
sender.close();
return Err(AcceptError::AuthenticationFailed(
"missing client id".into(),
));
let error = AcceptError::AuthenticationFailed("missing client id".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let bundle = match (self.config.get_existing_client)(cid, description.clone()).await {
Some(b) => b,
None => {
let rejection =
CommunicationValue::new(CommunicationType::IdentificationResponse)
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
.add_typed_default(
DataType::ErrorMessage,
DataValue::Str("unknown client id".into()),
);
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(
@ -391,17 +546,23 @@ impl HandshakeEngine {
CommunicationType::IdentificationResponse,
)
} else if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) {
let bundle = extract_register_bundle(&first_msg)?;
let bundle = match extract_register_bundle(&first_msg) {
Ok(bundle) => bundle,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let pk_bytes = bundle.as_bytes();
(
Flow::Register { bundle, pk_bytes },
CommunicationType::RegisterResponse,
)
} else {
sender.close();
return Err(AcceptError::AuthenticationFailed(
"unexpected authentication message".into(),
));
let error =
AcceptError::AuthenticationFailed("unexpected authentication message".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
};
self.complete_auth_handshake(
@ -434,7 +595,7 @@ impl HandshakeEngine {
) -> Result<HandshakeResult, AcceptError> {
use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth, verify_ed25519};
let tm = mtp_codec::TypeMap::latest();
let tm = codec.type_map();
// PQ preflight: host requiring PQ must have a PQ key
let pq_enabled = !self
@ -457,6 +618,7 @@ impl HandshakeEngine {
RejectionReason::AuthenticationFailed {
detail: "host requires PQ authentication but has no PQ signing key".into(),
},
Some(tm),
)
.await;
sender.close();
@ -468,11 +630,17 @@ impl HandshakeEngine {
// Initialize host signers
let host_pq_signer = if pq_enabled {
Some(Arc::new(
MlDsaSigner::new(
match MlDsaSigner::new(
&self.config.host_keyring.sig_pq_secret_key,
&self.config.host_keyring.sig_pq_public_key,
)
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?,
) {
Ok(signer) => signer,
Err(error) => {
let error = AcceptError::AuthenticationFailed(error.to_string());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
},
))
} else {
None
@ -505,14 +673,21 @@ impl HandshakeEngine {
let server_challenge: u128 = rand::random();
let (chal_sig, chal_pq_sig) =
host_sign(auth::challenge_payload(challenge_id, server_challenge)).await?;
match host_sign(auth::challenge_payload(challenge_id, server_challenge)).await {
Ok(signatures) => signatures,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let mut challenge_msg = CommunicationValue::new(CommunicationType::Challenge)
.add_typed_default(
DataType::ServerNonce,
DataValue::UnsignedNumber(server_challenge),
)
.add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig));
let mut challenge_msg =
CommunicationValue::new_with_type_map(CommunicationType::Challenge, tm)
.add_typed_default(
DataType::ServerNonce,
DataValue::UnsignedNumber(server_challenge),
)
.add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig));
challenge_msg = challenge_msg.add_typed_default(
DataType::RequirePq,
if self.config.require_pq {
@ -536,31 +711,28 @@ impl HandshakeEngine {
AcceptError::Receive(e)
})?;
if Some(proof.get_type()) != CommunicationType::ChallengeResponse.try_to_id(&tm) {
sender.close();
return Err(AcceptError::AuthenticationFailed(
"missing challenge response".into(),
));
let error = AcceptError::AuthenticationFailed("missing challenge response".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
let client_nonce = match proof.get_data(DataType::ClientNonce) {
DataValue::UnsignedNumber(n) => *n,
Some(DataValue::UnsignedNumber(n)) => *n,
_ => {
sender.close();
return Err(AcceptError::AuthenticationFailed(
"missing client nonce".into(),
));
let error = AcceptError::AuthenticationFailed("missing client nonce".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let sig_bytes = match proof.get_data(DataType::Signature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => {
sender.close();
return Err(AcceptError::AuthenticationFailed(
"missing challenge signature".into(),
));
let error = AcceptError::AuthenticationFailed("missing challenge signature".into());
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let pq_sig_bytes: Vec<u8> = match proof.get_data(DataType::PqSignature) {
DataValue::Bytes(b) => b.clone(),
Some(DataValue::Bytes(b)) => b.clone(),
_ => vec![],
};
@ -601,6 +773,7 @@ impl HandshakeEngine {
RejectionReason::AuthenticationFailed {
detail: "client proof signature invalid".into(),
},
Some(tm),
)
.await;
sender.close();
@ -613,21 +786,66 @@ impl HandshakeEngine {
let (assigned_id, client_bundle) = match flow {
Flow::Login { id, bundle } => (id, bundle),
Flow::Register { bundle, .. } => {
let new_id =
(self.config.complete_register)(bundle.clone(), description.clone()).await;
let _registration_guard = self.config.registration_lock.lock().await;
let identity = bundle.as_bytes();
let cached_id = self
.config
.registration_ids
.lock()
.ok()
.and_then(|registrations| registrations.get(&identity).copied());
let new_id = if let Some(id) = cached_id {
id
} else if let Some(lookup) = &self.config.find_registered_client {
match lookup(bundle.clone(), description.clone()).await {
Some(id) => id,
None => {
(self.config.complete_register)(bundle.clone(), description.clone())
.await
}
}
} else {
(self.config.complete_register)(bundle.clone(), description.clone()).await
};
if new_id != 0
&& let Ok(mut registrations) = self.config.registration_ids.lock()
{
registrations.insert(identity, new_id);
}
if new_id == 0 {
let error = AcceptError::AuthenticationFailed(
"registration callback returned reserved client id 0".into(),
);
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
(new_id, bundle)
}
};
if assigned_id == 0 {
let error = AcceptError::AuthenticationFailed(
"client id 0 is reserved for no authenticated identity".into(),
);
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
// Sign and send final response
let (host_sig, host_pq_sig) = host_sign(auth::host_final_payload(
let (host_sig, host_pq_sig) = match host_sign(auth::host_final_payload(
assigned_id,
client_nonce,
server_challenge,
))
.await?;
.await
{
Ok(signatures) => signatures,
Err(error) => {
reject_error_generic(sender, &error, tm).await;
return Err(error);
}
};
let mut response = CommunicationValue::new(response_type)
let mut response = CommunicationValue::new_with_type_map(response_type, tm)
.add_typed_default(DataType::Connected, DataValue::BoolTrue)
.add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128))
.add_typed_default(
@ -658,6 +876,7 @@ impl HandshakeEngine {
auth_state: crate::error::AuthState::Authenticated,
client_id: assigned_id,
client_public_key: Some(client_bundle),
guest_id_lease: None,
})
}
}
@ -682,36 +901,52 @@ enum Flow {
impl HandshakeEngine {
const GUEST_ID_MAX_RETRIES: u32 = 100;
async fn assign_guest_id(&self) -> Result<u64, AcceptError> {
async fn assign_guest_id(&self) -> Result<GuestIdLease, AcceptError> {
if let Some(ref generator) = self.config.guest_id_generator {
let id = generator().await.ok_or_else(|| {
AcceptError::AuthenticationFailed(
"guest id generator rejected the connection".into(),
)
})?;
if id > mtp_codec::MAX_WIRE_ID {
if id == 0 {
return Err(AcceptError::AuthenticationFailed(
"guest id exceeds wire limit".into(),
"guest id 0 is reserved for no authenticated identity".into(),
));
}
if (self.config.get_existing_client)(id, None).await.is_none() {
return Ok(id);
if let Some(lease) = self.try_reserve_guest_id(id).await? {
return Ok(lease);
}
}
self.random_guest_id().await
}
async fn random_guest_id(&self) -> Result<u64, AcceptError> {
async fn random_guest_id(&self) -> Result<GuestIdLease, AcceptError> {
for _ in 0..Self::GUEST_ID_MAX_RETRIES {
let id = rand::random::<u64>() & mtp_codec::MAX_WIRE_ID;
if (self.config.get_existing_client)(id, None).await.is_none() {
return Ok(id);
let id = rand::random::<u64>();
if let Some(lease) = self.try_reserve_guest_id(id).await? {
return Ok(lease);
}
}
Err(AcceptError::AuthenticationFailed(
"failed to allocate a unique guest id after retries".into(),
))
}
async fn try_reserve_guest_id(&self, id: u64) -> Result<Option<GuestIdLease>, AcceptError> {
if id == 0 || (self.config.get_existing_client)(id, None).await.is_some() {
return Ok(None);
}
let mut active_ids = self.config.active_guest_ids.lock().map_err(|_| {
AcceptError::AuthenticationFailed("guest ID registry is poisoned".into())
})?;
if !active_ids.insert(id) {
return Ok(None);
}
Ok(Some(GuestIdLease {
active_ids: Arc::clone(&self.config.active_guest_ids),
id,
}))
}
}
// ---------------------------------------------------------------------------
@ -723,7 +958,7 @@ fn extract_register_bundle(
msg: &CommunicationValue,
) -> Result<mtp_crypto::PublicKeyBundle, AcceptError> {
match msg.get_data(DataType::PublicKeys) {
DataValue::Bytes(b) => mtp_crypto::PublicKeyBundle::from_bytes(b)
Some(DataValue::Bytes(b)) => mtp_crypto::PublicKeyBundle::from_bytes(b)
.map_err(|_| AcceptError::AuthenticationFailed("invalid public key bundle".into())),
_ => Err(AcceptError::AuthenticationFailed(
"missing public keys".into(),
@ -731,32 +966,58 @@ fn extract_register_bundle(
}
}
async fn send_rejection_generic<S: HandshakeSender>(sender: &S, reason: RejectionReason) {
async fn send_rejection_generic<S: HandshakeSender>(
sender: &S,
reason: RejectionReason,
type_map: Option<&TypeMap>,
) {
let type_map = type_map.cloned().unwrap_or_else(TypeMap::latest);
let response = match &reason {
RejectionReason::BadVersion { supported_versions } => {
CommunicationValue::new(CommunicationType::ErrorBadVersion)
CommunicationValue::new_with_type_map(CommunicationType::ErrorBadVersion, &type_map)
.add_typed_default(
DataType::Version,
DataValue::Str(supported_versions.join(",")),
)
.add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string()))
}
_ => CommunicationValue::new(CommunicationType::IdentificationResponse)
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
.add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())),
_ => CommunicationValue::new_with_type_map(
CommunicationType::IdentificationResponse,
&type_map,
)
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
.add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())),
};
let _ = sender.send(&response).await;
}
#[cfg(feature = "crypto")]
async fn reject_error_generic<S: HandshakeSender>(
sender: &S,
error: &AcceptError,
type_map: &TypeMap,
) {
send_rejection_generic(
sender,
RejectionReason::AuthenticationFailed {
detail: error.to_string(),
},
Some(type_map),
)
.await;
sender.close();
}
async fn send_accepted_generic<S: HandshakeSender>(
sender: &S,
version: &Version,
type_map: &TypeMap,
assigned_id: Option<u64>,
) -> Result<(), CommunicationError> {
let mut response = CommunicationValue::new(CommunicationType::IdentificationResponse)
.add_typed_default(DataType::Connected, DataValue::BoolTrue)
.add_typed_default(DataType::Version, DataValue::Str(version.to_string()));
let mut response =
CommunicationValue::new_with_type_map(CommunicationType::IdentificationResponse, type_map)
.add_typed_default(DataType::Connected, DataValue::BoolTrue)
.add_typed_default(DataType::Version, DataValue::Str(version.to_string()));
if let Some(id) = assigned_id {
response = response.add_typed_default(DataType::Id, DataValue::UnsignedNumber(id as u128));
}
@ -780,6 +1041,9 @@ impl HandshakeSender for mtp_transport::Sender {
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
mtp_transport::Sender::finish_stream(self)
}
fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async move { self.set_type_map(type_map).await }
}
fn close(&self) {
let sender = self.clone();
tokio::spawn(async move { sender.close().await });
@ -793,6 +1057,10 @@ impl HandshakeReceiver for mtp_transport::Receiver {
{
mtp_transport::Receiver::receive(self)
}
fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async move { self.set_type_map(type_map).await }
}
}
impl<C: mtp_transport::TransportConnection> HandshakeSender for mtp_transport::GenericSender<C> {
@ -807,6 +1075,9 @@ impl<C: mtp_transport::TransportConnection> HandshakeSender for mtp_transport::G
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
mtp_transport::GenericSender::finish_stream(self)
}
fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async move { self.set_type_map(type_map).await }
}
fn close(&self) {
mtp_transport::GenericSender::close(self);
}
@ -821,4 +1092,68 @@ impl<C: mtp_transport::TransportConnection> HandshakeReceiver
{
mtp_transport::GenericReceiver::receive(self)
}
fn set_type_map(&self, type_map: &TypeMap) -> impl std::future::Future<Output = ()> + Send {
async move { self.set_type_map(type_map).await }
}
}
#[cfg(all(test, feature = "crypto"))]
mod tests {
use super::*;
use std::sync::Mutex;
#[tokio::test]
async fn guest_generator_accepts_full_width_id_after_collision_check()
-> Result<(), Box<dyn std::error::Error>> {
let lookups = Arc::new(Mutex::new(Vec::new()));
let recorded_lookups = Arc::clone(&lookups);
let mut config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new())
.with_guest_id_generator(Box::new(|| Box::pin(async { Some(u64::MAX) })));
config.get_existing_client = Box::new(move |id, description| {
let recorded_lookups = Arc::clone(&recorded_lookups);
Box::pin(async move {
recorded_lookups
.lock()
.expect("guest ID lookup mutex should not be poisoned")
.push((id, description));
None
})
});
let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));
assert_eq!(engine.assign_guest_id().await?.id, u64::MAX);
assert_eq!(
*lookups
.lock()
.expect("guest ID lookup mutex should not be poisoned"),
vec![(u64::MAX, None)]
);
Ok(())
}
#[tokio::test]
async fn active_guest_ids_are_unique_and_released() -> Result<(), Box<dyn std::error::Error>> {
let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new());
let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));
let first = engine.try_reserve_guest_id(1).await?.unwrap();
assert!(engine.try_reserve_guest_id(1).await?.is_none());
drop(first);
assert!(engine.try_reserve_guest_id(1).await?.is_some());
Ok(())
}
#[tokio::test]
async fn zero_is_rejected_as_a_guest_id() -> Result<(), Box<dyn std::error::Error>> {
let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new())
.with_guest_id_generator(Box::new(|| Box::pin(async { Some(0) })));
let engine = HandshakeEngine::new(Registry::builtin(), Arc::new(config));
assert!(engine.assign_guest_id().await.is_err());
Ok(())
}
}

7
host/src/error.rs
View file

@ -7,14 +7,13 @@ use mtp_codec::{CommunicationValue, DataType, DataValue};
#[cfg(feature = "crypto")]
pub(crate) fn random_client_id() -> u64 {
rand::random::<u64>() & mtp_codec::MAX_WIRE_ID
rand::random::<u64>()
}
#[cfg(test)]
pub(crate) fn extract_version(msg: &CommunicationValue) -> Option<Version> {
let value = msg.get_data(DataType::Version);
match value {
DataValue::Str(s) => Version::parse(s.as_str()),
match msg.get_data(DataType::Version) {
Some(DataValue::Str(s)) => Version::parse(s.as_str()),
_ => None,
}
}

12
host/src/handshake.rs
View file

@ -164,6 +164,8 @@ impl HandshakeContext {
let remote_addr = sender.handle().remote_addr();
receiver.set_max_message_size(self.config.policy.max_message_size);
#[cfg(feature = "pipes")]
let type_map = result.codec.type_map().clone();
#[cfg(feature = "pipes")]
{
if self.config.send_pongs {
receiver.respond_to_pings(sender.clone());
@ -177,6 +179,7 @@ impl HandshakeContext {
pending_creations: tokio::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(),
});
let dispatcher_clone = dispatcher.clone();
@ -202,9 +205,11 @@ impl HandshakeContext {
pipe_dispatcher: dispatcher,
description: result.description,
_dispatcher_task: task,
_connection_guard: None,
auth_state: result.auth_state,
client_id: result.client_id,
client_public_key: result.client_public_key,
guest_id_lease: result.guest_id_lease,
}
}
@ -226,9 +231,11 @@ impl HandshakeContext {
_pipe_stream: std::marker::PhantomData,
description: result.description,
_dispatcher_task: task,
_connection_guard: None,
auth_state: result.auth_state,
client_id: result.client_id,
client_public_key: result.client_public_key,
guest_id_lease: result.guest_id_lease,
}
}
}
@ -246,6 +253,8 @@ impl HandshakeContext {
let remote_addr = sender.handle().remote_addr();
receiver.set_max_message_size(self.config.policy.max_message_size);
#[cfg(feature = "pipes")]
let type_map = codec.type_map().clone();
#[cfg(feature = "pipes")]
{
if self.config.send_pongs {
receiver.respond_to_pings(sender.clone());
@ -259,6 +268,7 @@ impl HandshakeContext {
pending_creations: tokio::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,
});
let dispatcher_clone = dispatcher.clone();
@ -284,6 +294,7 @@ impl HandshakeContext {
pipe_dispatcher: dispatcher,
description,
_dispatcher_task: task,
_connection_guard: None,
}
}
@ -305,6 +316,7 @@ impl HandshakeContext {
_pipe_stream: std::marker::PhantomData,
description,
_dispatcher_task: task,
_connection_guard: None,
}
}
}

11
host/src/lib.rs
View file

@ -28,7 +28,10 @@ pub use pipe::PipeRequest;
pub use mtp_codec::registry::Registry;
#[cfg(feature = "crypto")]
pub use config::{AuthenticationPolicy, CompleteRegister, GetExistingClient, GuestIdGenerator};
pub use config::{
AuthenticationPolicy, CompleteRegister, FindRegisteredClient, GetExistingClient,
GuestIdGenerator,
};
#[cfg(feature = "crypto")]
pub use error::AuthState;
@ -51,9 +54,9 @@ mod tests {
fn version_extraction() {
let tm = mtp_codec::TypeMap::latest();
let msg = mtp_codec::CommunicationValue::from_comm(CommunicationType::Identification, &tm)
.add_typed(DataType::Version, &tm, DataValue::Str("2.0".to_string()));
.add_typed(DataType::Version, &tm, DataValue::Str("3.0".to_string()));
let version = error::extract_version(&msg);
assert_eq!(version, Some(mtp_codec::Version(2, 0)));
assert_eq!(version, Some(mtp_codec::Version(3, 0)));
}
#[test]
@ -92,7 +95,7 @@ mod tests {
#[tokio::test]
async fn alternative_transports_use_the_shared_connection_type() {
let registry = Registry::builtin();
let version = mtp_codec::Version(1, 0);
let version = mtp_codec::Version(3, 0);
let codec = VersionedCodec::for_version(registry, version.clone()).unwrap();
let connection: MTPConnection<AlternateSender, AlternateReceiver> =
MTPConnection::from_transport_parts(

98
host/src/pipe.rs
View file

@ -1,4 +1,4 @@
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, TypeMap};
use mtp_common::{CommunicationError, PipeError};
use mtp_transport::{PipeReader, PipeWriter, Policy, TransportEvent};
use std::collections::HashMap;
@ -152,24 +152,49 @@ where
.await
.insert(self.pipe_id, pipe_tx);
let response = CommunicationValue::new(CommunicationType::PipeResponse)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolTrue);
self.sender
.send_pipe_message(&response)
.await
.map_err(PipeError::from)?;
let response = 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_pipe_message(&response).await {
self.dispatcher
.pending_pipes
.lock()
.await
.remove(&self.pipe_id);
return Err(PipeError::from(error));
}
tokio::time::timeout(self.dispatcher.policy.read_timeout, pipe_rx)
.await
.map_err(|_| PipeError::HandshakeTimeout)?
.map_err(|_| PipeError::StreamClosed)
match tokio::time::timeout(self.dispatcher.policy.read_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 response = CommunicationValue::new(CommunicationType::PipeResponse)
.with_id(self.pipe_id)
.add_typed_default(DataType::Accepted, DataValue::BoolFalse);
let response = 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_pipe_message(&response)
.await
@ -182,6 +207,7 @@ pub(crate) struct PipeDispatcher<P> {
Mutex<HashMap<u32, tokio::sync::oneshot::Sender<Result<bool, PipeError>>>>,
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) async fn run_dispatcher<S, R, P>(
@ -195,15 +221,21 @@ pub(crate) async fn run_dispatcher<S, R, P>(
R: PipeReceiver<P>,
P: tokio::io::AsyncRead + Send + Unpin + 'static,
{
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_pipe_event().await {
Ok(TransportEvent::Message(message)) => {
if Some(message.get_type()) == pipe_req_type {
if message.is_type(CommunicationType::PipeRequest) {
let Some(pipe_id) = message.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 request = PipeRequest {
pipe_id: message.get_id(),
pipe_id,
description: message
.get_str(DataType::Description)
.unwrap_or("")
@ -214,14 +246,36 @@ pub(crate) async fn run_dispatcher<S, R, P>(
let _ = pipe_req_tx.send(request).await;
continue;
}
if Some(message.get_type()) == pipe_resp_type {
if message.is_type(CommunicationType::PipeResponse) {
let Some(pipe_id) = message.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 mut pending = dispatcher.pending_creations.lock().await;
if let Some(reply) = pending.remove(&message.get_id()) {
if let Some(reply) = pending.remove(&pipe_id) {
let _ =
reply.send(Ok(message.get_bool(DataType::Accepted).unwrap_or(false)));
}
continue;
}
if !matches!(message.id(), Some(id) if id != 0)
&& message
.get_type_name()
.is_some_and(|name| name.ends_with("Response"))
{
let error = CommunicationError::Other(
"response frame must contain a non-zero id".into(),
);
if app_tx.send(Err(error)).await.is_err() {
break;
}
continue;
}
if app_tx.send(Ok(message)).await.is_err() {
break;
}