rtc 0.21.0-alpha.2

use crate::peer_connection::configuration::setting_engine::ReplayProtection;
use crate::peer_connection::event::{RTCEventInternal, TaggedRTCEventInternal};
use crate::peer_connection::message::internal::{
    DTLSMessage, RTCMessageInternal, TaggedRTCMessageInternal,
};
use crate::peer_connection::transport::dtls::DtlsTransport;
use crate::peer_connection::transport::dtls::role::RTCDtlsRole;
use crate::peer_connection::transport::dtls::state::RTCDtlsTransportState;
use crate::statistics::accumulator::{CertificateStatsAccumulator, RTCStatsAccumulator};
use crypto::RTCCryptoProvider;
use dtls::endpoint::EndpointEvent;
use dtls::extension::extension_use_srtp::SrtpProtectionProfile;
use dtls::state::State;
use log::{debug, warn};
use shared::TransportContext;
use shared::error::{Error, Result};
use srtp::option::{srtcp_replay_protection, srtp_replay_protection};
use srtp::protection_profile::ProtectionProfile;
use std::collections::VecDeque;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Instant;

pub(crate) struct DtlsHandlerContext {
    pub(crate) dtls_transport: DtlsTransport,

    pub(crate) read_outs: VecDeque<TaggedRTCMessageInternal>,
    pub(crate) write_outs: VecDeque<TaggedRTCMessageInternal>,
    pub(crate) event_outs: VecDeque<TaggedRTCEventInternal>,
}

impl DtlsHandlerContext {
    pub(crate) fn new(dtls_transport: DtlsTransport) -> Self {
        Self {
            dtls_transport,
            read_outs: VecDeque::new(),
            write_outs: VecDeque::new(),
            event_outs: VecDeque::new(),
        }
    }
}

/// DtlsHandler implements DTLS Protocol handling
pub(crate) struct DtlsHandler<'a> {
    ctx: &'a mut DtlsHandlerContext,
    stats: &'a mut RTCStatsAccumulator,
}

impl<'a> DtlsHandler<'a> {
    pub(crate) fn new(ctx: &'a mut DtlsHandlerContext, stats: &'a mut RTCStatsAccumulator) -> Self {
        DtlsHandler { ctx, stats }
    }

