rtc 0.20.0-rc.2

use crate::peer_connection::event::RTCEventInternal;
use crate::peer_connection::event::RTCPeerConnectionEvent;
use crate::peer_connection::message::internal::{
    RTCMessageInternal, STUNMessage, TaggedRTCMessageInternal,
};
use crate::peer_connection::transport::ice::RTCIceTransport;
use crate::statistics::accumulator::RTCStatsAccumulator;
use crate::statistics::stats::ice_candidate_pair::RTCStatsIceCandidatePairState;
use log::{debug, trace};
use shared::error::{Error, Result};
use shared::{TransportContext, TransportMessage};
use std::collections::VecDeque;
use std::time::Instant;

#[derive(Default)]
pub(crate) struct IceHandlerContext {
    pub(crate) ice_transport: RTCIceTransport,

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

    /// Cached stats key of the currently selected ICE candidate pair
    /// (`RTCIceCandidatePair_<local>_<remote>`), refreshed only when the selected
    /// pair changes. The per-packet read/write path looks the pair's stats
    /// accumulator up with this key instead of re-`format!`ing it on every packet
    /// — that formatting was a heap allocation on the hot path in both directions.
    pub(crate) selected_pair_key: Option<String>,
}

impl IceHandlerContext {
    pub(crate) fn new(ice_transport: RTCIceTransport) -> Self {
        Self {
            ice_transport,

            read_outs: VecDeque::new(),
            write_outs: VecDeque::new(),
            event_outs: VecDeque::new(),

            selected_pair_key: None,
        }
    }
}

