openrtc 2.3.0

//! Native Draft 14 MoQ datagram pump beneath Iroh.

#![cfg(all(not(target_arch = "wasm32"), feature = "iroh-transport-moq"))]

use std::{
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex,
    },
    time::Instant,
};

use bytes::Bytes;
use tokio::task::JoinHandle;

use crate::{
    iroh_carrier::{
        segment_packet, CarrierControl, CarrierFrame, CarrierFrameExpectation, CarrierReassembler,
        CARRIER_CONTROL_PACKET_ID, CARRIER_HEADER_LEN,
    },
    packet_carrier_transport::PacketCarrierSession,
    transport::NativeMoQSession,
};

// Keep the encoded object below the effective QUIC datagram ceiling across
// both Draft 14 relay hops. The browser uses the same conservative ceiling;
// its locally reported WebTransport MTU does not describe the downstream hop.
const MOQ_CARRIER_OBJECT_CEILING: usize = 1_000;

pub struct NativeMoqCarrierSession {
    moq: Arc<NativeMoQSession>,
    packet_session: Arc<PacketCarrierSession>,
    outbound_task: JoinHandle<()>,
    expected: CarrierFrameExpectation,
    application_key: [u8; 32],
    terminal_reason: Arc<Mutex<Option<&'static str>>>,
    terminal_acknowledged: Arc<AtomicBool>,
}

impl std::fmt::Debug for NativeMoqCarrierSession {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("NativeMoqCarrierSession")
            .finish_non_exhaustive()
    }
}

impl NativeMoqCarrierSession {
    pub async fn attach(
        moq: Arc<NativeMoQSession>,
        session: PacketCarrierSession,
        expected: CarrierFrameExpectation,
        application_key: [u8; 32],
    ) -> anyhow::Result<Self> {
        anyhow::ensure!(
            moq.state() == crate::transport::NativeMoQState::Connected,
            "MoQ packet-carrier session must be started before attaching Iroh",
        );
        let carrier_message_ceiling = MOQ_CARRIER_OBJECT_CEILING;
        anyhow::ensure!(
            carrier_message_ceiling > CARRIER_HEADER_LEN,
            "MoQ carrier datagram MTU is too small for the carrier header"
        );
        let session = Arc::new(session);
        let reassembler = Arc::new(tokio::sync::Mutex::new(CarrierReassembler::default()));
        let reassembly_clock = Instant::now();
        let inbound_session = session.clone();
        let inbound_reassembler = reassembler;
        let terminal_reason = Arc::new(Mutex::new(None));
        let terminal_acknowledged = Arc::new(AtomicBool::new(false));
        let inbound_terminal_reason = terminal_reason.clone();
        let inbound_terminal_acknowledged = terminal_acknowledged.clone();
        let inbound_moq = moq.clone();
        moq.set_message_handler(Arc::new(move |message: Bytes| {
            let session = inbound_session.clone();
            let reassembler = inbound_reassembler.clone();
            let terminal_reason = inbound_terminal_reason.clone();
            let terminal_acknowledged = inbound_terminal_acknowledged.clone();
            let moq = inbound_moq.clone();
            Box::pin(async move {
                let frame = CarrierFrame::decode(&message, expected);
                let Ok(frame) = frame else {
                    return;
                };
                match frame.terminal_control_kind(expected, &application_key) {
                    Ok(Some(CarrierControl::SessionTokenRevokedAck)) => {
                        terminal_acknowledged.store(true, Ordering::Release);
                        return;
                    }
                    Ok(Some(control @ CarrierControl::SessionTokenRevoked)) => {
                        if let Ok(ack) = CarrierFrame::terminal_control(
                            expected,
                            &application_key,
                            CarrierControl::SessionTokenRevokedAck,
                        )
                        .encode()
                        {
                            let _ = moq.send(&ack).await;
                        }
                        if let (Some(lifecycle_reason), Ok(mut reason)) =
                            (control.lifecycle_reason(), terminal_reason.lock())
                        {
                            *reason = Some(lifecycle_reason);
                        }
                        session.close();
                        return;
                    }
                    Err(_) if frame.header.packet_id == CARRIER_CONTROL_PACKET_ID => return,
                    _ => {}
                }
                let now_ms =
                    u64::try_from(reassembly_clock.elapsed().as_millis()).unwrap_or(u64::MAX);
                let Ok(Some(packet)) = reassembler.lock().await.push(frame, now_ms) else {
                    return;
                };
                let _ = session.deliver_inbound(packet).await;
            })
        }))
        .await;

        let outbound_moq = moq.clone();
        let outbound_session = session.clone();
        let outbound_task = tokio::spawn(async move {
            // Activating the custom Iroh address can enqueue QUIC packets before
            // the peer's MoQ subscription reaches this relay. Keep those packets
            // in the existing bounded carrier queue until the track is writable;
            // treating this normal setup window as terminal strands the physical
            // replacement before either peer can prove it.
            if !outbound_moq
                .wait_for_peer_data_subscription(std::time::Duration::from_secs(25))
                .await
            {
                return;
            }
            let mut packet_id = 0_u64;
            while let Ok(packet) = outbound_session.recv_outbound().await {
                packet_id = packet_id.wrapping_add(1);
                if packet_id == CARRIER_CONTROL_PACKET_ID {
                    packet_id = 0;
                }
                let Ok(frames) =
                    segment_packet(&packet, expected, packet_id, carrier_message_ceiling)
                else {
                    break;
                };
                for frame in frames {
                    let Ok(encoded) = frame.encode() else {
                        return;
                    };
                    if let Err(error) = outbound_moq.send(&encoded).await {
                        eprintln!("[NativeMoQ] Iroh carrier datagram send failed: {error:#}");
                        return;
                    }
                }
            }
        });
        Ok(Self {
            moq,
            packet_session: session,
            outbound_task,
            expected,
            application_key,
            terminal_reason,
            terminal_acknowledged,
        })
    }

