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_send_stream, 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],
) {
if let Ok(mut keys) = self.connection_application_crypto_keys.write() {
keys.insert(connection_id.to_string(), key);
}
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);
}
}
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 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(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);
}
}
self.application_crypto_key_for_connection(connection_id)
}
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) fn wrap_peer_send_stream_for_connection(
&self,
connection_id: Option<&str>,
send: SendStream,
) -> PeerSendStream {
let key = self.application_crypto_key_for_connection(connection_id);
wrap_peer_send_stream(key, send)
}
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>> {
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, 0, 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_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 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)
}
}
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_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()))
}