    pub(crate) fn name(&self) -> &'static str {
        "DtlsHandler"
    }

    /// Update stats when DTLS handshake completes
    fn update_dtls_stats_from_profile(
        &mut self,
        srtp_profile: SrtpProtectionProfile,
        peer_certificates: &[Vec<u8>],
        dtls_cipher: Option<String>,
    ) -> Result<()> {
        // Update transport DTLS state
        self.stats
            .transport
            .on_dtls_state_changed(RTCDtlsTransportState::Connected);

        // Retain the peer's chain for `DtlsTransport::get_remote_certificates()`. The stats
        // registration below keeps only a fingerprint-derived id, which is not reversible.
        self.ctx.dtls_transport.remote_certificates = peer_certificates.to_vec();

        // Update DTLS role from transport
        self.stats.transport.dtls_role = self.ctx.dtls_transport.dtls_role;

        // Update SRTP cipher from DTLS negotiation
        let srtp_cipher = match srtp_profile {
            SrtpProtectionProfile::Srtp_Aead_Aes_128_Gcm => "SRTP_AEAD_AES_128_GCM",
            SrtpProtectionProfile::Srtp_Aead_Aes_256_Gcm => "SRTP_AEAD_AES_256_GCM",
            SrtpProtectionProfile::Srtp_Aes128_Cm_Hmac_Sha1_80 => "SRTP_AES128_CM_HMAC_SHA1_80",
            SrtpProtectionProfile::Srtp_Aes128_Cm_Hmac_Sha1_32 => "SRTP_AES128_CM_HMAC_SHA1_32",
            _ => "Unknown",
        };
        self.stats.transport.srtp_cipher = srtp_cipher.to_string();

        // Update TLS version
        self.stats.transport.tls_version = "DTLS 1.2".to_string();

        // Update DTLS cipher from the negotiated cipher suite
        if let Some(cipher) = dtls_cipher {
            self.stats.transport.dtls_cipher = cipher;
        }

        // Register local certificate and set local_certificate_id
        if let Some(local_cert) = self.ctx.dtls_transport.certificates.first() {
            let fingerprints =
                local_cert.get_fingerprints(self.ctx.dtls_transport.crypto_provider.crypto())?;
            if let Some(fp) = fingerprints.first() {
                // Register certificate in accumulator
                // Use hex encoding for certificate (base64 would need additional dependency)
                if let Some(der) = local_cert.dtls_certificate.certificate.first() {
                    self.stats.register_certificate(
                        fp.value.clone(),
                        CertificateStatsAccumulator {
                            fingerprint: fp.value.clone(),
                            fingerprint_algorithm: fp.algorithm.clone(),
                            base64_certificate: hex::encode(der.as_ref()),
                            issuer_certificate_id: String::new(),
                        },
                    );
                }

                // Set local certificate ID in transport stats
                self.stats.transport.local_certificate_id = fp.value.clone();
            }
        }

        // Register remote certificate and set remote_certificate_id
        if let Some(peer_cert_der) = peer_certificates.first() {
            // Compute fingerprint from peer certificate
            let hash = self
                .ctx
                .dtls_transport
                .crypto_provider
                .crypto()
                .hash(crypto::HashAlgorithm::Sha256, peer_cert_der)
                .map_err(|error| Error::Crypto(error.to_string()))?;
            let fingerprint: String = hash
                .iter()
                .map(|b| format!("{:02x}", b))
                .collect::<Vec<_>>()
                .join(":");

            let cert_id = format!("remote-certificate-{}", &fingerprint[..8]);

            // Register remote certificate in accumulator
            self.stats.register_certificate(
                cert_id.clone(),
                CertificateStatsAccumulator {
                    fingerprint,
                    fingerprint_algorithm: "sha-256".to_string(),
                    base64_certificate: hex::encode(peer_cert_der),
                    issuer_certificate_id: String::new(),
                },
            );

            // Set remote certificate ID in transport stats
            self.stats.transport.remote_certificate_id = cert_id;
        }
        Ok(())
    }
}

