1use alloc::collections::BTreeMap;
6use alloc::collections::VecDeque;
7use alloc::vec;
8use alloc::vec::Vec;
9
10use super::messages::{Dii, DownloadDataBlock};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14#[cfg_attr(feature = "serde", derive(serde::Serialize))]
15pub struct ModuleKey {
16 pub download_id: u32,
18 pub module_id: u16,
20 pub module_version: u8,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26#[cfg_attr(feature = "serde", derive(serde::Serialize))]
27pub struct Module {
28 pub key: ModuleKey,
30 pub data: Vec<u8>,
32}
33
34type SlotKey = (u32, u16); struct Slot {
43 module_version: u8,
44 block_size: usize,
45 data: Vec<u8>,
46 received: Vec<u64>,
47 n_blocks: usize,
48 remaining: usize,
49}
50
51impl Slot {
52 fn is_received(&self, n: usize) -> bool {
53 (self.received[n >> 6] >> (n & 63)) & 1 != 0
54 }
55 fn mark_received(&mut self, n: usize) {
56 self.received[n >> 6] |= 1 << (n & 63);
57 }
58}
59
60pub const DEFAULT_MAX_MODULE_SIZE: u32 = 64 * 1024 * 1024;
62pub const DEFAULT_MAX_TOTAL_BYTES: usize = 256 * 1024 * 1024;
66pub const DEFAULT_MAX_SLOTS: usize = 16 * 1024;
73
74pub struct ModuleReassembler {
98 slots: BTreeMap<SlotKey, Slot>,
99 completed: BTreeMap<SlotKey, u8>,
105 completed_order: VecDeque<SlotKey>,
106 max_module_size: u32,
107 max_total_bytes: usize,
108 max_slots: usize,
109 total_bytes: usize,
110}
111
112impl Default for ModuleReassembler {
113 fn default() -> Self {
114 Self::new()
115 }
116}
117
118impl ModuleReassembler {
119 #[must_use]
122 pub fn new() -> Self {
123 Self::with_limits(DEFAULT_MAX_MODULE_SIZE, DEFAULT_MAX_TOTAL_BYTES)
124 }
125
126 #[must_use]
129 pub fn with_max_module_size(max_module_size: u32) -> Self {
130 Self::with_limits(max_module_size, DEFAULT_MAX_TOTAL_BYTES)
131 }
132
133 #[must_use]
136 pub fn with_limits(max_module_size: u32, max_total_bytes: usize) -> Self {
137 Self {
138 slots: BTreeMap::new(),
139 completed: BTreeMap::new(),
140 completed_order: VecDeque::new(),
141 max_module_size,
142 max_total_bytes,
143 max_slots: DEFAULT_MAX_SLOTS,
144 total_bytes: 0,
145 }
146 }
147
148 #[must_use]
153 pub fn with_max_slots(mut self, max_slots: usize) -> Self {
154 self.max_slots = max_slots;
155 self
156 }
157
158 pub fn note_dii(&mut self, dii: &Dii<'_>) {
165 for m in &dii.modules {
166 if m.module_size == 0 || m.module_size > self.max_module_size || dii.block_size == 0 {
167 continue;
168 }
169 let key: SlotKey = (dii.download_id, m.module_id);
170 if self.completed.get(&key) == Some(&m.module_version) {
171 continue;
175 }
176 self.completed.remove(&key);
179 if let Some(existing) = self.slots.get(&key) {
180 if existing.module_version == m.module_version {
181 continue; }
183 let s = self.slots.remove(&key).expect("just found");
185 self.total_bytes -= s.data.len();
186 }
187 let size = m.module_size as usize;
188 if self.total_bytes + size > self.max_total_bytes || self.slots.len() >= self.max_slots
189 {
190 continue; }
192 let block_size = dii.block_size as usize;
193 let n_blocks = size.div_ceil(block_size).max(1);
194 self.total_bytes += size;
195 self.slots.insert(
196 key,
197 Slot {
198 module_version: m.module_version,
199 block_size,
200 data: vec![0u8; size],
201 received: vec![0u64; n_blocks.div_ceil(64)],
202 n_blocks,
203 remaining: n_blocks,
204 },
205 );
206 }
207 }
208
209 pub fn feed_ddb(&mut self, ddb: &DownloadDataBlock<'_>) -> Option<Module> {
214 let key: SlotKey = (ddb.download_id, ddb.module_id);
215 let slot = self.slots.get_mut(&key)?;
216 let n = ddb.block_number as usize;
217 if slot.module_version != ddb.module_version || n >= slot.n_blocks || slot.is_received(n) {
218 return None;
219 }
220 let offset = n * slot.block_size;
221 let expected = (slot.data.len() - offset).min(slot.block_size);
222 if ddb.block_data.len() != expected {
223 return None; }
225 slot.data[offset..offset + expected].copy_from_slice(ddb.block_data);
226 slot.mark_received(n);
227 slot.remaining -= 1;
228 if slot.remaining > 0 {
229 return None;
230 }
231 let slot = self.slots.remove(&key).expect("slot exists");
232 self.total_bytes -= slot.data.len();
233 self.record_completed(key, slot.module_version);
234 Some(Module {
235 key: ModuleKey {
236 download_id: ddb.download_id,
237 module_id: ddb.module_id,
238 module_version: slot.module_version,
239 },
240 data: slot.data,
241 })
242 }
243
244 fn record_completed(&mut self, key: SlotKey, version: u8) {
250 if !self.completed.contains_key(&key)
251 && self.completed.len() >= self.max_slots
252 && let Some(oldest) = self.completed_order.pop_front()
253 {
254 self.completed.remove(&oldest);
255 }
256 self.completed.insert(key, version);
257 self.completed_order.push_back(key);
258 }
259
260 pub fn clear(&mut self) {
262 self.slots.clear();
263 self.completed.clear();
264 self.completed_order.clear();
265 self.total_bytes = 0;
266 }
267
268 pub fn retain<F>(&mut self, mut keep: F)
278 where
279 F: FnMut(u32, u16) -> bool,
280 {
281 let mut freed = 0usize;
282 self.slots.retain(|&(download_id, module_id), slot| {
283 let keep = keep(download_id, module_id);
284 if !keep {
285 freed += slot.data.len();
286 }
287 keep
288 });
289 self.total_bytes -= freed;
290 self.completed
291 .retain(|&(download_id, module_id), _| keep(download_id, module_id));
292 let completed = &self.completed;
293 self.completed_order
294 .retain(|key| completed.contains_key(key));
295 }
296
297 #[must_use]
299 pub fn pending(&self) -> usize {
300 self.slots.len()
301 }
302
303 #[must_use]
312 pub fn pending_bytes(&self) -> usize {
313 self.total_bytes
314 }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::super::messages::DiiModule;
320 use super::*;
321 use crate::compatibility::CompatibilityDescriptor;
322
323 fn dii(download_id: u32, block_size: u16, modules: Vec<DiiModule<'static>>) -> Dii<'static> {
324 Dii {
325 transaction_id: 0x8000_0002,
326 adaptation: &[],
327 download_id,
328 block_size,
329 window_size: 0,
330 ack_period: 0,
331 t_c_download_window: 0,
332 t_c_download_scenario: 0,
333 compatibility_descriptor: CompatibilityDescriptor {
334 empty_count_present: false,
335 descriptors: vec![],
336 },
337 modules,
338 private_data: &[],
339 }
340 }
341
342 fn module(module_id: u16, module_size: u32, module_version: u8) -> DiiModule<'static> {
343 DiiModule {
344 module_id,
345 module_size,
346 module_version,
347 module_info: &[],
348 }
349 }
350
351 fn ddb(
352 download_id: u32,
353 module_id: u16,
354 module_version: u8,
355 block_number: u16,
356 block_data: &[u8],
357 ) -> DownloadDataBlock<'_> {
358 DownloadDataBlock {
359 download_id,
360 adaptation: &[],
361 module_id,
362 module_version,
363 block_number,
364 block_data,
365 }
366 }
367
368 #[test]
369 fn two_block_module_completes() {
370 let mut r = ModuleReassembler::new();
371 r.note_dii(&dii(1, 4, vec![module(7, 6, 0)]));
372 assert!(r.feed_ddb(&ddb(1, 7, 0, 0, &[1, 2, 3, 4])).is_none());
373 let m = r.feed_ddb(&ddb(1, 7, 0, 1, &[5, 6])).expect("complete");
374 assert_eq!(m.key.module_id, 7);
375 assert_eq!(m.data, vec![1, 2, 3, 4, 5, 6]);
376 assert_eq!(r.pending(), 0);
377 }
378
379 #[test]
380 fn out_of_order_blocks_complete() {
381 let mut r = ModuleReassembler::new();
382 r.note_dii(&dii(1, 2, vec![module(1, 4, 0)]));
383 assert!(r.feed_ddb(&ddb(1, 1, 0, 1, &[3, 4])).is_none());
384 let m = r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).expect("complete");
385 assert_eq!(m.data, vec![1, 2, 3, 4]);
386 }
387
388 #[test]
389 fn ddb_before_dii_is_ignored() {
390 let mut r = ModuleReassembler::new();
391 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none());
392 r.note_dii(&dii(1, 2, vec![module(1, 2, 0)]));
394 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_some());
395 }
396
397 #[test]
398 fn version_mismatch_ignored_and_new_version_restarts() {
399 let mut r = ModuleReassembler::new();
400 r.note_dii(&dii(1, 2, vec![module(1, 4, 0)]));
401 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none());
402 assert!(r.feed_ddb(&ddb(1, 1, 3, 1, &[9, 9])).is_none());
404 r.note_dii(&dii(1, 2, vec![module(1, 4, 3)]));
406 assert_eq!(r.pending(), 1);
407 assert!(r.feed_ddb(&ddb(1, 1, 3, 0, &[5, 6])).is_none());
408 let m = r.feed_ddb(&ddb(1, 1, 3, 1, &[7, 8])).expect("complete");
409 assert_eq!(m.key.module_version, 3);
410 assert_eq!(m.data, vec![5, 6, 7, 8]);
411 }
412
413 #[test]
414 fn repeated_dii_keeps_progress() {
415 let mut r = ModuleReassembler::new();
416 let d = dii(1, 2, vec![module(1, 4, 0)]);
417 r.note_dii(&d);
418 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none());
419 r.note_dii(&d); let m = r.feed_ddb(&ddb(1, 1, 0, 1, &[3, 4])).expect("complete");
421 assert_eq!(m.data, vec![1, 2, 3, 4]);
422 }
423
424 #[test]
428 fn completed_module_is_not_re_emitted_on_dii_repeat() {
429 let mut r = ModuleReassembler::new();
430 let d = dii(1, 2, vec![module(1, 4, 0)]);
431 r.note_dii(&d);
432 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none());
433 let m = r.feed_ddb(&ddb(1, 1, 0, 1, &[3, 4])).expect("complete");
434 assert_eq!(m.data, vec![1, 2, 3, 4]);
435 assert_eq!(r.pending(), 0);
436
437 r.note_dii(&d);
439 assert_eq!(
440 r.pending(),
441 0,
442 "a same-version DII repeat must not recreate a completed module's slot"
443 );
444 assert!(
445 r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none(),
446 "a completed module's DDB repeats must not be re-collected"
447 );
448 assert!(r.feed_ddb(&ddb(1, 1, 0, 1, &[3, 4])).is_none());
449
450 r.note_dii(&dii(1, 2, vec![module(1, 4, 1)]));
452 assert_eq!(r.pending(), 1);
453 assert!(r.feed_ddb(&ddb(1, 1, 1, 0, &[9, 9])).is_none());
454 let m2 = r.feed_ddb(&ddb(1, 1, 1, 1, &[8, 8])).expect("complete");
455 assert_eq!(m2.data, vec![9, 9, 8, 8]);
456 assert_eq!(m2.key.module_version, 1);
457 }
458
459 #[test]
462 fn retain_drops_slots_and_completed_records_for_withdrawn_modules() {
463 let mut r = ModuleReassembler::new();
464 r.note_dii(&dii(1, 2, vec![module(1, 4, 0), module(2, 2, 0)]));
465 let m = r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2]));
466 assert!(m.is_none());
467 let m = r
468 .feed_ddb(&ddb(1, 1, 0, 1, &[3, 4]))
469 .expect("module 1 complete");
470 assert_eq!(m.data, vec![1, 2, 3, 4]);
471 assert_eq!(r.pending(), 1); r.retain(|_, module_id| module_id == 1);
475 assert_eq!(
476 r.pending(),
477 0,
478 "module 2's in-progress slot must be dropped"
479 );
480 assert_eq!(r.pending_bytes(), 0);
481
482 r.note_dii(&dii(1, 2, vec![module(1, 4, 0)]));
485 assert_eq!(r.pending(), 0);
486 }
487
488 #[test]
489 fn duplicate_and_out_of_range_blocks_ignored() {
490 let mut r = ModuleReassembler::new();
491 r.note_dii(&dii(1, 2, vec![module(1, 4, 0)]));
492 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none());
493 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2])).is_none()); assert!(r.feed_ddb(&ddb(1, 1, 0, 9, &[9, 9])).is_none()); assert_eq!(r.pending(), 1);
496 }
497
498 #[test]
499 fn wrong_block_length_ignored() {
500 let mut r = ModuleReassembler::new();
501 r.note_dii(&dii(1, 4, vec![module(1, 6, 0)]));
502 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2, 3])).is_none());
504 assert!(r.feed_ddb(&ddb(1, 1, 0, 1, &[5, 6, 7])).is_none());
505 assert_eq!(r.pending(), 1);
506 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[1, 2, 3, 4])).is_none());
507 assert!(r.feed_ddb(&ddb(1, 1, 0, 1, &[5, 6])).is_some());
508 }
509
510 #[test]
511 fn oversize_module_skipped() {
512 let mut r = ModuleReassembler::with_max_module_size(8);
513 r.note_dii(&dii(1, 4, vec![module(1, 9, 0), module(2, 8, 0)]));
514 assert_eq!(r.pending(), 1); }
516
517 #[test]
518 fn zero_block_size_skipped() {
519 let mut r = ModuleReassembler::new();
520 r.note_dii(&dii(1, 0, vec![module(1, 4, 0)]));
521 assert_eq!(r.pending(), 0);
522 }
523
524 #[test]
528 fn zero_size_module_announcement_ignored() {
529 let mut r = ModuleReassembler::new();
530 r.note_dii(&dii(1, 4, vec![module(1, 0, 0)]));
531 assert_eq!(r.pending(), 0);
532 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[])).is_none());
533 }
534
535 #[test]
540 fn slot_count_capped() {
541 let mut r = ModuleReassembler::new().with_max_slots(3);
542 let modules: Vec<_> = (0..5).map(|i| module(i, 1, 0)).collect();
543 r.note_dii(&dii(1, 4, modules));
544 assert_eq!(r.pending(), 3); assert!(r.feed_ddb(&ddb(1, 0, 0, 0, &[0xAA])).is_some());
547 assert_eq!(r.pending(), 2);
548 r.note_dii(&dii(1, 4, vec![module(4, 1, 0)]));
550 assert_eq!(r.pending(), 3);
551 }
552
553 #[test]
557 fn aggregate_budget_bounds_total_memory() {
558 let mut r = ModuleReassembler::with_limits(8, 10);
559 r.note_dii(&dii(1, 4, vec![module(1, 8, 0)]));
560 assert_eq!(r.pending_bytes(), 8);
561 r.note_dii(&dii(2, 4, vec![module(1, 8, 0)]));
564 assert_eq!(r.pending(), 1);
565 assert_eq!(r.pending_bytes(), 8);
566 assert!(r.feed_ddb(&ddb(1, 1, 0, 0, &[0; 4])).is_none());
568 assert!(r.feed_ddb(&ddb(1, 1, 0, 1, &[0; 4])).is_some());
569 assert_eq!(r.pending_bytes(), 0);
570 r.note_dii(&dii(2, 4, vec![module(1, 8, 0)]));
572 assert_eq!(r.pending(), 1);
573 assert_eq!(r.pending_bytes(), 8);
574 }
575
576 #[test]
578 fn version_replacement_releases_budget() {
579 let mut r = ModuleReassembler::with_limits(8, 8);
580 r.note_dii(&dii(1, 4, vec![module(1, 8, 0)]));
581 assert_eq!(r.pending_bytes(), 8);
582 r.note_dii(&dii(1, 4, vec![module(1, 8, 1)]));
583 assert_eq!(r.pending(), 1);
584 assert_eq!(r.pending_bytes(), 8); }
586
587 #[test]
590 fn block_size_one_uses_bitset() {
591 let mut r = ModuleReassembler::new();
592 r.note_dii(&dii(1, 1, vec![module(1, 130, 0)]));
593 for i in 0..129u16 {
594 assert!(r.feed_ddb(&ddb(1, 1, 0, i, &[i as u8])).is_none());
595 assert!(r.feed_ddb(&ddb(1, 1, 0, i, &[i as u8])).is_none());
597 }
598 let m = r.feed_ddb(&ddb(1, 1, 0, 129, &[0x81])).expect("complete");
599 assert_eq!(m.data.len(), 130);
600 assert_eq!(m.data[129], 0x81);
601 }
602}