// WebRTC 传输模块 — 使用 str0m (Sans-IO) 将 H.264 编码帧推送到浏览器 use std::io::{Read, Write}; use std::net::{SocketAddr, TcpListener, UdpSocket}; use std::time::{Duration, Instant}; use anyhow::{bail, Result}; use str0m::bwe::{Bitrate, BweKind}; use str0m::change::SdpOffer; use str0m::format::Codec; use str0m::media::{Frequency, MediaKind, MediaTime, Mid, Pt}; use str0m::net::{Protocol, Receive}; use str0m::{Candidate, Event, IceConnectionState, Input, Output, Rtc, RtcConfig}; /// Minimum interval between honored keyframe productions, regardless of source /// (PLI from viewer, connect event, resolution change). Prevents PLI storms /// from causing back-to-back IDRs that swamp the network with multi-hundred-KB /// bursts. See issue #23. const FORCED_KEYFRAME_MIN_INTERVAL: Duration = Duration::from_secs(1); mod html_page; use html_page::HTML_PAGE; // ── WebRTC 状态 ─────────────────────────────────────────────────────────── pub struct WebRtcState { signal_listener: TcpListener, inner: Option, fps: u32, } struct WebRtcInner { rtc: Rtc, socket: UdpSocket, udp_addr: SocketAddr, video_mid: Option, video_pt: Option, connected: bool, need_keyframe: bool, force_keyframe_to_encode: bool, last_forced_keyframe_at: Option, current_bwe_estimate: Option, rtp_clock: u32, buf: Vec, } impl WebRtcState { pub fn new(port: u16, fps: u32) -> Result { let signal_listener = TcpListener::bind(format!("0.0.0.0:{port}"))?; signal_listener.set_nonblocking(true)?; tracing::info!("WebRTC signaling on http://0.0.0.0:{port}/"); Ok(Self { signal_listener, inner: None, fps, }) } pub fn handle_signaling(&mut self) -> Result { let mut handled = false; loop { let (mut stream, _addr) = match self.signal_listener.accept() { Ok(s) => s, Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break, Err(e) => bail!("TCP accept error: {e}"), }; handled = true; stream.set_nonblocking(true)?; let mut req = vec![0u8; 65536]; let n = match stream.read(&mut req) { Ok(n) => n, Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => continue, Err(e) => { tracing::warn!("TCP read error: {e}"); continue; } }; let req_str = String::from_utf8_lossy(&req[..n]); if req_str.starts_with("GET / ") || req_str.starts_with("GET /sdp ") && !req_str.contains("Content-Type: application/json") { let resp = format!( "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", HTML_PAGE.len(), HTML_PAGE ); if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } } else if req_str.starts_with("POST /sdp") { let body = extract_body(&req_str); if body.is_empty() { let resp = "HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\nempty body"; if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } continue; } match WebRtcInner::new(self.fps).and_then(|mut new_inner| { let answer_json = new_inner.handle_sdp_offer(body.as_bytes())?; Ok((new_inner, answer_json)) }) { Ok((new_inner, answer_json)) => { let replacing = self.inner.is_some(); self.inner = Some(new_inner); if replacing { tracing::info!("Replaced WebRTC connection (old dropped)"); } else { tracing::info!("New WebRTC connection"); } let resp = format!( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", answer_json.len(), answer_json ); if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } } Err(e) => { tracing::error!("SDP offer handling failed: {e}"); let resp = "HTTP/1.1 500 Internal Server Error\r\nConnection: close\r\n\r\n"; if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } } } } else { let resp = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n"; if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } } } Ok(handled) } pub fn poll_rtc(&mut self) -> Result<()> { if let Some(inner) = self.inner.as_mut() { if inner.poll_rtc()? { tracing::info!("WebRTC connection closed; clearing connection state"); self.inner = None; } } Ok(()) } pub fn feed_network(&mut self) -> Result<()> { if let Some(inner) = self.inner.as_mut() { inner.feed_network()?; } Ok(()) } pub fn poll_and_feed(&mut self) -> Result<()> { self.poll_rtc()?; self.feed_network()?; self.poll_rtc() } pub fn write_h264_frame(&mut self, data: &[u8], pts_ticks: i64) -> Result<()> { let should_destroy = if let Some(inner) = self.inner.as_mut() { inner.write_h264_frame(data, pts_ticks)? } else { false }; if should_destroy { tracing::info!("WebRTC connection failed during write; clearing connection state"); self.inner = None; } Ok(()) } pub fn is_connected(&self) -> bool { self.inner.as_ref().is_some_and(WebRtcInner::is_connected) } /// Returns the latest bandwidth estimation estimate in bits per second, if available. pub fn get_bwe_estimate(&self) -> Option { self.inner .as_ref() .and_then(|inner| inner.current_bwe_estimate.map(|b| b.as_u64())) } /// Internal keyframe request (connect, resolution change). Always honored, /// but updates last_forced_keyframe_at so a subsequent viewer PLI in the next /// second is throttled. pub fn set_need_keyframe(&mut self) { if let Some(inner) = self.inner.as_mut() { inner.set_need_keyframe(); } } /// External keyframe request from viewer (PLI/FIR via str0m /// `Event::KeyframeRequest`). Rate-limited to FORCED_KEYFRAME_MIN_INTERVAL /// to prevent PLI storms from swamping the network with IDR bursts. /// See issue #23. #[allow(dead_code)] pub fn request_keyframe_from_viewer(&mut self) { if let Some(inner) = self.inner.as_mut() { inner.request_keyframe_from_viewer(); } } pub fn take_force_keyframe(&mut self) -> bool { if let Some(inner) = self.inner.as_mut() { let v = inner.force_keyframe_to_encode; inner.force_keyframe_to_encode = false; v } else { false } } } impl WebRtcInner { fn new(fps: u32) -> Result { let _ = fps; let mut rtc = RtcConfig::new() .enable_bwe(Some(Bitrate::mbps(5))) .build(Instant::now()); let socket = UdpSocket::bind("0.0.0.0:0")?; socket.set_nonblocking(true)?; // Increase UDP send buffer to absorb IDR frame bursts (256KB IDR → ~145 RTP // packets in a single poll_rtc loop). Default Linux wmem is ~208KB which // causes EAGAIN on large keyframes. 2MB comfortably buffers several IDRs. const SND_BUF_REQ: usize = 2 * 1024 * 1024; // SAFETY: fd is a valid UDP socket; setsockopt/getsockopt with SOL_SOCKET + // SO_SNDBUF are safe on Linux. We check the return value and log the actual // kernel-assigned buffer (Linux may cap at wmem_max and/or double the value). unsafe { let fd = std::os::unix::io::AsRawFd::as_raw_fd(&socket); let val: libc::c_int = SND_BUF_REQ as libc::c_int; let ret = libc::setsockopt( fd, libc::SOL_SOCKET, libc::SO_SNDBUF, &val as *const libc::c_int as *const libc::c_void, std::mem::size_of::() as libc::socklen_t, ); if ret < 0 { tracing::warn!( "setsockopt SO_SNDBUF failed (errno {})", std::io::Error::last_os_error() ); } let mut actual: libc::c_int = 0; let mut actual_len: libc::socklen_t = std::mem::size_of::() as libc::socklen_t; let gret = libc::getsockopt( fd, libc::SOL_SOCKET, libc::SO_SNDBUF, &mut actual as *mut libc::c_int as *mut libc::c_void, &mut actual_len, ); if gret == 0 { tracing::info!( "UDP send buffer: requested {}KB, actual {}KB", SND_BUF_REQ / 1024, actual / 1024, ); } } let local_addr = socket.local_addr()?; let lan_ip = local_ip().unwrap_or_else(|| { tracing::debug!("Failed to detect LAN IP, falling back to 127.0.0.1"); "127.0.0.1".to_string() }); let candidate_addr: SocketAddr = format!("{lan_ip}:{}", local_addr.port()).parse()?; let candidate = Candidate::host(candidate_addr, "udp") .map_err(|e| anyhow::anyhow!("candidate: {e}"))?; rtc.add_local_candidate(candidate); tracing::info!("WebRTC UDP: {candidate_addr} (bound 0.0.0.0)"); Ok(Self { rtc, socket, udp_addr: candidate_addr, video_mid: None, video_pt: None, connected: false, need_keyframe: false, force_keyframe_to_encode: false, last_forced_keyframe_at: None, current_bwe_estimate: None, rtp_clock: 0, buf: vec![0u8; 65535], }) } fn handle_sdp_offer(&mut self, body: &[u8]) -> Result { let offer: SdpOffer = serde_json::from_slice(body).map_err(|e| anyhow::anyhow!("parse SDP offer: {e}"))?; let answer = self .rtc .sdp_api() .accept_offer(offer) .map_err(|e| anyhow::anyhow!("accept_offer: {e}"))?; self.need_keyframe = true; self.force_keyframe_to_encode = true; tracing::info!("SDP exchange complete, waiting for ICE/DTLS..."); self.discover_video_params(); let answer_json = serde_json::to_vec(&answer).map_err(|e| anyhow::anyhow!("serialize answer: {e}"))?; String::from_utf8(answer_json).map_err(|e| anyhow::anyhow!("answer utf8: {e}")) } fn discover_video_params(&mut self) { let mid = match self.video_mid { Some(m) => m, None => { tracing::debug!("discover_video_params: no video_mid yet"); return; } }; self.video_pt = None; // Disable str0m's LeakyBucketPacer for this video stream. Default pacing // adds ~100ms send latency per large IDR; our 8Mbps cap + VBV already // provide rate control. BWE stays enabled for adaptation feedback. if let Some(stream_tx) = self.rtc.direct_api().stream_tx_by_mid(mid, None) { stream_tx.set_unpaced(true); } if let Some(writer) = self.rtc.writer(mid) { for pp in writer.payload_params() { tracing::debug!("Codec: pt={:?} spec={:?}", pp.pt(), pp.spec()); if pp.spec().codec.is_video() && pp.spec().codec == Codec::H264 { self.video_pt = Some(pp.pt()); tracing::info!("H.264 payload type: {:?}", pp.pt()); break; } } } if self.video_pt.is_none() { tracing::warn!("discover_video_params: no H.264 codec found for mid={mid}"); } } /// Unthrottled keyframe trigger. Always sets the keyframe flags and refreshes /// `last_forced_keyframe_at` so a follow-up viewer PLI within the next /// `FORCED_KEYFRAME_MIN_INTERVAL` is dropped. fn set_need_keyframe(&mut self) { self.need_keyframe = true; self.force_keyframe_to_encode = true; self.last_forced_keyframe_at = Some(Instant::now()); } /// Throttled keyframe trigger used for viewer-originated PLI/FIR requests. /// Honored only if enough time has elapsed since the last forced keyframe. fn request_keyframe_from_viewer(&mut self) { let now = Instant::now(); let should_honor = self .last_forced_keyframe_at .is_none_or(|last| now.duration_since(last) >= FORCED_KEYFRAME_MIN_INTERVAL); if should_honor { self.last_forced_keyframe_at = Some(now); self.need_keyframe = true; self.force_keyframe_to_encode = true; } else { tracing::warn!( "PLI throttled (last forced keyframe {:?} ago, min interval {:?})", self.last_forced_keyframe_at.map(|t| now.duration_since(t)), FORCED_KEYFRAME_MIN_INTERVAL ); } } fn poll_rtc(&mut self) -> Result { loop { match self.rtc.poll_output() { Ok(Output::Transmit(t)) => { tracing::trace!("TX {} bytes -> {}", t.contents.len(), t.destination); if let Err(e) = self.socket.send_to(&t.contents, t.destination) { if e.kind() == std::io::ErrorKind::WouldBlock { tracing::debug!( "UDP send WouldBlock ({} bytes) — send buffer full", t.contents.len(), ); } else { tracing::warn!("UDP send error to {}: {e}", t.destination); } } } Ok(Output::Event(e)) => { tracing::debug!("RTC event: {e:?}"); match &e { Event::Connected => { tracing::info!("WebRTC connected!"); self.connected = true; self.set_need_keyframe(); self.discover_video_params(); } Event::IceConnectionStateChange(IceConnectionState::Disconnected) => { tracing::warn!("WebRTC disconnected"); self.connected = false; return Ok(true); } Event::MediaAdded(ma) => { tracing::info!("Media added: mid={} kind={:?}", ma.mid, ma.kind); if ma.kind == MediaKind::Video { if let Some(media) = self.rtc.media(ma.mid) { if media.direction().is_sending() && self.video_mid.is_none() { self.video_mid = Some(ma.mid); tracing::info!("Captured video mid: {}", ma.mid); self.discover_video_params(); } } } } Event::KeyframeRequest(_) => { tracing::info!("received keyframe request from viewer"); self.request_keyframe_from_viewer(); } Event::EgressBitrateEstimate(est) => { let bitrate = match est { BweKind::Twcc(b) => *b, BweKind::Remb(_, b) => *b, _ => { tracing::debug!("BWE estimate: unrecognized kind, skipping"); continue; } }; tracing::info!("BWE estimate: {bitrate}"); self.current_bwe_estimate = Some(bitrate); } _ => { tracing::debug!("WebRTC event: {:?}", e); } } } Ok(Output::Timeout(_t)) => break, Err(e) => { tracing::error!("rtc.poll_output error: {e}"); self.connected = false; return Ok(true); } } } Ok(false) } fn feed_network(&mut self) -> Result<()> { let mut recv_count = 0u32; loop { match self.socket.recv_from(&mut self.buf) { Ok((n, source)) => { recv_count += 1; if recv_count <= 5 { tracing::trace!("UDP recv {} bytes from {}", n, source); } let input = Input::Receive( Instant::now(), Receive { proto: Protocol::Udp, source, destination: self.udp_addr, contents: self.buf[..n] .try_into() .map_err(|e| anyhow::anyhow!("receive contents: {e}"))?, }, ); self.rtc.handle_input(input).map_err(|e| { anyhow::anyhow!("handle_input({n} bytes from {source}): {e}") })?; } Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break, Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue, Err(e) => bail!("UDP recv error: {e}"), } } self.rtc .handle_input(Input::Timeout(Instant::now())) .map_err(|e| anyhow::anyhow!("handle timeout: {e}"))?; Ok(()) } fn write_h264_frame(&mut self, data: &[u8], pts_ticks: i64) -> Result { if !self.connected { return Ok(false); } let mid = match self.video_mid { Some(m) => m, None => { tracing::debug!("write_h264: no video_mid"); return Ok(false); } }; let pt = match self.video_pt { Some(p) => p, None => { tracing::debug!("write_h264: no video_pt"); return Ok(false); } }; if self.need_keyframe { if !is_idr_nalu(data) { tracing::debug!( "write_h264: skipping non-IDR frame ({} bytes), waiting for keyframe", data.len() ); return Ok(false); } tracing::info!( "write_h264: got IDR keyframe ({} bytes), starting playback", data.len() ); self.need_keyframe = false; } let rtp_timestamp = rtp_timestamp_from_pts_ticks(pts_ticks); self.rtp_clock = rtp_timestamp as u32; let rtp_time = MediaTime::new(rtp_timestamp, Frequency::NINETY_KHZ); let writer = match self.rtc.writer(mid) { Some(w) => w, None => { tracing::debug!("write_h264: no writer for mid={mid}"); return Ok(false); } }; tracing::debug!( "write_h264: {} bytes, pt={:?}, rtp={}", data.len(), pt, self.rtp_clock ); writer .write(pt, Instant::now(), rtp_time, data) .map_err(|e| anyhow::anyhow!("writer.write: {e}"))?; let should_destroy = self.poll_rtc()?; Ok(should_destroy) } fn is_connected(&self) -> bool { self.connected } } /// Convert PTS in 90kHz media-clock ticks to RTP MediaTime ticks (u64). /// /// With WebRTC encoder time_base = 1/90000, pts_ticks ARE RTP timestamps. /// No fps-based conversion needed. Returned as u64 to feed MediaTime::new /// without premature 13.25-hour u32 wrap; str0m handles RTP u32 wrap internally. pub fn rtp_timestamp_from_pts_ticks(pts_ticks: i64) -> u64 { pts_ticks.max(0) as u64 } // ── 工具函数 ────────────────────────────────────────────────────────────── /// 从 HTTP 请求中提取 body(在 \r\n\r\n 之后) fn extract_body(req: &str) -> &str { if let Some(idx) = req.find("\r\n\r\n") { req.get(idx + 4..).unwrap_or("") } else { "" } } fn local_ip() -> Option { std::net::UdpSocket::bind("0.0.0.0:0").ok().and_then(|s| { s.connect("1.1.1.1:80").ok()?; let addr = s.local_addr().ok()?; drop(s); let ip = addr.ip().to_string(); if ip == "0.0.0.0" || ip.starts_with("127.") { return None; } Some(ip) }) } fn is_idr_nalu(data: &[u8]) -> bool { let mut i = 0; while i < data.len() { let tail = &data[i..]; if tail.starts_with(&[0, 0, 0, 1]) { let Some(&header) = tail.get(4) else { break }; if header & 0x1F == 5 { return true; } i += 5; } else if tail.starts_with(&[0, 0, 1]) { let Some(&header) = tail.get(3) else { break }; if header & 0x1F == 5 { return true; } i += 4; } else { i += 1; } } false } #[cfg(test)] mod tests { use super::*; #[test] fn empty_data() { assert!(!is_idr_nalu(&[])); } #[test] fn short_data_no_start_code() { assert!(!is_idr_nalu(&[0])); assert!(!is_idr_nalu(&[0, 0])); assert!(!is_idr_nalu(&[1, 2, 3])); } #[test] fn three_byte_start_code_no_nal_header() { assert!(!is_idr_nalu(&[0, 0, 1])); } #[test] fn four_byte_start_code_no_nal_header() { assert!(!is_idr_nalu(&[0, 0, 0, 1])); } #[test] fn three_byte_start_code_idr_at_tail() { assert!(is_idr_nalu(&[0, 0, 1, 0x65])); assert!(!is_idr_nalu(&[0, 0, 1, 0x01])); } #[test] fn four_byte_start_code_idr_at_tail() { assert!(is_idr_nalu(&[0, 0, 0, 1, 0x65])); assert!(!is_idr_nalu(&[0, 0, 0, 1, 0x01])); } #[test] fn idr_in_middle_of_frame() { let data: Vec = [ &[0, 0, 0, 1, 0x67][..], // SPS &[0, 0, 0, 1, 0x68][..], // PPS &[0, 0, 0, 1, 0x65][..], // IDR ] .concat(); assert!(is_idr_nalu(&data)); } #[test] fn no_idr_in_frame() { let data: Vec = [ &[0, 0, 0, 1, 0x67][..], // SPS &[0, 0, 0, 1, 0x68][..], // PPS ] .concat(); assert!(!is_idr_nalu(&data)); } #[test] fn mixed_start_code_lengths() { let data: Vec = [ &[0, 0, 0, 1, 0x67][..], // SPS (4-byte start code) &[0, 0, 1, 0x65][..], // IDR (3-byte start code) ] .concat(); assert!(is_idr_nalu(&data)); } #[test] fn all_zeros() { assert!(!is_idr_nalu(&[0, 0, 0, 0, 0, 0, 0, 0])); } // ── Task 5: BWE estimate handling ── #[test] fn bitrate_mbps_conversion() { let b = Bitrate::mbps(5); assert!(b.as_u64() > 0, "5 Mbps should be > 0 bps"); assert_eq!(b.as_u64(), 5_000_000, "5 Mbps should be 5,000,000 bps"); } #[test] fn bitrate_bps_conversion() { let b = Bitrate::bps(1_234_567); assert_eq!(b.as_u64(), 1_234_567); } #[test] fn default_bwe_is_5mbps() { let default = Bitrate::mbps(5); let bps = default.as_u64(); assert_eq!(bps, 5_000_000); } // ── RTP timestamp conversion (issue #24) ── #[test] fn rtp_timestamp_zero_pts() { assert_eq!(rtp_timestamp_from_pts_ticks(0), 0); } #[test] fn rtp_timestamp_one_frame_at_60fps() { // 16.7ms at 90kHz = ~1500 ticks. Real time maps directly to ticks now. assert_eq!(rtp_timestamp_from_pts_ticks(1500), 1500); } #[test] fn rtp_timestamp_one_second() { // 1 second at 90kHz = 90000 ticks assert_eq!(rtp_timestamp_from_pts_ticks(90_000), 90_000); } #[test] fn rtp_timestamp_negative_clamps_to_zero() { assert_eq!(rtp_timestamp_from_pts_ticks(-5), 0); } #[test] fn rtp_timestamp_u64_no_truncation() { // Value above u32::MAX should NOT truncate when feeding MediaTime let large = u32::MAX as i64 + 1000; assert_eq!(rtp_timestamp_from_pts_ticks(large), large as u64); } }