openrtc 1.0.4

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());
        }
    }

    /// Require reciprocal application protection for every trusted native
    /// user-device connection promoted after this call.
    ///
    /// Products with protected named channels should enable this before
    /// initializing presence or auto-connect on every participating runtime.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn set_trusted_user_device_application_crypto_required(&self, required: bool) {
        self.trusted_user_device_application_crypto_required
            .store(required, std::sync::atomic::Ordering::Release);
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub(crate) fn trusted_user_device_application_crypto_is_required(&self) -> bool {
        self.trusted_user_device_application_crypto_required
            .load(std::sync::atomic::Ordering::Acquire)
    }

    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()) {
            // A durable-device projection can briefly retain the prior browser
            // connection id while RTDB has already supplied the replacement
            // endpoint. Never bind that stale generation's crypto key to a
            // stream opened on the new endpoint; validate physical ownership
            // before preferring the projected connection id.
            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)
    }

    /// Wrap an incoming product stream against the exact logical connection
    /// selected by the native ingress router. Product hosts should prefer this
    /// over endpoint-only lookup so a retired same-endpoint record can never
    /// supply a stale application key.
    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}")
        })
    }

    /// Wrap a native product stream against the exact admitted physical
    /// transport generation that produced it. Native host bridges should use
    /// this boundary before exposing stream bytes outside Rust so application
    /// keys remain owned by the runtime and a same-endpoint replacement cannot
    /// lend its key to a retired stream.
    #[cfg(not(target_arch = "wasm32"))]
    pub async fn wrap_incoming_application_bi_stream_for_transport(
        &self,
        endpoint_id: &iroh::EndpointId,
        transport_stable_id: u64,
        send: SendStream,
        recv: RecvStream,
    ) -> anyhow::Result<(PeerSendStream, PeerRecvStream)> {
        let current_transport_stable_id = self
            .get_connection(*endpoint_id)
            .await
            .map(|connection| connection.stable_id() as u64);
        if current_transport_stable_id != Some(transport_stable_id) {
            anyhow::bail!(
                "incoming application stream belongs to a retired transport generation: incoming={} current={:?}",
                transport_stable_id,
                current_transport_stable_id,
            );
        }
        let local_node_id = self
            .current_node_id()
            .await
            .ok_or_else(|| anyhow::anyhow!("local node id unavailable for incoming application stream"))?;
        let connection_id =
            Self::deterministic_connection_id(&local_node_id, &endpoint_id.to_string());
        if !self.native_application_stream_admitted_for_transport(
            &connection_id,
            transport_stable_id,
        ) {
            anyhow::bail!(
                "incoming application stream transport generation is not admitted"
            );
        }
        self.wrap_incoming_application_bi_stream_for_connection(
            &connection_id,
            send,
            recv,
        )
    }

    /// Wrap an admitted incoming application stream with the same per-peer
    /// application-crypto boundary used by `open_peer_bi` on the sender.
    /// Product hosts should call this before parsing channel envelopes or
    /// application frames from the public incoming-stream queue.
    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()))
}