/// IceHandler implements ICE Protocol handling
pub(crate) struct IceHandler<'a> {
    ctx: &'a mut IceHandlerContext,
    stats: &'a mut RTCStatsAccumulator,
}

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

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

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

    fn handle_read(&mut self, mut msg: TaggedRTCMessageInternal) -> Result<()> {
        if let RTCMessageInternal::Stun(STUNMessage::Raw(message)) = msg.message {
            self.ctx.ice_transport.agent.handle_read(TransportMessage {
                now: msg.now,
                transport: msg.transport,
                message,
            })?;
        } else if self
            .ctx
            .ice_transport
            .agent
            .get_selected_candidate_pair()
            .is_some()
        {
            // only ICE connection is ready and bypass it
            debug!("bypass ice read {:?}", msg.transport.peer_addr);

            // Track packets/bytes received on the selected candidate pair. Look
            // the accumulator up via the cached key (set on the last
            // SelectedCandidatePairChange) to avoid allocating the key per packet.
            let bytes = msg.message.len();
            if let Some(key) = &self.ctx.selected_pair_key
                && let Some(pair) = self.stats.ice_candidate_pairs.get_mut(key)
            {
                pair.on_packet_received(bytes, msg.now);
            }

            // When ICE restarts and the selected candidate pair changes,
            // WebRTC treats this as a path migration, and DTLS continues unchanged, bound to the ICE transport, not to a fixed 5-tuple.
            // Use default for transport to make DTLS tunneled
            msg.transport = TransportContext::default();
            self.ctx.read_outs.push_back(msg);
        } else {
            trace!(
                "drop message from {:?} to {:?} before ICE connection is connected",
                msg.transport.peer_addr, msg.transport.local_addr
            );
        }

        Ok(())
    }

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

    fn handle_write(&mut self, mut msg: TaggedRTCMessageInternal) -> Result<()> {
        if let Some((local, remote)) = self.ctx.ice_transport.agent.get_selected_candidate_pair() {
            // use ICE selected candidate pair to replace local/peer addr and
            // transport protocol: DTLS/SCTP endpoints run on placeholder
            // transport contexts, so without the protocol rewrite their
            // transmits stay stamped UDP even when the selected pair is TCP,
            // and consumers routing poll_write() output by transport_protocol
            // send every non-STUN packet down the wrong transport.
            msg.transport.local_addr = local.addr();
            msg.transport.peer_addr = remote.addr();
            msg.transport.transport_protocol = local.network_type().to_protocol();
            debug!("Bypass ice write {:?}", msg.transport.peer_addr);

            // Track packets/bytes sent on the selected candidate pair. Look the
            // accumulator up via the cached key to avoid allocating it per packet.
            let bytes = msg.message.len();
            if let Some(key) = &self.ctx.selected_pair_key
                && let Some(pair) = self.stats.ice_candidate_pairs.get_mut(key)
            {
                pair.on_packet_sent(bytes, msg.now);
            }

            self.ctx.write_outs.push_back(msg);
        } else {
            trace!(
                "drop message from {:?} to {:?} before ICE connection is connected",
                msg.transport.local_addr, msg.transport.peer_addr,
            );
        }

        Ok(())
    }

    fn poll_write(&mut self) -> Option<Self::Wout> {
        while let Some(transmit) = self.ctx.ice_transport.agent.poll_write() {
            self.ctx.write_outs.push_back(TaggedRTCMessageInternal {
                now: transmit.now,
                transport: transmit.transport,
                message: RTCMessageInternal::Stun(STUNMessage::Raw(transmit.message)),
            });
        }

        self.ctx.write_outs.pop_front()
    }

    fn handle_event(&mut self, evt: RTCEventInternal) -> Result<()> {
        self.ctx.event_outs.push_back(evt);
        Ok(())
    }

    fn poll_event(&mut self) -> Option<Self::Eout> {
        if let Some(evt) = self.ctx.ice_transport.agent.poll_event() {
            match evt {
                ::ice::Event::ConnectionStateChange(state) => {
                    let ice_connection_state = state.into();
                    self.ctx.ice_transport.ice_connection_state = ice_connection_state;

                    // Update transport stats for ICE state change
                    // Use the original ice state which converts to RTCIceTransportState
                    self.stats.transport.on_ice_state_changed(state.into());

                    self.ctx
                        .event_outs
                        .push_back(RTCEventInternal::RTCPeerConnectionEvent(
                            RTCPeerConnectionEvent::OnIceConnectionStateChangeEvent(
                                ice_connection_state,
                            ),
                        ));
                }
                ::ice::Event::SelectedCandidatePairChange(local, remote) => {
                    debug!(
                        "ice selected candidate pair {:?} <-> {:?}",
                        local.addr(),
                        remote.addr()
                    );

                    // Create/update the candidate pair accumulator
                    let pair = self
                        .stats
                        .get_or_create_candidate_pair(local.id(), remote.id());
                    pair.nominated = true;
                    pair.state = RTCStatsIceCandidatePairState::Succeeded;

                    // Cache the pair's stats key so the per-packet path can look
                    // the accumulator up by borrow instead of re-formatting (and
                    // allocating) the key on every read/write.
                    let key = format!("RTCIceCandidatePair_{}_{}", local.id(), remote.id());
                    self.ctx.selected_pair_key = Some(key.clone());
                    // Update transport stats for selected candidate pair change
                    self.stats.transport.on_selected_candidate_pair_changed(key);

                    self.ctx
                        .event_outs
                        .push_back(RTCEventInternal::ICESelectedCandidatePairChange);
                }
                ::ice::Event::RoleChange(is_controlling) => {
                    self.stats.transport.on_ice_role_changed(is_controlling);
                }
            }
        }

        self.ctx.event_outs.pop_front()
    }

    fn handle_timeout(&mut self, now: Instant) -> Result<()> {
        self.ctx.ice_transport.agent.handle_timeout(now)
    }

    fn poll_timeout(&mut self) -> Option<Instant> {
        self.ctx.ice_transport.agent.poll_timeout()
    }

    fn close(&mut self) -> Result<()> {
        self.ctx.ice_transport.agent.close()
    }
}