    pub fn take_terminal_reason(&self) -> Option<&'static str> {
        self.terminal_reason.lock().ok()?.take()
    }

    pub async fn send_terminal(&self, reason: &str) -> anyhow::Result<()> {
        let control = match reason {
            crate::lifecycle_reason::REASON_SESSION_TOKEN_REVOKED => {
                CarrierControl::SessionTokenRevoked
            }
            _ => return Ok(()),
        };
        let encoded = CarrierFrame::terminal_control(self.expected, &self.application_key, control)
            .encode()?;
        self.terminal_acknowledged.store(false, Ordering::Release);
        for _ in 0..12 {
            self.moq.send(&encoded).await?;
            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
            if self.terminal_acknowledged.load(Ordering::Acquire) {
                break;
            }
        }
        Ok(())
    }
}

impl Drop for NativeMoqCarrierSession {
    fn drop(&mut self) {
        // NativeMoQSession retains the receive callback, which in turn keeps
        // an Arc to this packet session while it unwinds. Close the exact
        // session generation explicitly so replacement does not wait for the
        // callback Arc to drop or fail with DuplicateSession.
        self.packet_session.close();
        self.outbound_task.abort();
        self.moq.close();
    }
}

#[cfg(all(test, feature = "test-harness", feature = "iroh-protocols-wasm"))]
mod tests {
    use super::*;
    use crate::{
        client::MoQConfig, iroh_carrier_kind::EXPERIMENTAL_MOQ_TRANSPORT_ID,
        packet_carrier_transport::PacketCarrierTransport, transport::moq::test_relay,
    };
    use iroh::{endpoint::presets, Endpoint, EndpointAddr, RelayMode, TransportAddr};

    const TEST_ALPN: &[u8] = b"openrtc/test/moq-iroh-carrier/1";
    const TEST_SESSION_ID: [u8; 16] = [0x2a; 16];

