simple-instant-stream/src/webrtc.rs

use dashmap::DashMap;
use smol::{
    channel::{Receiver, Sender},
    net::UdpSocket,
};
use std::{
    error::Error,
    sync::Arc,
    time::{Duration, Instant},
};

use str0m::{
    Candidate, Event, IceConnectionState, Input, Output, Rtc,
    change::SdpOffer,
    media::{MediaKind, MediaTime, Mid},
    net::{Protocol, Receive},
};

use crate::{StreamSession, media::VideoFrame};

pub struct Webrtc {
    pub offer_rx: Receiver<(String, String)>,
    pub accept_tx: Sender<(String, String)>,
    pub sessions_ref: Arc<DashMap<String, StreamSession>>,
}

impl Webrtc {
    pub fn start(self) -> Result<(), Box<dyn Error>> {
        smol::spawn(async move {
            while let Ok(offer) = self.offer_rx.recv().await {
                let (stream_key, sdp_body) = offer;
                let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
                let local_addr = socket.local_addr().unwrap();

                let mut builder = Rtc::builder();
                {
                    let cc = builder.codec_config();
                    cc.enable_h264(false);
                    cc.add_h264(102.into(), None, true, 0x42e01f);
                    cc.add_h264(104.into(), None, true, 0x4d001f);
                    cc.add_h264(106.into(), None, true, 0x64001f);
                }
                let mut rtc = builder.build(Instant::now());

                let candidate = Candidate::host(local_addr, Protocol::Udp).unwrap();
                rtc.add_local_candidate(candidate);

                let offer_sdp = SdpOffer::from_sdp_string(&sdp_body).unwrap();
                let mut changes = rtc.sdp_api();
                let mid = changes.add_media(
                    MediaKind::Video,
                    str0m::media::Direction::SendOnly,
                    Some(stream_key.clone()),
                    Some("video0".to_string()),
                    None,
                );
                let offer_answer = match changes.accept_offer(offer_sdp) {
                    Ok(a) => a,
                    Err(e) => {
                        println!("accept_offer failed: {:?}", e);
                        continue;
                    }
                };
                let answer_sdp = offer_answer.to_sdp_string();

                self.accept_tx
                    .send((stream_key.clone(), answer_sdp))
                    .await
                    .unwrap();

                let sessions_ref = self.sessions_ref.clone();
                smol::spawn(async move {
                    Webrtc::detach_connection(socket, rtc, sessions_ref, mid).await;
                })
                .detach();
            }
        })
        .detach();
        Ok(())
    }

    async fn detach_connection(
        socket: UdpSocket,
        mut rtc: Rtc,
        sessions_ref: Arc<DashMap<String, StreamSession>>,
        _hint_mid: Mid,
    ) {
        let mut video_mid: Option<Mid> = None;
        let mut video_pt = None;
        let mut connected = false;
        let mut video_stream: Option<async_broadcast::Receiver<Arc<VideoFrame>>> = None;

        let mut recv_buf = vec![0u8; 65535];
        let local_addr = socket.local_addr().unwrap();

        loop {
            let deadline = loop {
                match rtc.poll_output() {
                    Ok(Output::Timeout(t)) => break t,
                    Ok(Output::Transmit(t)) => {
                        if socket.send_to(&t.contents, t.destination).await.is_err() {
                            return;
                        }
                    }
                    Ok(Output::Event(e)) => match e {
                        Event::MediaAdded(ma) => {
                            if ma.kind == MediaKind::Video {
                                if let Some(writer) = rtc.writer(ma.mid) {
                                    let best = writer
                                        .payload_params()
                                        .max_by_key(|p| {
                                            p.spec().format.profile_level_id.unwrap_or(0)
                                        });
                                    if let Some(params) = best {
                                        println!("Selected PT {:?}", params.pt());
                                        video_pt = Some(params.pt());
                                        video_mid = Some(ma.mid);
                                    }
                                }
                            }
                        }
                        Event::IceConnectionStateChange(state) => {
                            println!("ICE state: {:?}", state);
                        }
                        Event::Connected => {
                            println!("DTLS+ICE connected, ready for media");
                            connected = true;
                        }
                        _ => {}
                    },
                    Err(_) => return,
                }
            };

            if connected {
                if video_stream.is_none() {
                    if let Some(session) = sessions_ref.get("test") {
                        video_stream = Some(session.frame_channel.new_receiver());
                    }
                }
                if let Some(ref mut stream) = video_stream {
                    // Drain at most 8 frames per loop tick so the UDP socket
                    // (ICE keepalives, RTCP) is not starved by a backlog.
                    for _ in 0..8 {
                        match stream.try_recv() {
                            Ok(frame) => {
                                let now = Instant::now();
                                // Explicit 90 kHz clock for H.264 RTP timestamps.
                                let rtp_time = MediaTime::from_90khz(frame.timestamp_ms as u64 * 90);
                                if let (Some(pt), Some(writer)) =
                                    (video_pt, video_mid.and_then(|m| rtc.writer(m)))
                                {
                                    if let Err(e) =
                                        writer.write(pt, now, rtp_time, frame.data.to_vec())
                                    {
                                        println!("write error: {:?}", e);
                                    }
                                }
                            }
                            Err(async_broadcast::TryRecvError::Empty) => break,
                            Err(async_broadcast::TryRecvError::Closed) => return,
                            // Overflow means some frames were dropped; the next
                            // try_recv will give the oldest surviving frame, so
                            // continue draining rather than breaking.
                            Err(async_broadcast::TryRecvError::Overflowed(_)) => continue,
                        }
                    }
                }
            }

            // Cap wait to 20 ms so frame delivery stays timely even when
            // str0m's deadline is far out.
            let wait_until = deadline.min(Instant::now() + Duration::from_millis(20)).max(Instant::now());

            let input = smol::future::or(
                async {
                    smol::Timer::at(wait_until).await;
                    None
                },
                async {
                    let (n, from) = socket.recv_from(&mut recv_buf).await.ok()?;
                    Some((n, from))
                },
            )
            .await;

            match input {
                None => {
                    rtc.handle_input(Input::Timeout(Instant::now())).ok();
                }
                Some((n, from)) => {
                    let data = recv_buf[..n].to_vec();
                    if let Ok(contents) = data.as_slice().try_into() {
                        rtc.handle_input(Input::Receive(
                            Instant::now(),
                            Receive {
                                proto: Protocol::Udp,
                                source: from,
                                destination: local_addr,
                                contents,
                            },
                        ))
                        .ok();
                    }
                }
            }
        }
    }
}