impl<'a>
    sansio::Protocol<TaggedRTCMessageInternal, TaggedRTCMessageInternal, TaggedRTCEventInternal>
    for DtlsHandler<'a>
{
    type Rout = TaggedRTCMessageInternal;
    type Wout = TaggedRTCMessageInternal;
    type Eout = TaggedRTCEventInternal;
    type Error = Error;
    type Time = Instant;

    fn handle_read(&mut self, msg: TaggedRTCMessageInternal) -> Result<()> {
        if let RTCMessageInternal::Dtls(DTLSMessage::Raw(dtls_message)) = msg.message {
            debug!("recv dtls RAW {:?}", msg.transport.peer_addr);

            let mut messages = vec![];
            let mut srtp_contexts = None;
            let mut srtp_profile_for_stats: Option<SrtpProtectionProfile> = None;
            let mut peer_certificates_for_stats: Vec<Vec<u8>> = vec![];
            let mut dtls_cipher_for_stats: Option<String> = None;

            let dtls_endpoint = self
                .ctx
                .dtls_transport
                .dtls_endpoint
                .as_mut()
                .ok_or(Error::ErrDtlsTransportNotStarted)?;

            for message in dtls_endpoint.read(
                msg.now,
                msg.transport.peer_addr,
                msg.transport.ecn,
                dtls_message,
            )? {
                match message {
                    EndpointEvent::HandshakeComplete => {
                        if let Some(state) =
                            dtls_endpoint.get_connection_state(msg.transport.peer_addr)
                        {
                            debug!("dtls handshake complete");

                            // Save profile for stats update after borrow ends
                            srtp_profile_for_stats = Some(state.srtp_protection_profile());

                            // Save peer certificates for stats
                            peer_certificates_for_stats = state.peer_certificates.clone();

                            // Save DTLS cipher suite for stats
                            dtls_cipher_for_stats = state.cipher_suite().map(|cs| cs.to_string());

                            let (local_srtp_context, remote_srtp_context) =
                                DtlsHandler::update_srtp_contexts(
                                    state,
                                    &self.ctx.dtls_transport.replay_protection,
                                    &self.ctx.dtls_transport.crypto_provider,
                                )?;
                            srtp_contexts = Some((local_srtp_context, remote_srtp_context));
                        } else {
                            warn!(
                                "Unable to find connection state for {}",
                                msg.transport.peer_addr
                            );
                        }
                    }
                    EndpointEvent::ApplicationData(message) => {
                        debug!("recv dtls application RAW {:?}", msg.transport.peer_addr);
                        messages.push(message);
                    }
                    _ => {}
                }
            }

            while let Some(transmit) = dtls_endpoint.poll_transmit() {
                self.ctx.write_outs.push_back(TaggedRTCMessageInternal {
                    now: transmit.now,
                    transport: transmit.transport,
                    message: RTCMessageInternal::Dtls(DTLSMessage::Raw(transmit.message)),
                });
            }

            if let Some((local_srtp_context, remote_srtp_context)) = srtp_contexts {
                self.ctx
                    .dtls_transport
                    .state_change(RTCDtlsTransportState::Connected);

                self.ctx.event_outs.push_back(TaggedRTCEventInternal {
                    now: msg.now,
                    event: RTCEventInternal::DTLSHandshakeComplete(
                        Some(local_srtp_context),
                        Some(remote_srtp_context),
                    ),
                });
            }

            // Update stats after dtls_endpoint borrow ends
            if let Some(srtp_profile) = srtp_profile_for_stats {
                self.update_dtls_stats_from_profile(
                    srtp_profile,
                    &peer_certificates_for_stats,
                    dtls_cipher_for_stats,
                )?;
            }

            for message in messages {
                debug!("recv dtls application RAW {:?}", msg.transport.peer_addr);
                self.ctx.read_outs.push_back(TaggedRTCMessageInternal {
                    now: msg.now,
                    transport: msg.transport,
                    message: RTCMessageInternal::Dtls(DTLSMessage::Raw(message)),
                });
            }
        } else {
            // Bypass
            debug!("bypass dtls read {:?}", msg.transport.peer_addr);
            self.ctx.read_outs.push_back(msg);
        }
        Ok(())
    }

    fn poll_read(&mut self) -> Option<Self::Rout> {
        self.ctx.read_outs.pop_front()
    }

    fn handle_write(&mut self, msg: TaggedRTCMessageInternal) -> Result<()> {
        if let RTCMessageInternal::Dtls(DTLSMessage::Raw(dtls_message)) = msg.message {
            debug!("send dtls RAW {:?}", msg.transport.peer_addr);

            let dtls_endpoint = self
                .ctx
                .dtls_transport
                .dtls_endpoint
                .as_mut()
                .ok_or(Error::ErrDtlsTransportNotStarted)?;

            dtls_endpoint.write(msg.now, msg.transport.peer_addr, &dtls_message)?;
            while let Some(transmit) = dtls_endpoint.poll_transmit() {
                self.ctx.write_outs.push_back(TaggedRTCMessageInternal {
                    now: transmit.now,
                    transport: transmit.transport,
                    message: RTCMessageInternal::Dtls(DTLSMessage::Raw(transmit.message)),
                });
            }
        } else {
            // Bypass
            debug!("Bypass dtls write {:?}", msg.transport.peer_addr);
            self.ctx.write_outs.push_back(msg);
        }
        Ok(())
    }

    fn poll_write(&mut self) -> Option<Self::Wout> {
        self.ctx.write_outs.pop_front()
    }

    fn handle_event(&mut self, evt: TaggedRTCEventInternal) -> Result<()> {
        let now = evt.now;
        if let RTCEventInternal::ICESelectedCandidatePairChange = evt.event {
            if self.ctx.dtls_transport.dtls_role == RTCDtlsRole::Client {
                // dtls_endpoint only connect once when acts as DTLSRole::Client
                if let Some(dtls_handshake_config) =
                    self.ctx.dtls_transport.dtls_handshake_config.take()
                {
                    let dtls_endpoint = self
                        .ctx
                        .dtls_transport
                        .dtls_endpoint
                        .as_mut()
                        .ok_or(Error::ErrDtlsTransportNotStarted)?;
                    dtls_endpoint.connect(
                        now,
                        TransportContext::default().peer_addr, // always use default for transport to make DTLS tunneled
                        dtls_handshake_config,
                        None,
                    )?;
                };
            }
        } else {
            self.ctx.event_outs.push_back(evt);
        }
        Ok(())
    }

    fn poll_event(&mut self) -> Option<Self::Eout> {
        self.ctx.event_outs.pop_front()
    }

    fn handle_timeout(&mut self, now: Instant) -> Result<()> {
        let Some(dtls_endpoint) = self.ctx.dtls_transport.dtls_endpoint.as_mut() else {
            return Ok(());
        };

        let remotes: Vec<SocketAddr> = dtls_endpoint.get_connections_keys().copied().collect();
        for remote in remotes {
            dtls_endpoint.handle_timeout(remote, now)?;
        }
        while let Some(transmit) = dtls_endpoint.poll_transmit() {
            self.ctx.write_outs.push_back(TaggedRTCMessageInternal {
                now: transmit.now,
                transport: transmit.transport,
                message: RTCMessageInternal::Dtls(DTLSMessage::Raw(transmit.message)),
            });
        }

        Ok(())
    }

    fn poll_timeout(&mut self) -> Option<Instant> {
        if let Some(dtls_endpoint) = self.ctx.dtls_transport.dtls_endpoint.as_ref() {
            let mut eto = None;

            let remotes = dtls_endpoint.get_connections_keys();
            for remote in remotes {
                if let Some(timeout) = dtls_endpoint.poll_timeout(remote) {
                    eto = Some(eto.map_or(timeout, |e: Instant| e.min(timeout)));
                }
            }

            eto
        } else {
            None
        }
    }

    fn close(&mut self) -> Result<()> {
        self.ctx.dtls_transport.stop()
    }
}

