use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::{sync::oneshot, task::JoinHandle};
use super::RTP_VERSION;
#[derive(Default)]
pub(super) struct GatewayChildTasks {
heartbeat: Option<GatewayChildTask>,
udp_ping: Option<GatewayChildTask>,
media: Option<GatewayChildTask>,
}
struct GatewayChildTask {
task: JoinHandle<()>,
graceful_stop: Option<oneshot::Sender<()>>,
}
impl GatewayChildTask {
fn new(task: JoinHandle<()>) -> Self {
Self {
task,
graceful_stop: None,
}
}
fn graceful(task: JoinHandle<()>, graceful_stop: oneshot::Sender<()>) -> Self {
Self {
task,
graceful_stop: Some(graceful_stop),
}
}
fn abort(mut self) {
self.graceful_stop.take();
self.task.abort();
}
fn request_shutdown(&mut self) {
if let Some(stop_tx) = self.graceful_stop.take() {
let _ = stop_tx.send(());
} else {
self.task.abort();
}
}
async fn shutdown(mut self) {
self.request_shutdown();
let _ = self.task.await;
}
}
impl GatewayChildTasks {
pub(super) fn has_media(&self) -> bool {
self.media.is_some()
}
pub(super) async fn replace_heartbeat(&mut self, task: JoinHandle<()>) {
Self::replace(&mut self.heartbeat, GatewayChildTask::new(task)).await;
}
pub(super) async fn replace_udp_ping(&mut self, task: JoinHandle<()>) {
Self::replace(&mut self.udp_ping, GatewayChildTask::new(task)).await;
}
pub(super) async fn replace_media(&mut self, task: JoinHandle<()>) {
Self::replace(&mut self.media, GatewayChildTask::new(task)).await;
}
pub(super) async fn shutdown_media(&mut self) {
if let Some(media) = self.media.take() {
media.shutdown().await;
}
}
pub(super) fn install_media_gracefully(
&mut self,
task: JoinHandle<()>,
stop_tx: oneshot::Sender<()>,
) {
assert!(
self.media.is_none(),
"media task must be stopped before install"
);
self.media = Some(GatewayChildTask::graceful(task, stop_tx));
}
async fn replace(slot: &mut Option<GatewayChildTask>, task: GatewayChildTask) {
if let Some(previous) = slot.take() {
previous.shutdown().await;
}
*slot = Some(task);
}
pub(super) async fn shutdown(&mut self) {
let mut tasks = [
self.heartbeat.take(),
self.udp_ping.take(),
self.media.take(),
];
for task in tasks.iter_mut().flatten() {
task.request_shutdown();
}
for task in tasks.into_iter().flatten() {
let _ = task.task.await;
}
}
fn abort_all(&mut self) {
for task in [
self.heartbeat.take(),
self.udp_ping.take(),
self.media.take(),
]
.into_iter()
.flatten()
{
task.abort();
}
}
}
impl Drop for GatewayChildTasks {
fn drop(&mut self) {
self.abort_all();
}
}
pub(super) fn packetize_h264_payloads(frame: &[u8], max_payload_bytes: usize) -> Vec<Vec<u8>> {
if max_payload_bytes <= 2 {
return Vec::new();
}
let mut payloads = Vec::new();
for nal in annex_b_nals(frame) {
if nal.len() <= max_payload_bytes {
payloads.push(nal.to_vec());
continue;
}
let Some((&nal_header, body)) = nal.split_first() else {
continue;
};
let fu_indicator = (nal_header & 0xe0) | 28;
let nal_type = nal_header & 0x1f;
let chunks = body.chunks(max_payload_bytes - 2);
let chunk_count = chunks.len();
for (index, chunk) in chunks.enumerate() {
let mut payload = Vec::with_capacity(chunk.len() + 2);
payload.push(fu_indicator);
payload.push(
nal_type
| if index == 0 { 0x80 } else { 0 }
| if index + 1 == chunk_count { 0x40 } else { 0 },
);
payload.extend_from_slice(chunk);
payloads.push(payload);
}
}
payloads
}
pub(super) fn annex_b_nals(frame: &[u8]) -> impl Iterator<Item = &[u8]> {
let mut next_start = find_annex_b_start(frame, 0);
let mut raw_frame_pending = next_start.is_none() && !frame.is_empty();
std::iter::from_fn(move || {
loop {
let Some((start, start_len)) = next_start.take() else {
return raw_frame_pending.then(|| {
raw_frame_pending = false;
frame
});
};
let nal_start = start + start_len;
next_start = find_annex_b_start(frame, nal_start);
let nal_end = next_start.map_or(frame.len(), |(next, _)| next);
if nal_start < nal_end {
return Some(&frame[nal_start..nal_end]);
}
}
})
}
fn find_annex_b_start(frame: &[u8], mut index: usize) -> Option<(usize, usize)> {
while index + 3 <= frame.len() {
if frame.get(index..index + 4) == Some(&[0, 0, 0, 1]) {
return Some((index, 4));
}
if frame.get(index..index + 3) == Some(&[0, 0, 1]) {
return Some((index, 3));
}
index += 1;
}
None
}
pub(super) fn current_unix_time() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}
pub(super) fn build_rtcp_sender_report(
sender_ssrc: u32,
unix_time: Duration,
rtp_timestamp: u32,
packet_count: u32,
octet_count: u32,
) -> [u8; 28] {
const NTP_UNIX_EPOCH_OFFSET_SECONDS: u64 = 2_208_988_800;
const RTCP_SENDER_REPORT_PACKET_TYPE: u8 = 200;
const RTCP_SENDER_REPORT_LENGTH_WORDS_MINUS_ONE: u16 = 6;
let ntp_seconds = unix_time
.as_secs()
.wrapping_add(NTP_UNIX_EPOCH_OFFSET_SECONDS) as u32;
let ntp_fraction = ((u64::from(unix_time.subsec_nanos()) << 32) / 1_000_000_000) as u32;
let mut report = [0u8; 28];
report[0] = RTP_VERSION << 6;
report[1] = RTCP_SENDER_REPORT_PACKET_TYPE;
report[2..4].copy_from_slice(&RTCP_SENDER_REPORT_LENGTH_WORDS_MINUS_ONE.to_be_bytes());
report[4..8].copy_from_slice(&sender_ssrc.to_be_bytes());
report[8..12].copy_from_slice(&ntp_seconds.to_be_bytes());
report[12..16].copy_from_slice(&ntp_fraction.to_be_bytes());
report[16..20].copy_from_slice(&rtp_timestamp.to_be_bytes());
report[20..24].copy_from_slice(&packet_count.to_be_bytes());
report[24..28].copy_from_slice(&octet_count.to_be_bytes());
report
}
#[cfg(test)]
mod tests {
use super::annex_b_nals;
#[test]
fn annex_b_iterator_handles_three_and_four_byte_start_codes() {
let frame = [0, 0, 1, 0x67, 1, 0, 0, 0, 1, 0x68, 2, 0, 0, 1, 0x65, 3];
let mut nals = annex_b_nals(&frame);
assert_eq!(nals.next(), Some([0x67, 1].as_slice()));
assert_eq!(nals.next(), Some([0x68, 2].as_slice()));
assert_eq!(nals.next(), Some([0x65, 3].as_slice()));
assert_eq!(nals.next(), None);
}
#[test]
fn annex_b_iterator_preserves_raw_nal_fallback() {
let frame = [0x65, 1, 2];
let mut nals = annex_b_nals(&frame);
assert_eq!(nals.next(), Some(frame.as_slice()));
assert_eq!(nals.next(), None);
assert_eq!(annex_b_nals(&[]).next(), None);
}
}