// WebRTC 传输模块 — 使用 str0m (Sans-IO) 将 H.264 编码帧推送到浏览器
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, UdpSocket};
use std::time::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};
// ── 嵌入式 HTML 测试页面 ──────────────────────────────────────────────────
const HTML_PAGE: &str = r#"
wl-webrtc P0
Connecting...
"#;
// ── 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,
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], frame_number: u64, fps: u32) -> Result<()> {
let should_destroy = if let Some(inner) = self.inner.as_mut() {
inner.write_h264_frame(data, frame_number, fps)?
} 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()))
}
pub fn set_need_keyframe(&mut self) {
if let Some(inner) = self.inner.as_mut() {
inner.need_keyframe = true;
inner.force_keyframe_to_encode = true;
}
}
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,
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;
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}");
}
}
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.need_keyframe = true;
self.force_keyframe_to_encode = true;
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.need_keyframe = true;
self.force_keyframe_to_encode = true;
}
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], frame_number: u64, fps: u32) -> 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 ticks_per_second = 90_000u64;
let fps = fps.max(1) as u64;
let rtp_timestamp = frame_number.saturating_mul(ticks_per_second) / fps;
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
}
}
// ── 工具函数 ──────────────────────────────────────────────────────────────
/// 从 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);
}
}