406 lines
14 KiB
Rust
406 lines
14 KiB
Rust
#[cfg(feature = "crypto")]
|
|
use mtp_codec::DataType;
|
|
use mtp_codec::{
|
|
CommunicationValue, DataTypeId, DataValue, Version,
|
|
registry::{Registry, VersionedCodec},
|
|
};
|
|
use mtp_common::CommunicationError;
|
|
use mtp_transport::{Policy, Receiver, Sender};
|
|
use std::net::IpAddr;
|
|
|
|
// Host configuration.
|
|
pub struct HostConfig {
|
|
pub ip: IpAddr,
|
|
pub port: u16,
|
|
pub tls_fullchain: Vec<u8>,
|
|
pub tls_key: Vec<u8>,
|
|
|
|
#[cfg(feature = "crypto")]
|
|
pub require_authentication: bool,
|
|
#[cfg(feature = "crypto")]
|
|
pub host_id: u64,
|
|
#[cfg(feature = "crypto")]
|
|
pub host_keyring: mtp_crypto::Keyring,
|
|
#[cfg(feature = "crypto")]
|
|
pub get_existing_user: Box<dyn Fn(u64) -> Option<mtp_crypto::PublicKeyBundle> + Send>,
|
|
#[cfg(feature = "crypto")]
|
|
pub complete_register: Box<dyn Fn(mtp_crypto::PublicKeyBundle) -> u64 + Send>,
|
|
}
|
|
|
|
#[cfg(feature = "crypto")]
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum AuthState {
|
|
Pending,
|
|
Authenticated,
|
|
Failed,
|
|
}
|
|
|
|
// A connection that has completed version negotiation.
|
|
pub struct MTPConnection {
|
|
pub version: Version,
|
|
pub codec: VersionedCodec,
|
|
pub sender: Sender,
|
|
pub receiver: Receiver,
|
|
#[cfg(feature = "crypto")]
|
|
pub auth_state: AuthState,
|
|
#[cfg(feature = "crypto")]
|
|
pub client_id: u64,
|
|
#[cfg(feature = "crypto")]
|
|
pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
|
|
}
|
|
|
|
// High-level MTP host with built-in version negotiation.
|
|
pub struct MTPHost {
|
|
transport: mtp_transport::Host,
|
|
registry: Registry,
|
|
#[cfg(feature = "crypto")]
|
|
config: HostConfig,
|
|
}
|
|
|
|
impl MTPHost {
|
|
pub async fn new(config: HostConfig) -> Result<Self, CommunicationError> {
|
|
let registry = Registry::builtin();
|
|
|
|
let transport = mtp_transport::host(
|
|
config.ip,
|
|
config.port,
|
|
config.tls_fullchain.clone(),
|
|
config.tls_key.clone(),
|
|
Policy::default(),
|
|
)
|
|
.await?;
|
|
|
|
Ok(Self {
|
|
transport,
|
|
registry,
|
|
#[cfg(feature = "crypto")]
|
|
config,
|
|
})
|
|
}
|
|
|
|
/*
|
|
* Accept an incoming connection, negotiate the protocol version,
|
|
* and return a ready-to-use `MTPConnection`.
|
|
*
|
|
* Returns `None` if the connection is closed or the version is
|
|
* incompatible.
|
|
*/
|
|
pub async fn accept(&mut self) -> Option<MTPConnection> {
|
|
let (sender, receiver) = self.transport.next().await?;
|
|
|
|
#[cfg(feature = "crypto")]
|
|
if self.config.require_authentication {
|
|
return self.accept_authenticated(sender, receiver).await;
|
|
}
|
|
|
|
// Read the first message (always encoded with reserved types).
|
|
let first_msg = receiver.receive().await.ok()?;
|
|
|
|
/*
|
|
* Extract the client's version from the first message.
|
|
* The client is expected to send DataType::Version (reserved ID 3)
|
|
* as a DataValue::Str("X.Y").
|
|
*
|
|
* Then negotiate the version for single-version clients
|
|
*/
|
|
let client_version = extract_version(&first_msg)?;
|
|
|
|
let negotiated = self.registry.negotiate(&[client_version])?;
|
|
|
|
let codec = VersionedCodec::new(self.registry.clone());
|
|
|
|
Some(MTPConnection {
|
|
version: negotiated,
|
|
codec,
|
|
sender,
|
|
receiver,
|
|
#[cfg(feature = "crypto")]
|
|
auth_state: AuthState::Authenticated,
|
|
#[cfg(feature = "crypto")]
|
|
client_id: 0,
|
|
#[cfg(feature = "crypto")]
|
|
client_public_key: None,
|
|
})
|
|
}
|
|
|
|
pub fn local_addr(&self) -> std::net::SocketAddr {
|
|
self.transport.local_addr()
|
|
}
|
|
|
|
pub fn registry(&self) -> &Registry {
|
|
&self.registry
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "crypto")]
|
|
impl MTPHost {
|
|
async fn accept_authenticated(
|
|
&mut self,
|
|
sender: Sender,
|
|
receiver: Receiver,
|
|
) -> Option<MTPConnection> {
|
|
use mtp_crypto::{
|
|
Ed25519Signer, PublicKeyBundle, SignatureScheme, verify_ed25519, verify_ml_dsa,
|
|
};
|
|
|
|
// 1. Receive client message first (no host greeting)
|
|
let msg = receiver.receive().await.ok()?;
|
|
let version_str = match msg.get_data(DataTypeId(3)) {
|
|
DataValue::Str(s) => s.clone(),
|
|
_ => {
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let client_version = Version::parse(&version_str)?;
|
|
|
|
let client_nonce = match msg.get_data(DataTypeId(7)) {
|
|
DataValue::UnsignedNumber(n) => *n,
|
|
_ => {
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let sig_bytes = match msg.get_data(DataTypeId(10)) {
|
|
DataValue::Bytes(b) => b.clone(),
|
|
_ => {
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let pq_sig_bytes: Vec<u8> = match msg.get_data(DataTypeId(12)) {
|
|
DataValue::Bytes(b) => b.clone(),
|
|
_ => vec![],
|
|
};
|
|
|
|
let (assigned_id, client_bundle) = if msg.get_type() == mtp_codec::CommunicationTypeId(15) {
|
|
// LOGIN
|
|
let cid = match msg.get_data(DataTypeId(6)) {
|
|
DataValue::UnsignedNumber(n) => *n as u64,
|
|
_ => {
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let bundle = match (self.config.get_existing_user)(cid) {
|
|
Some(b) => b,
|
|
None => {
|
|
let rejection = CommunicationValue::new(
|
|
mtp_codec::CommunicationType::IdentificationResponse,
|
|
)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
);
|
|
let _ = sender.send(&rejection).await;
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let mut sig_payload = Vec::new();
|
|
sig_payload.extend_from_slice(version_str.as_bytes());
|
|
sig_payload.extend_from_slice(&cid.to_be_bytes());
|
|
sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
|
|
|
|
/* ===== Signature ===== */
|
|
if verify_ed25519(&bundle.sig_cl_public_key, &sig_payload, &sig_bytes).is_err() {
|
|
let rejection =
|
|
CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
);
|
|
let _ = sender.send(&rejection).await;
|
|
sender.close();
|
|
return None;
|
|
}
|
|
if !pq_sig_bytes.is_empty()
|
|
&& verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err()
|
|
{
|
|
let rejection =
|
|
CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
);
|
|
let _ = sender.send(&rejection).await;
|
|
sender.close();
|
|
return None;
|
|
}
|
|
/* ===== End Signature ===== */
|
|
|
|
(cid, bundle)
|
|
} else if msg.get_type() == mtp_codec::CommunicationTypeId(17) {
|
|
// REGISTER
|
|
let bundle = match msg.get_data(DataTypeId(9)) {
|
|
DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).ok()?,
|
|
_ => {
|
|
sender.close();
|
|
return None;
|
|
}
|
|
};
|
|
let pk_bytes = bundle.as_bytes();
|
|
|
|
let mut sig_payload = Vec::new();
|
|
sig_payload.extend_from_slice(version_str.as_bytes());
|
|
sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
|
|
sig_payload.extend_from_slice(&pk_bytes);
|
|
|
|
/* ===== Signature ===== */
|
|
if verify_ed25519(&bundle.sig_cl_public_key, &sig_payload, &sig_bytes).is_err() {
|
|
let rejection =
|
|
CommunicationValue::new(mtp_codec::CommunicationType::RegisterResponse)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
);
|
|
let _ = sender.send(&rejection).await;
|
|
sender.close();
|
|
return None;
|
|
}
|
|
if !pq_sig_bytes.is_empty()
|
|
&& verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err()
|
|
{
|
|
let rejection =
|
|
CommunicationValue::new(mtp_codec::CommunicationType::RegisterResponse)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolFalse)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
);
|
|
let _ = sender.send(&rejection).await;
|
|
sender.close();
|
|
return None;
|
|
}
|
|
/* ===== End Signature ===== */
|
|
|
|
let new_id = (self.config.complete_register)(bundle.clone());
|
|
(new_id, bundle)
|
|
} else {
|
|
sender.close();
|
|
return None;
|
|
};
|
|
|
|
// 2. Send success response (single host message)
|
|
let new_nonce: u128 = rand::random();
|
|
|
|
let mut host_sig_payload = Vec::new();
|
|
host_sig_payload.push(0x01);
|
|
host_sig_payload.extend_from_slice(&assigned_id.to_be_bytes());
|
|
host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes());
|
|
host_sig_payload.extend_from_slice(&new_nonce.to_be_bytes());
|
|
|
|
let host_signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key).ok()?;
|
|
|
|
/* ===== Signature ===== */
|
|
let host_sig = host_signer.sign(&host_sig_payload).ok()?;
|
|
|
|
let mut response =
|
|
CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse)
|
|
.add_typed_default(DataType::Connected, DataValue::BoolTrue)
|
|
.add_typed_default(
|
|
DataType::ClientNonce,
|
|
DataValue::UnsignedNumber(client_nonce),
|
|
)
|
|
.add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(new_nonce))
|
|
.add_typed_default(DataType::Signature, DataValue::Bytes(host_sig))
|
|
.add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128));
|
|
|
|
if !self
|
|
.config
|
|
.host_keyring
|
|
.sig_pq_secret_key
|
|
.as_bytes()
|
|
.is_empty()
|
|
{
|
|
use mtp_crypto::MlDsaSigner;
|
|
let host_pq_signer = MlDsaSigner::new(
|
|
&self.config.host_keyring.sig_pq_secret_key,
|
|
&self.config.host_keyring.sig_pq_public_key,
|
|
)
|
|
.ok()?;
|
|
let host_pq_sig = host_pq_signer.sign(&host_sig_payload).ok()?;
|
|
response =
|
|
response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig));
|
|
}
|
|
/* ===== End Signature ===== */
|
|
|
|
sender.send(&response).await.ok()?;
|
|
|
|
// 3. Version negotiation
|
|
let negotiated = self.registry.negotiate(&[client_version])?;
|
|
let codec = VersionedCodec::new(self.registry.clone());
|
|
|
|
Some(MTPConnection {
|
|
version: negotiated,
|
|
codec,
|
|
sender,
|
|
receiver,
|
|
auth_state: AuthState::Authenticated,
|
|
client_id: assigned_id,
|
|
client_public_key: Some(client_bundle),
|
|
})
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Extract the protocol version from an initial `CommunicationValue`.
|
|
*
|
|
* The client's first message must contain a `Version` data entry
|
|
* (reserved ID 3) mapping to `DataValue::Str("major.minor")`.
|
|
*/
|
|
fn extract_version(msg: &CommunicationValue) -> Option<Version> {
|
|
let value = msg.get_data(DataTypeId(3));
|
|
match value {
|
|
DataValue::Str(s) => Version::parse(s.as_str()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
/* ================================ TESTS ================================ */
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn version_extraction() {
|
|
let tm = mtp_codec::TypeMap::latest();
|
|
let msg = mtp_codec::CommunicationValue::from_comm(
|
|
mtp_codec::CommunicationType::Identification,
|
|
&tm,
|
|
)
|
|
.add_data(DataTypeId(3), DataValue::Str("2.0".to_string()));
|
|
let version = extract_version(&msg);
|
|
assert_eq!(version, Some(Version(2, 0)));
|
|
}
|
|
|
|
#[test]
|
|
fn version_extraction_returns_none_for_missing() {
|
|
let tm = mtp_codec::TypeMap::latest();
|
|
let msg = mtp_codec::CommunicationValue::from_comm(
|
|
mtp_codec::CommunicationType::Identification,
|
|
&tm,
|
|
);
|
|
assert!(extract_version(&msg).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn version_extraction_bad_format() {
|
|
let tm = mtp_codec::TypeMap::latest();
|
|
let msg = mtp_codec::CommunicationValue::from_comm(
|
|
mtp_codec::CommunicationType::Identification,
|
|
&tm,
|
|
)
|
|
.add_data(DataTypeId(3), DataValue::UnsignedNumber(42));
|
|
assert!(extract_version(&msg).is_none());
|
|
}
|
|
}
|