impl<'a> DtlsHandler<'a> {
    const DEFAULT_SESSION_SRTP_REPLAY_PROTECTION_WINDOW: usize = 64;
    const DEFAULT_SESSION_SRTCP_REPLAY_PROTECTION_WINDOW: usize = 64;
    pub(crate) fn update_srtp_contexts(
        state: &State,
        replay_protection: &ReplayProtection,
        crypto_provider: &Arc<dyn RTCCryptoProvider>,
    ) -> Result<(srtp::context::Context, srtp::context::Context)> {
        let profile = match state.srtp_protection_profile() {
            SrtpProtectionProfile::Srtp_Aead_Aes_128_Gcm => ProtectionProfile::AeadAes128Gcm,
            SrtpProtectionProfile::Srtp_Aead_Aes_256_Gcm => ProtectionProfile::AeadAes256Gcm,
            SrtpProtectionProfile::Srtp_Aes128_Cm_Hmac_Sha1_80 => {
                ProtectionProfile::Aes128CmHmacSha1_80
            }
            SrtpProtectionProfile::Srtp_Aes128_Cm_Hmac_Sha1_32 => {
                ProtectionProfile::Aes128CmHmacSha1_32
            }
            _ => return Err(Error::ErrNoSuchSrtpProfile),
        };

        let mut srtp_config = srtp::config::Config {
            profile,
            ..Default::default()
        };
        if replay_protection.srtp != 0 {
            srtp_config.remote_rtp_options =
                Some(srtp::option::srtp_replay_protection(replay_protection.srtp));
        }

        if replay_protection.srtcp != 0 {
            srtp_config.remote_rtcp_options = Some(srtp::option::srtcp_replay_protection(
                replay_protection.srtcp,
            ));
        }

        let keying_material = state.export_keying_material(
            srtp::config::LABEL_EXTRACTOR_DTLS_SRTP,
            &[],
            srtp_config.keying_material_len(),
        )?;
        srtp_config
            .set_session_keys_from_keying_material(keying_material.as_ref(), state.is_client())?;

        let local_context = srtp::context::Context::new(
            &srtp_config.keys.local_master_key,
            &srtp_config.keys.local_master_salt,
            srtp_config.profile,
            srtp_config.local_rtp_options,
            srtp_config.local_rtcp_options,
            crypto_provider.crypto(),
        )?;

        let remote_context = srtp::context::Context::new(
            &srtp_config.keys.remote_master_key,
            &srtp_config.keys.remote_master_salt,
            srtp_config.profile,
            if srtp_config.remote_rtp_options.is_none() {
                Some(srtp_replay_protection(
                    Self::DEFAULT_SESSION_SRTP_REPLAY_PROTECTION_WINDOW,
                ))
            } else {
                srtp_config.remote_rtp_options
            },
            if srtp_config.remote_rtcp_options.is_none() {
                Some(srtcp_replay_protection(
                    Self::DEFAULT_SESSION_SRTCP_REPLAY_PROTECTION_WINDOW,
                ))
            } else {
                srtp_config.remote_rtcp_options
            },
            crypto_provider.crypto(),
        )?;

        Ok((local_context, remote_context))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::peer_connection::transport::{RTCTransportId, TransportKind};

    /// A fixed nonce keeps test ids deterministic while still distinguishing the kinds.
    fn test_transport_id(kind: TransportKind) -> RTCTransportId {
        RTCTransportId::new(0xabcd_ef01_2345_6789, kind)
    }

    #[test]
    fn timeout_before_dtls_starts_is_a_noop() {
        let provider =
            crypto::default_provider().expect("a built-in crypto provider is enabled for tests");
        let transport = DtlsTransport::new(
            crate::peer_connection::transport::dtls::RTCDtlsTransportConfig {
                id: test_transport_id(TransportKind::Dtls),
                ice_transport_id: test_transport_id(TransportKind::Ice),
                certificates: vec![],
                answering_dtls_role: Default::default(),
                srtp_protection_profiles: vec![],
                dtls_cipher_suites: vec![],
                allow_insecure_verification_algorithm: false,
                disable_certificate_fingerprint_verification: false,
                replay_protection: Default::default(),
                crypto_provider: provider,
            },
        )
        .expect("transport");
        let mut context = DtlsHandlerContext::new(transport);
        let mut stats = RTCStatsAccumulator::default();
        let mut handler = DtlsHandler::new(&mut context, &mut stats);

        sansio::Protocol::handle_timeout(&mut handler, Instant::now()).unwrap();
    }

    // The handshake path is the only writer of `remote_certificates`. A test that assigns the
    // field and reads it back would keep passing if this assignment were deleted, so drive the
    // writer instead.
    #[test]
    fn handshake_completion_retains_the_peer_certificate_chain() -> Result<()> {
        let provider =
            crypto::default_provider().expect("a built-in crypto provider is enabled for tests");
        let transport = DtlsTransport::new(
            crate::peer_connection::transport::dtls::RTCDtlsTransportConfig {
                id: test_transport_id(TransportKind::Dtls),
                ice_transport_id: test_transport_id(TransportKind::Ice),
                certificates: vec![],
                answering_dtls_role: Default::default(),
                srtp_protection_profiles: vec![],
                dtls_cipher_suites: vec![],
                allow_insecure_verification_algorithm: false,
                disable_certificate_fingerprint_verification: false,
                replay_protection: Default::default(),
                crypto_provider: provider,
            },
        )
        .expect("transport");
        let mut context = DtlsHandlerContext::new(transport);
        let mut stats = RTCStatsAccumulator::default();

        assert!(context.dtls_transport.get_remote_certificates().is_empty());

        // The fingerprint is hashed from these bytes, so any content works.
        let peer_chain = vec![vec![0x30u8, 0x82, 0x01, 0x0a], vec![0x30, 0x82, 0x02, 0x0b]];

        DtlsHandler::new(&mut context, &mut stats).update_dtls_stats_from_profile(
            SrtpProtectionProfile::Srtp_Aead_Aes_128_Gcm,
            &peer_chain,
            None,
        )?;

        assert_eq!(
            peer_chain,
            context.dtls_transport.get_remote_certificates(),
            "the whole chain is retained, not just the first entry the stats path consumes"
        );
        Ok(())
    }
}