297 lines
9.7 KiB
Rust
297 lines
9.7 KiB
Rust
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use base64::Engine as _;
|
|
use base64::engine::general_purpose::STANDARD as BASE64_STD;
|
|
|
|
use crate::{
|
|
app_state::AppState,
|
|
log, log_err,
|
|
omega::omega_connection::OmegaConnection,
|
|
rho::connection::{ConnectionKind, GeneralConnection, OptionalDataValueCompat},
|
|
util::{file_util::load_file_vec, logger::PrintType},
|
|
};
|
|
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
|
|
use mtp::crypto::PublicKeyBundle;
|
|
use mtp::host::{AuthenticationPolicy, HostConfig, Policy, SendMode};
|
|
use mtp::webserver::{MTPWebServer, WebServerConfig};
|
|
|
|
fn web_config(max_connections: usize) -> Result<WebServerConfig, mtp::webserver::RouterError> {
|
|
WebServerConfig::new()
|
|
.max_connections(max_connections)
|
|
.route("/", |_request, response| async move { response.body("OK") })
|
|
}
|
|
|
|
fn rho_policy() -> Policy {
|
|
Policy::default()
|
|
.with_send_mode(SendMode::SingleStreamPerMessage)
|
|
.with_timeouts(
|
|
Duration::from_millis(2_000),
|
|
Duration::from_millis(2_000),
|
|
Duration::from_millis(30_000),
|
|
)
|
|
.with_keep_alive(Some(Duration::from_secs(6)))
|
|
.with_max_idle_timeout(Some(Duration::from_secs(30)))
|
|
.with_receiver_queue_capacity(1000)
|
|
.with_max_concurrent_stream_tasks(10)
|
|
.with_persistent_stream_retries(5, Duration::from_secs(5))
|
|
}
|
|
|
|
/*
|
|
* Resolves the PublicKeyBundle mtp needs to verify a login's signed
|
|
* challenge response. "iota"/"client" ids are looked up through Omega, the
|
|
* source of truth for both kinds of registered keys. Any other description
|
|
* (including "anonymous") never resolves, which routes the connection
|
|
* through mtp's unauthenticated fallback instead of rejecting it outright.
|
|
*/
|
|
pub async fn get_by_connector_id(
|
|
omega: Arc<OmegaConnection>,
|
|
client_id: u64,
|
|
description: Option<String>,
|
|
) -> Option<PublicKeyBundle> {
|
|
let request = match description.as_deref() {
|
|
Some("iota") => {
|
|
println!("Iota connection request for client_id: {}", client_id);
|
|
CommunicationValue::new(CommunicationType::GetIotaData)
|
|
.add_typed_default(DataType::IotaId, DataValue::SignedNumber(client_id as i128))
|
|
}
|
|
Some("client") => CommunicationValue::new(CommunicationType::GetUserData)
|
|
.add_typed_default(DataType::UserId, DataValue::SignedNumber(client_id as i128)),
|
|
_ => return None,
|
|
};
|
|
|
|
let response = match omega
|
|
.await_response(&request, Some(Duration::from_secs(20)))
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(e) => {
|
|
log_err!(
|
|
client_id as i64,
|
|
PrintType::General,
|
|
"Failed to look up public key for connector (description={:?}): {}",
|
|
description,
|
|
e
|
|
);
|
|
return None;
|
|
}
|
|
};
|
|
|
|
let bytes = BASE64_STD
|
|
.decode(response.get_data(DataType::PublicKey).as_str()?)
|
|
.ok()?;
|
|
PublicKeyBundle::from_bytes(&bytes).ok()
|
|
}
|
|
|
|
/* Only Iota registration goes through mtp's Register flow; users are registered out of band. */
|
|
pub async fn complete_register(
|
|
omega: Arc<OmegaConnection>,
|
|
pub_key: PublicKeyBundle,
|
|
description: Option<String>,
|
|
) -> u64 {
|
|
println!("Iota register start");
|
|
if description.as_deref() != Some("iota") {
|
|
return 0;
|
|
}
|
|
println!("Iota connection request");
|
|
|
|
let pub_key_bytes = match pub_key.try_as_bytes() {
|
|
Ok(bytes) => bytes,
|
|
Err(error) => {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Failed to serialize Iota public key: {}",
|
|
error
|
|
);
|
|
return 0;
|
|
}
|
|
};
|
|
|
|
let request = CommunicationValue::new(CommunicationType::CompleteRegisterIota)
|
|
.add_typed_default(
|
|
DataType::PublicKey,
|
|
DataValue::Str(BASE64_STD.encode(pub_key_bytes)),
|
|
);
|
|
|
|
let response = match omega
|
|
.await_response(&request, Some(Duration::from_secs(20)))
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(e) => {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Failed to complete Iota registration: {}",
|
|
e
|
|
);
|
|
return 0;
|
|
}
|
|
};
|
|
|
|
match response.get_data(DataType::IotaId) {
|
|
Some(DataValue::SignedNumber(id)) => *id as u64,
|
|
_ => 0,
|
|
}
|
|
}
|
|
|
|
pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error>> {
|
|
let cert_pem = load_file_vec("certs", "cert.pem").expect("Error loading Pemfile");
|
|
|
|
let key_pem = load_file_vec("certs", "key.pem").expect("Error loading Keyfile");
|
|
|
|
let host_config = HostConfig::new(
|
|
state.config.bind_address,
|
|
state.config.rho_port,
|
|
cert_pem,
|
|
key_pem,
|
|
)
|
|
.with_policy(rho_policy())
|
|
.with_authentication(
|
|
state
|
|
.keyring_for_host()
|
|
.map_err(|error| format!("Unable to copy keyring for host: {error}"))?,
|
|
Box::new({
|
|
let omega = state.omega.clone();
|
|
move |user_id, description| {
|
|
Box::pin(get_by_connector_id(omega.clone(), user_id, description))
|
|
}
|
|
}),
|
|
Box::new({
|
|
let omega = state.omega.clone();
|
|
move |pub_key, description| {
|
|
Box::pin(complete_register(omega.clone(), pub_key, description))
|
|
}
|
|
}),
|
|
)
|
|
.with_authentication_policy(AuthenticationPolicy::AllowAuthentication);
|
|
|
|
let web_config = web_config(state.config.rho_max_connections)?;
|
|
let mut host = MTPWebServer::new(host_config, web_config).await?;
|
|
log!(
|
|
0,
|
|
PrintType::General,
|
|
"Server listening on {}:{}.",
|
|
state.config.bind_address,
|
|
state.config.rho_port
|
|
);
|
|
|
|
loop {
|
|
let conn = match host.accept().await {
|
|
Ok(Some(conn)) => {
|
|
println!("Accepted connection");
|
|
conn
|
|
}
|
|
Ok(None) => break,
|
|
Err(e) => {
|
|
/*
|
|
* A single client's failed/aborted handshake (bad auth, a
|
|
* probe, a mid-handshake disconnect) must not take down the
|
|
* whole listener, only that connection attempt is lost.
|
|
*/
|
|
log_err!(0, PrintType::General, "Rejected connection: {}", e);
|
|
continue;
|
|
}
|
|
};
|
|
|
|
let global_permit = match state.rho_connection_limits.all.clone().try_acquire_owned() {
|
|
Ok(permit) => permit,
|
|
Err(_) => {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Rejected connection: global Rho connection limit reached"
|
|
);
|
|
continue;
|
|
}
|
|
};
|
|
let peer_ip = conn.remote_addr.map(|address| address.ip());
|
|
let state = state.clone();
|
|
tokio::spawn(async move {
|
|
let Some(conn) = GeneralConnection::new(conn, state.clone()) else {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Rejected connection: unrecognized or unauthenticated description"
|
|
);
|
|
return;
|
|
};
|
|
|
|
let anonymous_permit = if conn.connection_kind() == ConnectionKind::AnonymousClient {
|
|
match state
|
|
.rho_connection_limits
|
|
.anonymous
|
|
.clone()
|
|
.try_acquire_owned()
|
|
{
|
|
Ok(permit) => Some(permit),
|
|
Err(_) => {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Rejected anonymous connection: anonymous limit reached"
|
|
);
|
|
return;
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let anonymous_ip_permit = if conn.connection_kind() == ConnectionKind::AnonymousClient {
|
|
let Some(peer_ip) = peer_ip else {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Rejected anonymous connection: peer address unavailable"
|
|
);
|
|
return;
|
|
};
|
|
match state
|
|
.rho_connection_limits
|
|
.try_acquire_anonymous_per_ip(peer_ip)
|
|
{
|
|
Some(permit) => Some(permit),
|
|
None => {
|
|
log_err!(
|
|
0,
|
|
PrintType::General,
|
|
"Rejected anonymous connection: per-IP limit reached"
|
|
);
|
|
return;
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let _global_permit = global_permit;
|
|
let _anonymous_permit = anonymous_permit;
|
|
let _anonymous_ip_permit = anonymous_ip_permit;
|
|
conn.handle().await;
|
|
});
|
|
}
|
|
log!(0, PrintType::General, "Server stopped");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{rho_policy, web_config};
|
|
use std::time::Duration;
|
|
|
|
#[test]
|
|
fn rho_connection_budget_configures_mtp_admission() {
|
|
let config = web_config(7).expect("health route is valid");
|
|
assert_eq!(config.max_connections, 7);
|
|
}
|
|
|
|
#[test]
|
|
fn rho_policy_keeps_idle_peers_alive_and_detects_dead_peers() {
|
|
let policy = rho_policy();
|
|
assert_eq!(policy.keep_alive_interval, Some(Duration::from_secs(6)));
|
|
assert_eq!(policy.max_idle_timeout, Some(Duration::from_secs(30)));
|
|
}
|
|
}
|