use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock as StdRwLock};
use iroh::endpoint::{RecvStream, SendStream};
use crate::application_crypto::{self, APPLICATION_KEY_BYTES};
use crate::application_crypto_streams::{wrap_peer_streams, PeerRecvStream, PeerSendStream};
use crate::client::Client;
#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
impl Client {
pub fn set_connection_application_crypto_key(
&self,
connection_id: &str,
key: [u8; APPLICATION_KEY_BYTES],
) {
let key_changed = self
.connection_application_crypto_keys
.write()
.map(|mut keys| {
if keys.get(connection_id) == Some(&key) {
false
} else {
keys.insert(connection_id.to_string(), key);
true
}
})
.unwrap_or(false);
if key_changed {
self.clear_connection_application_crypto_confirmation(connection_id);
if let Ok(mut sequences) = self
.connection_application_crypto_outbound_sequences
.write()
{
sequences.insert(connection_id.to_string(), 0);
}
}
}
pub fn set_connection_application_crypto_required(&self, connection_id: &str) {
if let Ok(mut required) = self.connection_application_crypto_required.write() {
required.insert(connection_id.to_string());
}
}
pub fn clear_connection_application_crypto_key(&self, connection_id: &str) {
if let Ok(mut keys) = self.connection_application_crypto_keys.write() {
keys.remove(connection_id);
}
if let Ok(mut required) = self.connection_application_crypto_required.write() {
required.remove(connection_id);
}
if let Ok(mut sequences) = self
.connection_application_crypto_outbound_sequences
.write()
{
sequences.remove(connection_id);
}
if let Ok(mut agreements) = self.connection_application_key_agreements.write() {
agreements.remove(connection_id);
}
self.clear_connection_application_crypto_confirmation(connection_id);
}
pub fn connection_application_crypto_key(
&self,
connection_id: &str,
) -> Option<[u8; APPLICATION_KEY_BYTES]> {
self.application_crypto_key_for_connection(Some(connection_id))
}
pub fn connection_requires_application_crypto(&self, connection_id: &str) -> bool {
self.connection_application_crypto_required
.read()
.ok()
.map(|required| required.contains(connection_id))
.unwrap_or(false)
}
pub(crate) fn connection_requires_application_crypto_confirmation(
&self,
connection_id: &str,
) -> bool {
self.connection_requires_application_crypto(connection_id)
&& self
.connection_application_key_agreements
.read()
.ok()
.is_some_and(|agreements| agreements.contains_key(connection_id))
}
pub(crate) fn confirm_connection_application_crypto(&self, connection_id: &str) {
if let Ok(mut confirmed) = self.connection_application_crypto_confirmed.write() {
confirmed.insert(connection_id.to_string());
}
}
pub(crate) fn connection_application_crypto_is_confirmed(&self, connection_id: &str) -> bool {
self.connection_application_crypto_confirmed
.read()
.ok()
.is_some_and(|confirmed| confirmed.contains(connection_id))
}
pub(crate) fn clear_connection_application_crypto_confirmation(&self, connection_id: &str) {
if let Ok(mut confirmed) = self.connection_application_crypto_confirmed.write() {
confirmed.remove(connection_id);
}
}
pub(crate) fn application_crypto_key_for_connection(
&self,
connection_id: Option<&str>,
) -> Option<[u8; APPLICATION_KEY_BYTES]> {
let connection_id = connection_id?.trim();
if connection_id.is_empty() {
return None;
}
self.connection_application_crypto_keys
.read()
.ok()?
.get(connection_id)
.copied()
}
pub(crate) fn get_or_create_connection_key_agreement(
&self,
connection_id: &str,
) -> Result<crate::key_agreement::EphemeralKeyAgreement, crate::key_agreement::KeyAgreementError>
{
let connection_id = connection_id.trim();
if let Some(existing) = self
.connection_application_key_agreements
.read()
.ok()
.and_then(|agreements| agreements.get(connection_id).cloned())
{
return Ok(existing);
}
let generated = crate::key_agreement::EphemeralKeyAgreement::generate()?;
if let Ok(mut agreements) = self.connection_application_key_agreements.write() {
Ok(agreements
.entry(connection_id.to_string())
.or_insert_with(|| generated.clone())
.clone())
} else {
Ok(generated)
}
}
#[cfg_attr(
any(target_arch = "wasm32", not(feature = "transport-webrtc")),
allow(dead_code)
)]
pub(crate) fn connection_key_agreement_public_key(
&self,
connection_id: &str,
) -> Option<[u8; crate::key_agreement::KEY_AGREEMENT_PUBLIC_KEY_BYTES]> {
self.connection_application_key_agreements
.read()
.ok()?
.get(connection_id.trim())
.map(|agreement| agreement.public_key_bytes())
}
pub(crate) async fn application_crypto_key_for_endpoint(
&self,
endpoint_id: &iroh::EndpointId,
) -> Option<[u8; APPLICATION_KEY_BYTES]> {
let endpoint_str = endpoint_id.to_string();
let records = self
.connection_manager
.get_by_endpoint_id(&endpoint_str)
.await;
for record in records {
if let Some(key) =
self.application_crypto_key_for_connection(Some(&record.connection_id))
{
return Some(key);
}
}
None
}
pub(crate) async fn application_crypto_key_for_connection_or_endpoint(
&self,
connection_id: Option<&str>,
endpoint_id: &iroh::EndpointId,
) -> Option<[u8; APPLICATION_KEY_BYTES]> {
let endpoint_str = endpoint_id.to_string();
if let Some(connection_id) = connection_id.map(str::trim).filter(|id| !id.is_empty()) {
if let Some(record) = self
.connection_manager
.get_by_connection_id(connection_id)
.await
{
let record_matches_endpoint = record
.endpoint_id
.as_deref()
.or(record.node_id.as_deref())
.is_some_and(|value| value == endpoint_str);
if record_matches_endpoint {
if let Some(key) =
self.application_crypto_key_for_connection(Some(connection_id))
{
return Some(key);
}
}
}
}
if let Some(record) = self
.connection_manager
.best_connection_for_peer(&endpoint_str)
.await
{
if let Some(key) =
self.application_crypto_key_for_connection(Some(&record.connection_id))
{
return Some(key);
}
}
let mut records = self
.connection_manager
.get_by_endpoint_id(&endpoint_str)
.await;
records.sort_by(|left, right| {
let left_connected = matches!(
left.state,
crate::connection_manager::ConnectionState::Connected
);
let right_connected = matches!(
right.state,
crate::connection_manager::ConnectionState::Connected
);
left_connected
.cmp(&right_connected)
.then_with(|| left.transport_generation.cmp(&right.transport_generation))
.then_with(|| left.updated_at_ms.cmp(&right.updated_at_ms))
});
for record in records.into_iter().rev() {
if let Some(key) =
self.application_crypto_key_for_connection(Some(&record.connection_id))
{
return Some(key);
}
}
None
}
async fn application_crypto_required_for_connection_or_endpoint(
&self,
connection_id: Option<&str>,
endpoint_id: &iroh::EndpointId,
) -> bool {
if connection_id
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some_and(|value| self.connection_requires_application_crypto(value))
{
return true;
}
let endpoint_id = endpoint_id.to_string();
self.connection_manager
.get_by_endpoint_id(&endpoint_id)
.await
.iter()
.any(|record| self.connection_requires_application_crypto(&record.connection_id))
}
pub(crate) async fn required_application_crypto_key_for_connection_or_endpoint(
&self,
connection_id: Option<&str>,
endpoint_id: &iroh::EndpointId,
operation: &str,
) -> anyhow::Result<Option<[u8; APPLICATION_KEY_BYTES]>> {
if let Some(connection_id) = connection_id
.map(str::trim)
.filter(|value| !value.is_empty())
{
let endpoint_id_text = endpoint_id.to_string();
if self
.connection_manager
.get_by_connection_id(connection_id)
.await
.is_some_and(|record| {
record
.endpoint_id
.as_deref()
.or(record.node_id.as_deref())
.is_some_and(|value| value == endpoint_id_text)
})
{
return enforce_application_crypto_requirement(
self.connection_requires_application_crypto(connection_id),
self.application_crypto_key_for_connection(Some(connection_id)),
operation,
);
}
}
let key = self
.application_crypto_key_for_connection_or_endpoint(connection_id, endpoint_id)
.await;
let required = self
.application_crypto_required_for_connection_or_endpoint(connection_id, endpoint_id)
.await;
enforce_application_crypto_requirement(required, key, operation)
}
pub fn wrap_incoming_application_bi_stream_for_connection(
&self,
connection_id: &str,
send: SendStream,
recv: RecvStream,
) -> anyhow::Result<(PeerSendStream, PeerRecvStream)> {
let key = enforce_application_crypto_requirement(
self.connection_requires_application_crypto(connection_id),
self.application_crypto_key_for_connection(Some(connection_id)),
"incoming application stream",
)?;
wrap_peer_streams(key, send, recv).map_err(|error| {
anyhow::anyhow!("incoming application crypto stream wrap failed: {error}")
})
}
pub async fn wrap_incoming_application_bi_stream(
&self,
endpoint_id: &iroh::EndpointId,
send: SendStream,
recv: RecvStream,
) -> anyhow::Result<(PeerSendStream, PeerRecvStream)> {
let key = self
.required_application_crypto_key_for_connection_or_endpoint(
None,
endpoint_id,
"incoming application stream",
)
.await?;
wrap_peer_streams(key, send, recv).map_err(|error| {
anyhow::anyhow!("incoming application crypto stream wrap failed: {error}")
})
}
pub(crate) fn wrap_peer_streams_for_connection(
&self,
connection_id: Option<&str>,
send: SendStream,
recv: RecvStream,
) -> anyhow::Result<(PeerSendStream, PeerRecvStream)> {
let key = self.application_crypto_key_for_connection(connection_id);
wrap_peer_streams(key, send, recv)
.map_err(|error| anyhow::anyhow!("application crypto stream wrap failed: {error}"))
}
pub(crate) async fn assert_raw_peer_stream_allowed(
&self,
endpoint_id: &iroh::EndpointId,
) -> anyhow::Result<()> {
if self
.application_crypto_key_for_endpoint(endpoint_id)
.await
.is_some()
{
anyhow::bail!(
"raw iroh stream open is not allowed while application crypto is active for this peer; use open_peer_bi/open_peer_uni instead"
);
}
Ok(())
}
pub(crate) fn protect_outbound_application_payload(
&self,
connection_id: &str,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
self.protect_outbound_application_payload_with_type(connection_id, 0, payload)
}
pub(crate) fn protect_outbound_direct_moq_payload(
&self,
connection_id: &str,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
if self
.application_crypto_key_for_connection(Some(connection_id))
.is_none()
{
anyhow::bail!(
"application crypto is required for MoQ route proof, but connection {connection_id} has no installed key"
);
}
self.protect_outbound_application_payload_with_type(
connection_id,
application_crypto::RAW_STREAM_TYPE_ID,
payload,
)
}
fn protect_outbound_application_payload_with_type(
&self,
connection_id: &str,
type_id: u8,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
let keys = self
.connection_application_crypto_keys
.read()
.map_err(|_| anyhow::anyhow!("application crypto key map poisoned"))?;
let Some(key) = keys.get(connection_id) else {
if self.connection_requires_application_crypto(connection_id) {
anyhow::bail!(
"application crypto required for connection {connection_id} but no key is installed"
);
}
return Ok(payload.to_vec());
};
if application_crypto::is_application_encrypted_payload(payload) {
return Ok(payload.to_vec());
}
let sequence = self
.connection_application_crypto_outbound_sequences
.write()
.ok()
.and_then(|mut sequences| {
let entry = sequences.entry(connection_id.to_string()).or_insert(0);
let current = *entry;
*entry = current.saturating_add(1);
Some(current)
})
.unwrap_or(0);
let mut nonce = [0u8; application_crypto::APPLICATION_NONCE_BYTES];
getrandom::getrandom(&mut nonce)
.map_err(|_| anyhow::anyhow!("application crypto random nonce failed"))?;
application_crypto::protect_application_payload_with_nonce_and_sequence(
key, type_id, payload, &nonce, sequence,
)
.map_err(|error| anyhow::anyhow!("application crypto protect failed: {:?}", error))
}
#[cfg_attr(
not(any(feature = "transport-webrtc", feature = "transport-moq")),
allow(dead_code)
)]
pub(crate) fn open_inbound_application_payload(
&self,
connection_id: &str,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
let Some(key) = self.application_crypto_key_for_connection(Some(connection_id)) else {
if self.connection_requires_application_crypto(connection_id) {
anyhow::bail!(
"application crypto required for connection {connection_id} but no key is installed"
);
}
return Ok(payload.to_vec());
};
application_crypto::open_application_payload(&key, 0, payload, true)
.map_err(|error| anyhow::anyhow!("application crypto open failed: {:?}", error))
}
#[cfg(not(target_arch = "wasm32"))]
#[deprecated(
note = "route provenance is required; use handle_inbound_peer_application_frame_for_transport"
)]
pub fn handle_inbound_peer_application_frame(
&self,
_connection_id: &str,
_remote_node_id: Option<&str>,
_transport: &str,
frame: &[u8],
) -> anyhow::Result<bool> {
if frame.first().copied() != Some(0) {
return Ok(false);
}
anyhow::bail!(
"native peer-data route provenance is required; use handle_inbound_peer_application_frame_for_transport"
)
}
#[cfg(not(target_arch = "wasm32"))]
pub async fn handle_inbound_peer_application_frame_for_transport(
&self,
connection_id: &str,
remote_node_id: Option<&str>,
transport: &str,
expected_transport_stable_id: Option<u64>,
frame: &[u8],
) -> anyhow::Result<bool> {
let Some((&type_id, protected_payload)) = frame.split_first() else {
return Ok(false);
};
if type_id != 0 {
return Ok(false);
}
if !application_crypto::is_application_encrypted_payload(protected_payload)
&& self.connection_requires_application_crypto(connection_id)
{
return Ok(false);
}
let payload = self.open_inbound_application_payload(connection_id, protected_payload)?;
let generation = self
.current_native_peer_data_generation(connection_id, expected_transport_stable_id)
.await
.ok_or_else(|| anyhow::anyhow!("native peer-data route generation is stale"))?;
self.emit_native_peer_data(crate::client::NativePeerDataEvent {
connection_id: connection_id.to_string(),
remote_node_id: remote_node_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
transport: transport.to_string(),
transport_stable_id: generation.transport_stable_id,
transport_generation: generation.transport_generation,
route_generation: generation.route_generation,
payload,
});
Ok(true)
}
#[cfg_attr(
not(any(feature = "transport-webrtc", feature = "transport-moq")),
allow(dead_code)
)]
pub(crate) fn open_inbound_direct_moq_payload(
&self,
connection_id: &str,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
let Some(key) = self.application_crypto_key_for_connection(Some(connection_id)) else {
if self.connection_requires_application_crypto(connection_id) {
anyhow::bail!(
"application crypto required for connection {connection_id} but no key is installed"
);
}
return Ok(payload.to_vec());
};
application_crypto::open_application_payload(
&key,
application_crypto::RAW_STREAM_TYPE_ID,
payload,
true,
)
.or_else(|raw_error| {
application_crypto::open_application_payload(&key, 0, payload, true).map_err(
|generic_error| {
anyhow::anyhow!(
"application crypto open failed: raw={:?} generic={:?}",
raw_error,
generic_error,
)
},
)
})
}
pub(crate) fn open_inbound_direct_moq_route_proof(
&self,
connection_id: &str,
payload: &[u8],
) -> anyhow::Result<Vec<u8>> {
let key = self
.application_crypto_key_for_connection(Some(connection_id))
.ok_or_else(|| {
anyhow::anyhow!(
"application crypto is required for MoQ route proof, but connection {connection_id} has no installed key"
)
})?;
application_crypto::open_application_payload(
&key,
application_crypto::RAW_STREAM_TYPE_ID,
payload,
true,
)
.map_err(|error| anyhow::anyhow!("MoQ route proof open failed: {:?}", error))
}
pub(crate) fn connection_ids_requiring_application_crypto(
&self,
connection_ids: &[String],
) -> bool {
self.connection_application_crypto_required
.read()
.ok()
.map(|required| {
connection_ids
.iter()
.any(|connection_id| required.contains(connection_id))
})
.unwrap_or(false)
}
}
fn enforce_application_crypto_requirement(
required: bool,
key: Option<[u8; APPLICATION_KEY_BYTES]>,
operation: &str,
) -> anyhow::Result<Option<[u8; APPLICATION_KEY_BYTES]>> {
if required && key.is_none() {
anyhow::bail!(
"application crypto is required for {operation}, but the current connection has no installed key"
);
}
Ok(key)
}
#[cfg(test)]
mod security_tests {
use super::enforce_application_crypto_requirement;
use crate::application_crypto::APPLICATION_KEY_BYTES;
#[test]
fn required_application_streams_never_downgrade_to_plaintext() {
let error = enforce_application_crypto_requirement(true, None, "test product stream")
.expect_err("required crypto without a key must fail closed");
assert!(error.to_string().contains("no installed key"));
let key = [7_u8; APPLICATION_KEY_BYTES];
assert_eq!(
enforce_application_crypto_requirement(true, Some(key), "test product stream")
.expect("installed key should satisfy the requirement"),
Some(key),
);
assert_eq!(
enforce_application_crypto_requirement(false, None, "manual low-level stream")
.expect("manual low-level plaintext remains an explicit opt-in"),
None,
);
}
}
pub(crate) fn new_connection_application_crypto_key_map(
) -> Arc<StdRwLock<HashMap<String, [u8; APPLICATION_KEY_BYTES]>>> {
Arc::new(StdRwLock::new(HashMap::new()))
}
pub(crate) fn new_connection_application_crypto_required_set() -> Arc<StdRwLock<HashSet<String>>> {
Arc::new(StdRwLock::new(HashSet::new()))
}
pub(crate) fn new_connection_application_crypto_confirmed_set() -> Arc<StdRwLock<HashSet<String>>> {
Arc::new(StdRwLock::new(HashSet::new()))
}
pub(crate) fn new_connection_application_crypto_outbound_sequences(
) -> Arc<StdRwLock<HashMap<String, u64>>> {
Arc::new(StdRwLock::new(HashMap::new()))
}
pub(crate) fn new_connection_application_key_agreement_map(
) -> Arc<StdRwLock<HashMap<String, crate::key_agreement::EphemeralKeyAgreement>>> {
Arc::new(StdRwLock::new(HashMap::new()))
}