1use std::collections::BTreeMap;2use std::time::Duration;345use sigil_protocol::MediaFrame;6use sigil_protocol::{7 FrameFlags, KeyframeRequestReasonV3, MAX_MEDIA_GROUP_BYTES_V3, MediaObjectV3, ProtocolError,8 read_media_object_v3,9};1011const CLIENT_MEDIA_OBJECT_CAPACITY: usize = 4;12const CLIENT_MEDIA_OBJECT_READ_TIMEOUT: Duration = Duration::from_secs(1);131415pub(crate) enum MediaObjectReadOutcomeV3 {16 Object {17 accept_index: u64,18 object: MediaObjectV3,19 },20 Dropped {21 accept_index: u64,22 reason: KeyframeRequestReasonV3,23 },24 Malformed(String),25}2627impl MediaObjectReadOutcomeV3 {28 fn accept_index(&self) -> Option<u64> {29 match self {30 Self::Object { accept_index, .. } | Self::Dropped { accept_index, .. } => {31 Some(*accept_index)32 }33 Self::Malformed(_) => None,34 }35 }3637 fn is_fast_forward_barrier(&self) -> bool {38 let Self::Object { object, .. } = self else {39 return false;40 };41 object.header.object_id == 042 && object.header.flags.contains(FrameFlags::KEYFRAME)43 && object.header.flags.contains(FrameFlags::CODEC_CONFIG)44 && object.header.flags.contains(FrameFlags::DISCONTINUITY)45 }46}474849struct MediaObjectReorderV3 {50 next_accept_index: u64,51 completed: BTreeMap<u64, MediaObjectReadOutcomeV3>,52}5354impl MediaObjectReorderV3 {55 fn new(first_accept_index: u64) -> Self {56 Self {57 next_accept_index: first_accept_index,58 completed: BTreeMap::new(),59 }60 }6162 fn pending_len(&self) -> usize {63 self.completed.len()64 }6566 fn push(67 &mut self,68 outcome: MediaObjectReadOutcomeV3,69 ) -> Result<Option<MediaObjectReadOutcomeV3>, String> {70 let Some(accept_index) = outcome.accept_index() else {71 return Ok(Some(outcome));72 };73 if accept_index < self.next_accept_index {74 // A discontinuity barrier may advance beyond older in-flight75 // reads. Their eventual timeout/reset outcomes belong to the76 // superseded GOP and must not poison the recovered sequence.77 return Ok(None);78 }79 if self.completed.insert(accept_index, outcome).is_some() {80 return Err(format!(81 "Duplicate media v3 accept index {accept_index} completed"82 ));83 }84 if self85 .completed86 .get(&accept_index)87 .is_some_and(MediaObjectReadOutcomeV3::is_fast_forward_barrier)88 {89 self.completed.retain(|index, _| *index >= accept_index);90 self.next_accept_index = accept_index;91 }92 self.take_next()93 }9495 fn take_next(&mut self) -> Result<Option<MediaObjectReadOutcomeV3>, String> {96 let Some(outcome) = self.completed.remove(&self.next_accept_index) else {97 return Ok(None);98 };99 self.next_accept_index = self100 .next_accept_index101 .checked_add(1)102 .ok_or_else(|| "Media v3 accept index overflowed".to_string())?;103 Ok(Some(outcome))104 }105}106107pub(crate) struct MediaObjectReceiverV3 {108 connection: iroh::endpoint::Connection,109 reads: tokio::task::JoinSet<MediaObjectReadOutcomeV3>,110 reorder: MediaObjectReorderV3,111 next_accept_index: u64,112 connection_closed: bool,113}114115impl MediaObjectReceiverV3 {116 pub(crate) fn new(connection: iroh::endpoint::Connection) -> Self {117 Self {118 connection,119 reads: tokio::task::JoinSet::new(),120 reorder: MediaObjectReorderV3::new(1),121 next_accept_index: 0,122 connection_closed: false,123 }124 }125126 pub(crate) async fn next(&mut self) -> Result<Option<MediaObjectReadOutcomeV3>, String> {127 loop {128 if let Some(completed) = self.reorder.take_next()? {129 return Ok(Some(completed));130 }131 if self.connection_closed && self.reads.is_empty() {132 if self.reorder.pending_len() != 0 {133 return Err("Media v3 connection closed with incomplete object order".into());134 }135 return Ok(None);136 }137138 tokio::select! {139 biased;140 completed = self.reads.join_next(), if !self.reads.is_empty() => {141 let completed = completed142 .ok_or_else(|| "Media v3 object reader ended unexpectedly".to_string())?143 .map_err(|error| format!("Media v3 object reader task failed: {error}"))?;144 if let Some(completed) = self.reorder.push(completed)? {145 return Ok(Some(completed));146 }147 }148 accepted = self.connection.accept_uni(), if !self.connection_closed149 && self.reads.len() + self.reorder.pending_len()150 < CLIENT_MEDIA_OBJECT_CAPACITY => {151 let mut stream = match accepted {152 Ok(stream) => stream,153 Err(_) => {154 self.connection_closed = true;155 continue;156 }157 };158 self.next_accept_index = self.next_accept_index.checked_add(1)159 .ok_or_else(|| "Media v3 accept index overflowed".to_string())?;160 let accept_index = self.next_accept_index;161 self.reads.spawn(async move {162 match tokio::time::timeout(163 CLIENT_MEDIA_OBJECT_READ_TIMEOUT,164 read_media_object_v3(&mut stream),165 )166 .await167 {168 Err(_) => MediaObjectReadOutcomeV3::Dropped {169 accept_index,170 reason: KeyframeRequestReasonV3::DeliveryTimeout,171 },172 Ok(Err(ProtocolError::Io(_))) => MediaObjectReadOutcomeV3::Dropped {173 accept_index,174 reason: KeyframeRequestReasonV3::TransportGap,175 },176 Ok(Err(error)) => MediaObjectReadOutcomeV3::Malformed(format!(177 "Invalid media v3 object: {error}"178 )),179 Ok(Ok(object)) => MediaObjectReadOutcomeV3::Object {180 accept_index,181 object,182 },183 }184 });185 }186 }187 }188 }189}190191impl Drop for MediaObjectReceiverV3 {192 fn drop(&mut self) {193 self.reads.abort_all();194 }195}196197198pub(crate) enum MediaObjectSequenceDecisionV3 {199 Deliver { discontinuity: bool },200 DropLate,201 DropUntilKeyframe,202}203204205pub(crate) struct MediaObjectSequenceV3 {206 group_id: Option<u64>,207 last_object_id: Option<u32>,208 last_sequence: Option<u64>,209 group_payload_bytes: usize,210 waiting_for_keyframe: bool,211}212213impl MediaObjectSequenceV3 {214 pub(crate) fn new() -> Self {215 Self {216 waiting_for_keyframe: true,217 ..Self::default()218 }219 }220221 pub(crate) fn last_sequence(&self) -> Option<u64> {222 self.last_sequence223 }224225 pub(crate) fn waiting_for_keyframe(&self) -> bool {226 self.waiting_for_keyframe227 }228229 pub(crate) fn note_dropped_object(&mut self) -> bool {230 let entered = !self.waiting_for_keyframe;231 self.waiting_for_keyframe = true;232 entered233 }234235 pub(crate) fn classify(&mut self, object: &MediaObjectV3) -> MediaObjectSequenceDecisionV3 {236 let header = &object.header;237 if self238 .group_id239 .is_some_and(|group_id| header.group_id < group_id)240 || self241 .last_sequence242 .is_some_and(|sequence| header.sequence <= sequence)243 {244 return MediaObjectSequenceDecisionV3::DropLate;245 }246247 let new_group = self.group_id != Some(header.group_id);248 let recovery_keyframe = header.object_id == 0249 && header.flags.contains(FrameFlags::KEYFRAME)250 && header.flags.contains(FrameFlags::CODEC_CONFIG);251 if new_group && !recovery_keyframe {252 self.waiting_for_keyframe = true;253 return MediaObjectSequenceDecisionV3::DropUntilKeyframe;254 }255 if !new_group && self.waiting_for_keyframe {256 return MediaObjectSequenceDecisionV3::DropUntilKeyframe;257 }258259 let sequence_contiguous = self260 .last_sequence261 .is_none_or(|sequence| sequence.checked_add(1) == Some(header.sequence));262 let object_contiguous = new_group263 || self264 .last_object_id265 .is_some_and(|object_id| object_id.checked_add(1) == Some(header.object_id));266 let next_group_bytes = if new_group {267 object.payload.len()268 } else {269 self.group_payload_bytes270 .saturating_add(object.payload.len())271 };272 if (!sequence_contiguous && !new_group)273 || !object_contiguous274 || next_group_bytes > MAX_MEDIA_GROUP_BYTES_V3275 {276 self.waiting_for_keyframe = true;277 return MediaObjectSequenceDecisionV3::DropUntilKeyframe;278 }279280 let discontinuity = header.flags.contains(FrameFlags::DISCONTINUITY)281 || self.waiting_for_keyframe282 || !sequence_contiguous;283 self.group_id = Some(header.group_id);284 self.last_object_id = Some(header.object_id);285 self.last_sequence = Some(header.sequence);286 self.group_payload_bytes = next_group_bytes;287 self.waiting_for_keyframe = false;288 MediaObjectSequenceDecisionV3::Deliver { discontinuity }289 }290}291292293pub(crate) fn media_object_frame(sequence: u64, flags: FrameFlags) -> MediaFrame {294 let payload = vec![sequence as u8];295 let header = sigil_protocol::MediaFrameHeader::h264(296 1280,297 800,298 payload.len(),299 sequence,300 sequence * 1_000,301 sequence as i64 * 1_000,302 flags,303 )304 .unwrap();305 MediaFrame::new(header, payload).unwrap()306}307308309mod tests {310 use super::*;311312 fn media_object_v3(313 group_id: u64,314 object_id: u32,315 sequence: u64,316 flags: FrameFlags,317 ) -> MediaObjectV3 {318 let payload = vec![sequence as u8];319 let header = sigil_protocol::MediaObjectHeaderV3::h264(320 1280,321 800,322 payload.len(),323 if object_id == 0 { 0 } else { 128 },324 flags,325 object_id,326 group_id,327 sequence,328 sequence * 1_000,329 sequence as i64 * 1_000,330 100,331 )332 .unwrap();333 MediaObjectV3::new(header, payload).unwrap()334 }335336 fn media_object_outcome_v3(337 accept_index: u64,338 group_id: u64,339 object_id: u32,340 sequence: u64,341 flags: FrameFlags,342 ) -> MediaObjectReadOutcomeV3 {343 MediaObjectReadOutcomeV3::Object {344 accept_index,345 object: media_object_v3(group_id, object_id, sequence, flags),346 }347 }348349 350 fn media_v3_groups_use_wire_object_identity_and_recover_on_new_group_zero() {351 let keyframe = FrameFlags::KEYFRAME.union(FrameFlags::CODEC_CONFIG);352 let mut sequence = MediaObjectSequenceV3::new();353354 assert_eq!(355 sequence.classify(&media_object_v3(10, 0, 10, keyframe)),356 MediaObjectSequenceDecisionV3::Deliver {357 discontinuity: true358 }359 );360 assert_eq!(361 sequence.classify(&media_object_v3(10, 1, 11, FrameFlags::NONE)),362 MediaObjectSequenceDecisionV3::Deliver {363 discontinuity: false364 }365 );366 assert_eq!(367 sequence.classify(&media_object_v3(10, 3, 13, FrameFlags::NONE)),368 MediaObjectSequenceDecisionV3::DropUntilKeyframe369 );370 assert_eq!(371 sequence.classify(&media_object_v3(10, 2, 12, FrameFlags::NONE)),372 MediaObjectSequenceDecisionV3::DropUntilKeyframe373 );374 assert_eq!(375 sequence.classify(&media_object_v3(376 20,377 0,378 20,379 keyframe.union(FrameFlags::DISCONTINUITY),380 )),381 MediaObjectSequenceDecisionV3::Deliver {382 discontinuity: true383 }384 );385 assert_eq!(386 sequence.classify(&media_object_v3(10, 2, 12, FrameFlags::NONE)),387 MediaObjectSequenceDecisionV3::DropLate388 );389 }390391 392 fn late_dropped_completion_after_v3_barrier_cannot_poison_recovered_group() {393 let barrier = FrameFlags::KEYFRAME394 .union(FrameFlags::CODEC_CONFIG)395 .union(FrameFlags::DISCONTINUITY);396 let mut reorder = MediaObjectReorderV3::new(1);397 let mut sequence = MediaObjectSequenceV3::new();398399 let recovered = reorder400 .push(media_object_outcome_v3(5, 20, 0, 20, barrier))401 .unwrap()402 .unwrap();403 let MediaObjectReadOutcomeV3::Object { object, .. } = recovered else {404 panic!("recovery barrier must remain an object");405 };406 assert_eq!(407 sequence.classify(&object),408 MediaObjectSequenceDecisionV3::Deliver {409 discontinuity: true410 }411 );412413 assert!(414 reorder415 .push(MediaObjectReadOutcomeV3::Dropped {416 accept_index: 1,417 reason: KeyframeRequestReasonV3::DeliveryTimeout,418 })419 .unwrap()420 .is_none()421 );422423 let delta = reorder424 .push(media_object_outcome_v3(6, 20, 1, 21, FrameFlags::NONE))425 .unwrap()426 .unwrap();427 let MediaObjectReadOutcomeV3::Object { object, .. } = delta else {428 panic!("new group delta must remain an object");429 };430 assert_eq!(431 sequence.classify(&object),432 MediaObjectSequenceDecisionV3::Deliver {433 discontinuity: false434 }435 );436 }437438 439 fn media_v3_receiver_rejects_group_payload_growth_beyond_the_shared_cap() {440 let keyframe = FrameFlags::KEYFRAME.union(FrameFlags::CODEC_CONFIG);441 let mut sequence = MediaObjectSequenceV3::new();442 assert!(matches!(443 sequence.classify(&media_object_v3(1, 0, 1, keyframe)),444 MediaObjectSequenceDecisionV3::Deliver { .. }445 ));446 sequence.group_payload_bytes = MAX_MEDIA_GROUP_BYTES_V3;447 assert_eq!(448 sequence.classify(&media_object_v3(1, 1, 2, FrameFlags::NONE)),449 MediaObjectSequenceDecisionV3::DropUntilKeyframe450 );451 }452}