    #[tokio::test]
    async fn two_iroh_endpoints_cross_draft14_datagram_relay() {
        let probe = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = probe.local_addr().unwrap().port();
        drop(probe);
        let relay = tokio::spawn(async move {
            let _ = test_relay::run(port).await;
        });
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;

        let secret_a = iroh::SecretKey::generate();
        let secret_b = iroh::SecretKey::generate();
        let endpoint_id_a = secret_a.public();
        let endpoint_id_b = secret_b.public();
        let transport_a = Arc::new(
            PacketCarrierTransport::new(EXPERIMENTAL_MOQ_TRANSPORT_ID, endpoint_id_a).unwrap(),
        );
        let transport_b = Arc::new(
            PacketCarrierTransport::new(EXPERIMENTAL_MOQ_TRANSPORT_ID, endpoint_id_b).unwrap(),
        );
        let packet_session_a = transport_a
            .attach_session(transport_a.peer_addr(endpoint_id_b))
            .unwrap();
        let remote_addr_b = packet_session_a.remote_addr().clone();
        let packet_session_b = transport_b
            .attach_session(transport_b.peer_addr(endpoint_id_a))
            .unwrap();

        let config = MoQConfig {
            relay_url: format!("https://localhost:{port}/moq"),
            access_token: None,
        };
        let carrier_session_id = hex::encode(TEST_SESSION_ID);
        let moq_a = Arc::new(
            NativeMoQSession::new_packet_carrier(
                &endpoint_id_a.to_string(),
                &endpoint_id_b.to_string(),
                &config,
                &carrier_session_id,
            )
            .await
            .unwrap(),
        );
        let moq_b = Arc::new(
            NativeMoQSession::new_packet_carrier(
                &endpoint_id_b.to_string(),
                &endpoint_id_a.to_string(),
                &config,
                &carrier_session_id,
            )
            .await
            .unwrap(),
        );
        moq_a.start().await.unwrap();
        moq_b.start().await.unwrap();
        let expected = CarrierFrameExpectation {
            transport_id: EXPERIMENTAL_MOQ_TRANSPORT_ID,
            carrier_session_id: TEST_SESSION_ID,
            transport_generation: 7,
        };
        let application_key = [0x44; 32];
        let carrier_a =
            NativeMoqCarrierSession::attach(moq_a, packet_session_a, expected, application_key)
                .await
                .unwrap();
        let carrier_b =
            NativeMoqCarrierSession::attach(moq_b, packet_session_b, expected, application_key)
                .await
                .unwrap();

        // Attach before either side has necessarily observed its peer's
        // subscription. The bounded pump must hold Iroh's first QUIC packets
        // through this normal MoQ setup window instead of terminating.
        let endpoint_a = Endpoint::builder(presets::N0)
            .secret_key(secret_a)
            .relay_mode(RelayMode::Disabled)
            .clear_ip_transports()
            .alpns(vec![
                TEST_ALPN.to_vec(),
                iroh_blobs::ALPN.to_vec(),
                iroh_gossip::ALPN.to_vec(),
                iroh_docs::ALPN.to_vec(),
            ])
            .add_custom_transport(transport_a)
            .bind()
            .await
            .unwrap();
        let endpoint_b = Endpoint::builder(presets::N0)
            .secret_key(secret_b)
            .relay_mode(RelayMode::Disabled)
            .clear_ip_transports()
            .alpns(vec![
                TEST_ALPN.to_vec(),
                iroh_blobs::ALPN.to_vec(),
                iroh_gossip::ALPN.to_vec(),
                iroh_docs::ALPN.to_vec(),
            ])
            .add_custom_transport(transport_b)
            .bind()
            .await
            .unwrap();

        let server = {
            let endpoint_b = endpoint_b.clone();
            tokio::spawn(async move {
                let incoming = endpoint_b.accept().await.expect("incoming MoQ carrier");
                let connection = incoming.await.expect("accepted MoQ carrier connection");
                assert_eq!(connection.alpn(), TEST_ALPN);
                assert!(connection.paths().iter().any(|path| {
                    path.is_selected()
                        && matches!(
                            path.remote_addr(),
                            TransportAddr::Custom(address)
                                if address.id() == EXPERIMENTAL_MOQ_TRANSPORT_ID
                        )
                }));
                let (mut send, mut recv) = connection.accept_bi().await.unwrap();
                assert_eq!(
                    recv.read_to_end(1024).await.unwrap(),
                    b"iroh-over-draft14-moq"
                );
                send.write_all(b"bilateral-iroh-response").await.unwrap();
                send.finish().unwrap();
                let _ = connection.closed().await;
            })
        };

        let address_b =
            EndpointAddr::new(endpoint_id_b).with_addrs([TransportAddr::Custom(remote_addr_b)]);
        let connection = tokio::time::timeout(
            std::time::Duration::from_secs(15),
            endpoint_a.connect(address_b, TEST_ALPN),
        )
        .await
        .expect("MoQ-backed Iroh connect timed out")
        .expect("MoQ-backed Iroh connect failed");
        assert!(connection.paths().iter().any(|path| {
            path.is_selected()
                && matches!(
                    path.remote_addr(),
                    TransportAddr::Custom(address)
                        if address.id() == EXPERIMENTAL_MOQ_TRANSPORT_ID
                )
        }));

        let (mut send, mut recv) = connection.open_bi().await.unwrap();
        send.write_all(b"iroh-over-draft14-moq").await.unwrap();
        send.finish().unwrap();
        assert_eq!(
            recv.read_to_end(1024).await.unwrap(),
            b"bilateral-iroh-response"
        );

        connection.close(0u8.into(), b"test-complete");
        server.await.unwrap();
        crate::native_carrier_protocol_test::assert_docs_blobs_gossip_over_custom_transport(
            endpoint_a.clone(),
            endpoint_b.clone(),
        )
        .await
        .unwrap();
        carrier_a
            .send_terminal(crate::lifecycle_reason::REASON_SESSION_TOKEN_REVOKED)
            .await
            .unwrap();
        tokio::time::timeout(std::time::Duration::from_secs(2), async {
            loop {
                if carrier_b.take_terminal_reason()
                    == Some(crate::lifecycle_reason::REASON_SESSION_TOKEN_REVOKED)
                {
                    break;
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("authenticated carrier terminal marker was not delivered");
        endpoint_a.close().await;
        endpoint_b.close().await;
        drop(carrier_a);
        drop(carrier_b);
        relay.abort();
    }
}