1use crate::constants::DEFAULT_CHUNK_SIZE;
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6use std::collections::hash_map::RandomState;
7use std::collections::{HashMap, HashSet, VecDeque};
8use std::hash::BuildHasher;
9use std::time::{Duration, Instant};
10
11mod base64_bytes {
12 use base64::{engine::general_purpose::STANDARD, Engine};
13 use serde::{Deserialize, Deserializer, Serializer};
14
15 pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
16 where
17 S: Serializer,
18 {
19 serializer.serialize_str(&STANDARD.encode(bytes))
20 }
21
22 pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
23 where
24 D: Deserializer<'de>,
25 {
26 let s = String::deserialize(deserializer)?;
27 STANDARD.decode(&s).map_err(serde::de::Error::custom)
28 }
29}
30
31#[derive(Debug, Clone)]
33pub struct FileTransferConfig {
34 pub chunk_size: usize,
41
42 pub max_file_size: u64,
44
45 pub max_concurrent_assemblies: usize,
49
50 pub max_assemblies_per_sender: usize,
57
58 pub max_total_buffered_bytes: u64,
67}
68
69impl Default for FileTransferConfig {
70 fn default() -> Self {
71 Self {
72 chunk_size: DEFAULT_CHUNK_SIZE,
73 max_file_size: 100 * 1024 * 1024, max_concurrent_assemblies: 32,
75 max_assemblies_per_sender: 16,
76 max_total_buffered_bytes: 128 * 1024 * 1024, }
78 }
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum ChunkRejection {
84 Invalid,
88 MetadataMismatch,
90 SenderMismatch,
92 TooManyTransfers,
94 SenderQuotaExceeded,
96 BufferBudgetExhausted,
101 PreviouslyFailed,
108}
109
110impl ChunkRejection {
111 pub fn is_resource_exhaustion(self) -> bool {
115 matches!(
116 self,
117 Self::TooManyTransfers | Self::SenderQuotaExceeded | Self::BufferBudgetExhausted
118 )
119 }
120
121 pub fn as_str(self) -> &'static str {
124 match self {
125 Self::Invalid => "invalid_chunk",
126 Self::MetadataMismatch => "metadata_mismatch",
127 Self::SenderMismatch => "sender_mismatch",
128 Self::TooManyTransfers => "too_many_transfers",
129 Self::SenderQuotaExceeded => "sender_quota_exceeded",
130 Self::BufferBudgetExhausted => "buffer_budget_exhausted",
131 Self::PreviouslyFailed => "previously_failed",
132 }
133 }
134}
135
136#[non_exhaustive]
140#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
141pub enum ChunkDecodeError {
142 #[error("unexpected end of data reading {what}")]
144 UnexpectedEof {
145 what: &'static str,
147 },
148 #[error("{field} {len} exceeds maximum {max}")]
150 FieldTooLong {
151 field: &'static str,
153 len: usize,
155 max: usize,
157 },
158 #[error("invalid {field} UTF-8: {source}")]
160 InvalidUtf8 {
161 field: &'static str,
163 source: std::string::FromUtf8Error,
165 },
166}
167
168#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct FileChunk {
171 pub file_id: String,
173
174 pub file_name: String,
176
177 pub file_size: u64,
179
180 pub total_chunks: u32,
182
183 pub chunk_index: u32,
185
186 #[serde(with = "base64_bytes")]
188 pub chunk_data: Vec<u8>,
189
190 pub file_checksum: String,
192}
193
194impl FileChunk {
195 pub fn to_json(&self) -> Result<String, serde_json::Error> {
197 serde_json::to_string(self)
198 }
199
200 pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
202 serde_json::from_str(json)
203 }
204
205 pub fn to_bytes(&self) -> Vec<u8> {
215 let file_id = self.file_id.as_bytes();
216 let file_name = self.file_name.as_bytes();
217 let checksum = self.file_checksum.as_bytes();
218
219 let capacity = 4
220 + file_id.len()
221 + 4
222 + file_name.len()
223 + 8
224 + 4
225 + 4
226 + 4
227 + self.chunk_data.len()
228 + 4
229 + checksum.len();
230 let mut buf = Vec::with_capacity(capacity);
231
232 buf.extend_from_slice(&(file_id.len() as u32).to_le_bytes());
233 buf.extend_from_slice(file_id);
234 buf.extend_from_slice(&(file_name.len() as u32).to_le_bytes());
235 buf.extend_from_slice(file_name);
236 buf.extend_from_slice(&self.file_size.to_le_bytes());
237 buf.extend_from_slice(&self.total_chunks.to_le_bytes());
238 buf.extend_from_slice(&self.chunk_index.to_le_bytes());
239 buf.extend_from_slice(&(self.chunk_data.len() as u32).to_le_bytes());
240 buf.extend_from_slice(&self.chunk_data);
241 buf.extend_from_slice(&(checksum.len() as u32).to_le_bytes());
242 buf.extend_from_slice(checksum);
243
244 buf
245 }
246
247 pub(crate) const MAX_STRING_FIELD_LEN: usize = 4096;
252 const MAX_CHUNK_DATA_LEN: usize = 128 * 1024 * 1024; pub fn from_bytes(data: &[u8]) -> Result<Self, ChunkDecodeError> {
257 let mut pos = 0;
258
259 let read_u32 = |pos: &mut usize| -> Result<u32, ChunkDecodeError> {
260 if *pos + 4 > data.len() {
261 return Err(ChunkDecodeError::UnexpectedEof { what: "u32" });
262 }
263 let mut bytes = [0u8; 4];
264 bytes.copy_from_slice(&data[*pos..*pos + 4]);
265 let val = u32::from_le_bytes(bytes);
266 *pos += 4;
267 Ok(val)
268 };
269
270 let read_u64 = |pos: &mut usize| -> Result<u64, ChunkDecodeError> {
271 if *pos + 8 > data.len() {
272 return Err(ChunkDecodeError::UnexpectedEof { what: "u64" });
273 }
274 let mut bytes = [0u8; 8];
275 bytes.copy_from_slice(&data[*pos..*pos + 8]);
276 let val = u64::from_le_bytes(bytes);
277 *pos += 8;
278 Ok(val)
279 };
280
281 let read_bytes = |pos: &mut usize, len: usize| -> Result<Vec<u8>, ChunkDecodeError> {
282 if *pos + len > data.len() {
283 return Err(ChunkDecodeError::UnexpectedEof { what: "bytes" });
284 }
285 let val = data[*pos..*pos + len].to_vec();
286 *pos += len;
287 Ok(val)
288 };
289
290 let file_id_len = read_u32(&mut pos)? as usize;
291 if file_id_len > Self::MAX_STRING_FIELD_LEN {
292 return Err(ChunkDecodeError::FieldTooLong {
293 field: "file_id_len",
294 len: file_id_len,
295 max: Self::MAX_STRING_FIELD_LEN,
296 });
297 }
298 let file_id = String::from_utf8(read_bytes(&mut pos, file_id_len)?).map_err(|e| {
299 ChunkDecodeError::InvalidUtf8 {
300 field: "file_id",
301 source: e,
302 }
303 })?;
304
305 let file_name_len = read_u32(&mut pos)? as usize;
306 if file_name_len > Self::MAX_STRING_FIELD_LEN {
307 return Err(ChunkDecodeError::FieldTooLong {
308 field: "file_name_len",
309 len: file_name_len,
310 max: Self::MAX_STRING_FIELD_LEN,
311 });
312 }
313 let file_name = String::from_utf8(read_bytes(&mut pos, file_name_len)?).map_err(|e| {
314 ChunkDecodeError::InvalidUtf8 {
315 field: "file_name",
316 source: e,
317 }
318 })?;
319
320 let file_size = read_u64(&mut pos)?;
321 let total_chunks = read_u32(&mut pos)?;
322 let chunk_index = read_u32(&mut pos)?;
323
324 let chunk_data_len = read_u32(&mut pos)? as usize;
325 if chunk_data_len > Self::MAX_CHUNK_DATA_LEN {
326 return Err(ChunkDecodeError::FieldTooLong {
327 field: "chunk_data_len",
328 len: chunk_data_len,
329 max: Self::MAX_CHUNK_DATA_LEN,
330 });
331 }
332 let chunk_data = read_bytes(&mut pos, chunk_data_len)?;
333
334 let checksum_len = read_u32(&mut pos)? as usize;
335 if checksum_len > Self::MAX_STRING_FIELD_LEN {
336 return Err(ChunkDecodeError::FieldTooLong {
337 field: "checksum_len",
338 len: checksum_len,
339 max: Self::MAX_STRING_FIELD_LEN,
340 });
341 }
342 let file_checksum =
343 String::from_utf8(read_bytes(&mut pos, checksum_len)?).map_err(|e| {
344 ChunkDecodeError::InvalidUtf8 {
345 field: "checksum",
346 source: e,
347 }
348 })?;
349
350 Ok(Self {
351 file_id,
352 file_name,
353 file_size,
354 total_chunks,
355 chunk_index,
356 chunk_data,
357 file_checksum,
358 })
359 }
360}
361
362#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct FileProgress {
365 pub file_id: String,
367
368 pub file_name: String,
370
371 pub file_size: u64,
373
374 pub chunks_completed: u32,
376
377 pub total_chunks: u32,
379
380 pub percentage: u8,
382}
383
384impl FileProgress {
385 pub fn calculate_percentage(chunks_completed: u32, total_chunks: u32) -> u8 {
387 if total_chunks == 0 {
388 return 0;
389 }
390 ((chunks_completed as f32 / total_chunks as f32) * 100.0) as u8
391 }
392}
393
394#[derive(Debug, Clone, PartialEq, Eq)]
397pub struct StaleTransfer {
398 pub file_id: String,
400 pub file_name: String,
402 pub sender: String,
404}
405
406struct FileAssembly {
408 sender: String,
411 file_name: String,
412 file_size: u64,
413 total_chunks: u32,
414 file_checksum: String,
415 received_chunks: HashMap<u32, Vec<u8>>,
416 received_bytes: u64,
418 last_updated_at: Instant,
419}
420
421impl FileAssembly {
422 fn new(sender: &str, chunk: &FileChunk) -> Self {
423 Self {
424 sender: sender.to_string(),
425 file_name: chunk.file_name.clone(),
426 file_size: chunk.file_size,
427 total_chunks: chunk.total_chunks,
428 file_checksum: chunk.file_checksum.clone(),
429 received_chunks: HashMap::new(),
430 received_bytes: 0,
431 last_updated_at: Instant::now(),
432 }
433 }
434
435 fn add_chunk(&mut self, chunk_index: u32, data: Vec<u8>) {
436 if let Some(old) = self.received_chunks.get(&chunk_index) {
437 self.received_bytes -= old.len() as u64;
438 }
439 self.received_bytes += data.len() as u64;
440 self.received_chunks.insert(chunk_index, data);
441 self.last_updated_at = Instant::now();
442 }
443
444 fn charged_bytes(&self) -> u64 {
448 self.received_bytes
449 + self.received_chunks.len() as u64 * FileTransferManager::PER_CHUNK_OVERHEAD
450 }
451
452 fn is_complete(&self) -> bool {
453 if self.total_chunks == 0 || self.received_chunks.len() != self.total_chunks as usize {
454 return false;
455 }
456
457 (0..self.total_chunks).all(|chunk_index| self.received_chunks.contains_key(&chunk_index))
458 }
459
460 fn progress(&self) -> FileProgress {
461 FileProgress {
462 file_id: String::new(), file_name: self.file_name.clone(),
464 file_size: self.file_size,
465 chunks_completed: self.received_chunks.len() as u32,
466 total_chunks: self.total_chunks,
467 percentage: FileProgress::calculate_percentage(
468 self.received_chunks.len() as u32,
469 self.total_chunks,
470 ),
471 }
472 }
473
474 fn reassemble(&self) -> Option<Vec<u8>> {
475 if !self.is_complete() {
476 return None;
477 }
478
479 let mut file_data = Vec::with_capacity(self.received_bytes as usize);
482
483 for i in 0..self.total_chunks {
485 let chunk_data = self.received_chunks.get(&i)?; file_data.extend_from_slice(chunk_data);
487 }
488
489 Some(file_data)
490 }
491}
492
493pub struct FileTransferManager {
495 config: FileTransferConfig,
496 active_assemblies: HashMap<String, FileAssembly>,
497 failed_transfers: HashMap<u64, Instant>,
504 failed_transfer_order: VecDeque<u64>,
507 tombstone_hasher: RandomState,
511}
512
513impl FileTransferManager {
514 const MIN_ACCEPTED_CHUNK_PAYLOAD: u64 = 1024;
519
520 const MIN_TOTAL_CHUNKS_CAP: u64 = 1024;
523
524 const PER_CHUNK_OVERHEAD: u64 = 64;
529
530 pub const MANUAL_SENDER: &'static str = "manual:ffi";
536
537 const TOMBSTONE_CAPACITY: usize = 1024;
543
544 pub fn new() -> Self {
546 Self::with_config(FileTransferConfig::default())
547 }
548
549 pub fn with_config(mut config: FileTransferConfig) -> Self {
555 let min_budget = Self::min_viable_budget(&config);
556 if config.max_total_buffered_bytes < min_budget {
557 tracing::warn!(
558 configured = config.max_total_buffered_bytes,
559 clamped_to = min_budget,
560 "max_total_buffered_bytes below what one max_file_size transfer charges; clamping"
561 );
562 config.max_total_buffered_bytes = min_budget;
563 }
564 Self {
565 config,
566 active_assemblies: HashMap::new(),
567 failed_transfers: HashMap::new(),
568 failed_transfer_order: VecDeque::new(),
569 tombstone_hasher: RandomState::new(),
570 }
571 }
572
573 fn tombstone_key(&self, file_id: &str) -> u64 {
576 self.tombstone_hasher.hash_one(file_id)
577 }
578
579 fn tombstone(&mut self, file_id: &str) {
583 let key = self.tombstone_key(file_id);
584 if self.failed_transfers.insert(key, Instant::now()).is_none() {
585 self.failed_transfer_order.push_back(key);
586 while self.failed_transfers.len() > Self::TOMBSTONE_CAPACITY {
587 match self.failed_transfer_order.pop_front() {
588 Some(oldest) => {
589 self.failed_transfers.remove(&oldest);
590 }
591 None => break,
592 }
593 }
594 }
595 }
596
597 fn is_tombstoned(&mut self, file_id: &str) -> bool {
601 let key = self.tombstone_key(file_id);
602 match self.failed_transfers.get_mut(&key) {
603 Some(last_seen) => {
604 *last_seen = Instant::now();
605 true
606 }
607 None => false,
608 }
609 }
610
611 fn derived_total_chunks_cap(max_file_size: u64) -> u64 {
615 (max_file_size / Self::MIN_ACCEPTED_CHUNK_PAYLOAD)
616 .saturating_add(1)
617 .max(Self::MIN_TOTAL_CHUNKS_CAP)
618 }
619
620 fn min_viable_budget(config: &FileTransferConfig) -> u64 {
625 let max_chunks =
626 Self::derived_total_chunks_cap(config.max_file_size).min(config.max_file_size);
627 config
628 .max_file_size
629 .saturating_add(max_chunks.saturating_mul(Self::PER_CHUNK_OVERHEAD))
630 }
631
632 pub fn chunk_file(
651 &self,
652 file_id: String,
653 file_name: String,
654 file_data: Vec<u8>,
655 chunk_size_override: Option<usize>,
656 ) -> crate::Result<Vec<FileChunk>> {
657 let file_size = file_data.len() as u64;
658
659 if file_size > self.config.max_file_size {
660 return Err(crate::Error::InvalidArgument(format!(
661 "File size {} exceeds maximum {}",
662 file_size, self.config.max_file_size
663 )));
664 }
665
666 if file_data.is_empty() {
667 return Err(crate::Error::InvalidArgument("File is empty".to_string()));
668 }
669
670 let chunk_size = chunk_size_override.unwrap_or(self.config.chunk_size);
671 if chunk_size == 0 {
672 return Err(crate::Error::InvalidArgument(
673 "Chunk size must be greater than zero".to_string(),
674 ));
675 }
676
677 let total_chunks = file_size.div_ceil(chunk_size as u64);
681 let max_total_chunks = Self::derived_total_chunks_cap(self.config.max_file_size);
682 if total_chunks > max_total_chunks {
683 return Err(crate::Error::InvalidArgument(format!(
684 "Chunk size {} splits {} bytes into {} chunks, exceeding the receive-side cap of {}; use a larger chunk size",
685 chunk_size, file_size, total_chunks, max_total_chunks
686 )));
687 }
688 let total_chunks = total_chunks as u32;
689
690 let file_checksum = format!("{:x}", Sha256::digest(&file_data));
691
692 let mut chunks = Vec::new();
693
694 for chunk_index in 0..total_chunks {
695 let start = (chunk_index as usize) * chunk_size;
696 let end = ((chunk_index as usize + 1) * chunk_size).min(file_data.len());
697
698 let chunk = FileChunk {
699 file_id: file_id.clone(),
700 file_name: file_name.clone(),
701 file_size,
702 total_chunks,
703 chunk_index,
704 chunk_data: file_data[start..end].to_vec(),
705 file_checksum: file_checksum.clone(),
706 };
707
708 chunks.push(chunk);
709 }
710
711 Ok(chunks)
712 }
713
714 pub fn process_chunk(
733 &mut self,
734 sender: &str,
735 chunk: FileChunk,
736 ) -> Result<FileProgress, ChunkRejection> {
737 if chunk.total_chunks == 0 || chunk.chunk_index >= chunk.total_chunks {
738 return Err(ChunkRejection::Invalid);
739 }
740
741 if chunk.file_size == 0 || chunk.file_size > self.config.max_file_size {
742 tracing::warn!(
743 file_id = %chunk.file_id,
744 file_size = chunk.file_size,
745 max_file_size = self.config.max_file_size,
746 "Rejecting file chunk: claimed file_size out of bounds"
747 );
748 return Err(ChunkRejection::Invalid);
749 }
750
751 let max_total_chunks = Self::derived_total_chunks_cap(self.config.max_file_size);
752 if chunk.total_chunks as u64 > max_total_chunks
753 || chunk.total_chunks as u64 > chunk.file_size
754 {
755 tracing::warn!(
756 file_id = %chunk.file_id,
757 total_chunks = chunk.total_chunks,
758 file_size = chunk.file_size,
759 max_total_chunks,
760 "Rejecting file chunk: claimed total_chunks out of bounds"
761 );
762 return Err(ChunkRejection::Invalid);
763 }
764
765 if chunk.chunk_data.is_empty() || chunk.chunk_data.len() as u64 > chunk.file_size {
766 tracing::warn!(
767 file_id = %chunk.file_id,
768 chunk_len = chunk.chunk_data.len(),
769 file_size = chunk.file_size,
770 "Rejecting file chunk: chunk data length inconsistent with claimed file_size"
771 );
772 return Err(ChunkRejection::Invalid);
773 }
774
775 if self.is_tombstoned(&chunk.file_id) {
779 tracing::debug!(
780 file_id = %chunk.file_id,
781 "Dropping chunk of a previously failed transfer"
782 );
783 return Err(ChunkRejection::PreviouslyFailed);
784 }
785
786 let file_id = chunk.file_id.clone();
787
788 let replaced_charge = match self.active_assemblies.get(&file_id) {
792 Some(assembly) => {
793 if assembly.sender != sender {
794 tracing::warn!(
795 file_id = %file_id,
796 sender = %sender,
797 "Rejecting file chunk: sender does not match the assembly's sender"
798 );
799 return Err(ChunkRejection::SenderMismatch);
800 }
801 if assembly.total_chunks != chunk.total_chunks
802 || assembly.file_size != chunk.file_size
803 || assembly.file_checksum != chunk.file_checksum
804 {
805 tracing::warn!(
806 file_id = %file_id,
807 "Rejecting file chunk: metadata does not match the existing assembly"
808 );
809 return Err(ChunkRejection::MetadataMismatch);
810 }
811 let replaced = assembly
815 .received_chunks
816 .get(&chunk.chunk_index)
817 .map(|c| c.len() as u64);
818 if assembly.received_bytes - replaced.unwrap_or(0) + chunk.chunk_data.len() as u64
819 > assembly.file_size
820 {
821 tracing::warn!(
822 file_id = %file_id,
823 received_bytes = assembly.received_bytes,
824 chunk_len = chunk.chunk_data.len(),
825 file_size = assembly.file_size,
826 "Rejecting file chunk: received bytes exceed claimed file_size"
827 );
828 return Err(ChunkRejection::Invalid);
829 }
830 replaced.map_or(0, |len| len + Self::PER_CHUNK_OVERHEAD)
833 }
834 None => {
835 if self.active_assemblies.len() >= self.config.max_concurrent_assemblies {
836 tracing::warn!(
837 file_id = %file_id,
838 active = self.active_assemblies.len(),
839 max = self.config.max_concurrent_assemblies,
840 "Rejecting file chunk: too many concurrent inbound transfers"
841 );
842 self.tombstone(&file_id);
846 return Err(ChunkRejection::TooManyTransfers);
847 }
848 let sender_active = self
849 .active_assemblies
850 .values()
851 .filter(|a| a.sender == sender)
852 .count();
853 if sender_active >= self.config.max_assemblies_per_sender {
854 tracing::warn!(
855 file_id = %file_id,
856 sender = %sender,
857 sender_active,
858 max = self.config.max_assemblies_per_sender,
859 "Rejecting file chunk: sender's concurrent-transfer quota reached"
860 );
861 self.tombstone(&file_id);
862 return Err(ChunkRejection::SenderQuotaExceeded);
863 }
864 0
865 }
866 };
867
868 let buffered: u64 = self
872 .active_assemblies
873 .values()
874 .map(FileAssembly::charged_bytes)
875 .sum();
876 if buffered - replaced_charge + chunk.chunk_data.len() as u64 + Self::PER_CHUNK_OVERHEAD
877 > self.config.max_total_buffered_bytes
878 {
879 tracing::warn!(
880 file_id = %file_id,
881 buffered,
882 chunk_len = chunk.chunk_data.len(),
883 budget = self.config.max_total_buffered_bytes,
884 "Rejecting file chunk: buffered-bytes budget exhausted, dropping transfer"
885 );
886 self.active_assemblies.remove(&file_id);
892 self.tombstone(&file_id);
893 return Err(ChunkRejection::BufferBudgetExhausted);
894 }
895
896 let assembly = self
898 .active_assemblies
899 .entry(file_id.clone())
900 .or_insert_with(|| FileAssembly::new(sender, &chunk));
901
902 assembly.add_chunk(chunk.chunk_index, chunk.chunk_data);
904
905 let mut progress = assembly.progress();
907 progress.file_id = file_id;
908
909 Ok(progress)
910 }
911
912 pub fn is_complete(&self, file_id: &str) -> bool {
914 self.active_assemblies
915 .get(file_id)
916 .map(|assembly| assembly.is_complete())
917 .unwrap_or(false)
918 }
919
920 pub fn finalize_file(&mut self, file_id: &str) -> Option<Vec<u8>> {
933 let assembly = self.active_assemblies.get(file_id)?;
934 let file_data = assembly.reassemble()?;
935
936 if file_data.len() as u64 != assembly.file_size {
938 tracing::warn!(
939 file_id = %file_id,
940 claimed = assembly.file_size,
941 actual = file_data.len(),
942 "File size mismatch — corrupted or tampered transfer"
943 );
944 self.active_assemblies.remove(file_id);
945 self.tombstone(file_id);
946 return None;
947 }
948
949 let actual_checksum = format!("{:x}", Sha256::digest(&file_data));
951 if actual_checksum != assembly.file_checksum {
952 tracing::warn!(
953 file_id = %file_id,
954 expected = %assembly.file_checksum,
955 actual = %actual_checksum,
956 "File checksum mismatch — corrupted or tampered transfer"
957 );
958 self.active_assemblies.remove(file_id);
959 self.tombstone(file_id);
960 return None;
961 }
962
963 self.active_assemblies.remove(file_id);
964 Some(file_data)
965 }
966
967 pub fn cleanup_stale_transfers(&mut self, max_age: Duration) -> Vec<StaleTransfer> {
973 let now = Instant::now();
974
975 self.failed_transfers
976 .retain(|_, last_seen| now.duration_since(*last_seen) <= max_age);
977 let failed_transfers = &self.failed_transfers;
978 self.failed_transfer_order
979 .retain(|key| failed_transfers.contains_key(key));
980
981 let stale_file_ids: Vec<String> = self
982 .active_assemblies
983 .iter()
984 .filter_map(|(file_id, assembly)| {
985 if now.duration_since(assembly.last_updated_at) > max_age {
986 return Some(file_id.clone());
987 }
988 None
989 })
990 .collect();
991
992 stale_file_ids
993 .into_iter()
994 .filter_map(|file_id| {
995 let assembly = self.active_assemblies.remove(&file_id)?;
996 self.tombstone(&file_id);
997 Some(StaleTransfer {
998 file_id,
999 file_name: assembly.file_name,
1000 sender: assembly.sender,
1001 })
1002 })
1003 .collect()
1004 }
1005
1006 pub fn get_progress(&self, file_id: &str) -> Option<FileProgress> {
1008 self.active_assemblies.get(file_id).map(|assembly| {
1009 let mut progress = assembly.progress();
1010 progress.file_id = file_id.to_string();
1011 progress
1012 })
1013 }
1014
1015 pub fn cancel_transfer(&mut self, file_id: &str) -> bool {
1017 self.active_assemblies.remove(file_id).is_some()
1018 }
1019
1020 pub(crate) fn refuse_transfer(&mut self, file_id: &str) {
1029 self.active_assemblies.remove(file_id);
1030 self.tombstone(file_id);
1031 }
1032
1033 pub fn active_transfer_count(&self) -> usize {
1035 self.active_assemblies.len()
1036 }
1037
1038 #[cfg(test)]
1041 pub(crate) fn backdate_transfers(&mut self, by: Duration) {
1042 let backdated = Instant::now()
1046 .checked_sub(by)
1047 .expect("backdate_transfers needs system uptime longer than the backdate duration");
1048 for assembly in self.active_assemblies.values_mut() {
1049 assembly.last_updated_at = backdated;
1050 }
1051 for last_seen in self.failed_transfers.values_mut() {
1052 *last_seen = backdated;
1053 }
1054 }
1055}
1056
1057impl Default for FileTransferManager {
1058 fn default() -> Self {
1059 Self::new()
1060 }
1061}
1062
1063pub struct OutboundTransferState {
1069 chunks: Vec<FileChunk>,
1070 next_unsent: usize,
1071 in_flight: HashSet<u32>,
1072 acked: HashSet<u32>,
1073 max_in_flight: usize,
1074}
1075
1076impl OutboundTransferState {
1077 pub fn new(chunks: Vec<FileChunk>, max_in_flight: usize) -> Self {
1079 Self {
1080 chunks,
1081 next_unsent: 0,
1082 in_flight: HashSet::new(),
1083 acked: HashSet::new(),
1084 max_in_flight,
1085 }
1086 }
1087
1088 pub fn next_chunks_to_send(&mut self) -> Vec<FileChunk> {
1091 let mut batch = Vec::new();
1092 while self.in_flight.len() < self.max_in_flight && self.next_unsent < self.chunks.len() {
1093 let chunk = self.chunks[self.next_unsent].clone();
1094 self.in_flight.insert(chunk.chunk_index);
1095 self.next_unsent += 1;
1096 batch.push(chunk);
1097 }
1098 batch
1099 }
1100
1101 pub fn on_chunk_ack(&mut self, chunk_index: u32) -> bool {
1104 self.acked.insert(chunk_index);
1105 self.in_flight.remove(&chunk_index)
1106 }
1107
1108 pub fn on_chunk_failed(&mut self, chunk_index: u32) {
1111 self.in_flight.remove(&chunk_index);
1112 }
1113
1114 pub fn total_chunks(&self) -> u32 {
1116 self.chunks.len() as u32
1117 }
1118
1119 pub fn acked_count(&self) -> u32 {
1121 self.acked.len() as u32
1122 }
1123
1124 pub fn is_fully_acked(&self) -> bool {
1126 self.acked.len() == self.chunks.len()
1127 }
1128
1129 pub fn all_sent(&self) -> bool {
1132 self.next_unsent >= self.chunks.len()
1133 }
1134
1135 pub fn has_capacity(&self) -> bool {
1137 self.in_flight.len() < self.max_in_flight && self.next_unsent < self.chunks.len()
1138 }
1139}
1140
1141#[cfg(test)]
1142mod tests {
1143 use super::*;
1144
1145 fn sample_chunk() -> FileChunk {
1146 FileChunk {
1147 file_id: "file1".to_string(),
1148 file_name: "test.txt".to_string(),
1149 file_size: 13,
1150 total_chunks: 1,
1151 chunk_index: 0,
1152 chunk_data: b"Hello, World!".to_vec(),
1153 file_checksum: "abc123".to_string(),
1154 }
1155 }
1156
1157 #[test]
1158 fn test_from_bytes_truncated_payload_is_unexpected_eof() {
1159 let bytes = sample_chunk().to_bytes();
1160 let err = FileChunk::from_bytes(&bytes[..2]).unwrap_err();
1162 assert!(matches!(
1163 err,
1164 ChunkDecodeError::UnexpectedEof { what: "u32" }
1165 ));
1166 let after_strings = 4 + "file1".len() + 4 + "test.txt".len();
1169 let err = FileChunk::from_bytes(&bytes[..after_strings + 4]).unwrap_err();
1170 assert!(matches!(
1171 err,
1172 ChunkDecodeError::UnexpectedEof { what: "u64" }
1173 ));
1174 let err = FileChunk::from_bytes(&bytes[..bytes.len() - 1]).unwrap_err();
1176 assert!(matches!(
1177 err,
1178 ChunkDecodeError::UnexpectedEof { what: "bytes" }
1179 ));
1180 }
1181
1182 #[test]
1183 fn test_from_bytes_oversized_length_prefix_is_field_too_long() {
1184 let mut bytes = sample_chunk().to_bytes();
1185 let huge = (FileChunk::MAX_STRING_FIELD_LEN as u32 + 1).to_le_bytes();
1187 bytes[..4].copy_from_slice(&huge);
1188 let err = FileChunk::from_bytes(&bytes).unwrap_err();
1189 assert!(matches!(
1190 err,
1191 ChunkDecodeError::FieldTooLong {
1192 field: "file_id_len",
1193 ..
1194 }
1195 ));
1196 }
1197
1198 #[test]
1199 fn test_from_bytes_invalid_utf8_names_the_field() {
1200 let mut bytes = sample_chunk().to_bytes();
1201 bytes[4] = 0xFF;
1203 let err = FileChunk::from_bytes(&bytes).unwrap_err();
1204 assert!(matches!(
1205 err,
1206 ChunkDecodeError::InvalidUtf8 {
1207 field: "file_id",
1208 ..
1209 }
1210 ));
1211 }
1212
1213 #[test]
1214 fn test_chunk_small_file() {
1215 let manager = FileTransferManager::new();
1216 let file_data = b"Hello, World!".to_vec();
1217
1218 let chunks = manager
1219 .chunk_file(
1220 "file1".to_string(),
1221 "test.txt".to_string(),
1222 file_data.clone(),
1223 None,
1224 )
1225 .unwrap();
1226
1227 assert_eq!(chunks.len(), 1); assert_eq!(chunks[0].chunk_index, 0);
1229 assert_eq!(chunks[0].total_chunks, 1);
1230 assert_eq!(chunks[0].file_size, file_data.len() as u64);
1231 assert_eq!(chunks[0].chunk_data, file_data);
1232 }
1233
1234 #[test]
1235 fn test_chunk_large_file() {
1236 let config = FileTransferConfig {
1237 chunk_size: 100, max_file_size: 1024 * 1024,
1239 ..FileTransferConfig::default()
1240 };
1241 let manager = FileTransferManager::with_config(config);
1242
1243 let file_data = vec![0u8; 350]; let chunks = manager
1246 .chunk_file("file1".to_string(), "test.bin".to_string(), file_data, None)
1247 .unwrap();
1248
1249 assert_eq!(chunks.len(), 4); assert_eq!(chunks[0].chunk_data.len(), 100);
1251 assert_eq!(chunks[1].chunk_data.len(), 100);
1252 assert_eq!(chunks[2].chunk_data.len(), 100);
1253 assert_eq!(chunks[3].chunk_data.len(), 50);
1254 }
1255
1256 #[test]
1257 fn test_file_too_large() {
1258 let config = FileTransferConfig {
1259 chunk_size: 1024,
1260 max_file_size: 1000,
1261 ..FileTransferConfig::default()
1262 };
1263 let manager = FileTransferManager::with_config(config);
1264
1265 let file_data = vec![0u8; 2000]; let result = manager.chunk_file(
1268 "file1".to_string(),
1269 "large.bin".to_string(),
1270 file_data,
1271 None,
1272 );
1273 assert!(result.is_err());
1274 }
1275
1276 #[test]
1277 fn test_empty_file() {
1278 let manager = FileTransferManager::new();
1279 let result = manager.chunk_file("file1".to_string(), "empty.txt".to_string(), vec![], None);
1280 assert!(result.is_err());
1281 }
1282
1283 #[test]
1284 fn test_zero_chunk_size_override_rejected() {
1285 let manager = FileTransferManager::new();
1286 let result = manager.chunk_file(
1287 "file1".to_string(),
1288 "test.txt".to_string(),
1289 b"hello".to_vec(),
1290 Some(0),
1291 );
1292 assert!(result.is_err());
1293 }
1294
1295 #[test]
1296 fn test_reassemble_file() {
1297 let config = FileTransferConfig {
1298 chunk_size: 10,
1299 max_file_size: 1024,
1300 ..FileTransferConfig::default()
1301 };
1302 let mut manager = FileTransferManager::with_config(config);
1303
1304 let original_data = b"Hello, World!".to_vec();
1305 let chunks = manager
1306 .chunk_file(
1307 "file1".to_string(),
1308 "test.txt".to_string(),
1309 original_data.clone(),
1310 None,
1311 )
1312 .unwrap();
1313
1314 for chunk in chunks {
1316 manager.process_chunk("peer", chunk).unwrap();
1317 }
1318
1319 assert!(manager.is_complete("file1"));
1320
1321 let reassembled = manager.finalize_file("file1").unwrap();
1322 assert_eq!(reassembled, original_data);
1323 }
1324
1325 #[test]
1326 fn test_reassemble_out_of_order() {
1327 let config = FileTransferConfig {
1328 chunk_size: 10,
1329 max_file_size: 1024,
1330 ..FileTransferConfig::default()
1331 };
1332 let mut manager = FileTransferManager::with_config(config);
1333
1334 let original_data = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ".to_vec();
1335 let mut chunks = manager
1336 .chunk_file(
1337 "file1".to_string(),
1338 "alphabet.txt".to_string(),
1339 original_data.clone(),
1340 None,
1341 )
1342 .unwrap();
1343
1344 chunks.reverse();
1346
1347 for chunk in chunks {
1349 manager.process_chunk("peer", chunk).unwrap();
1350 }
1351
1352 assert!(manager.is_complete("file1"));
1353
1354 let reassembled = manager.finalize_file("file1").unwrap();
1355 assert_eq!(reassembled, original_data);
1356 }
1357
1358 #[test]
1359 fn test_progress_tracking() {
1360 let config = FileTransferConfig {
1361 chunk_size: 100,
1362 max_file_size: 1024,
1363 ..FileTransferConfig::default()
1364 };
1365 let mut manager = FileTransferManager::with_config(config);
1366
1367 let file_data = vec![0u8; 250]; let chunks = manager
1369 .chunk_file("file1".to_string(), "test.bin".to_string(), file_data, None)
1370 .unwrap();
1371
1372 assert_eq!(chunks.len(), 3);
1373
1374 let progress = manager.process_chunk("peer", chunks[0].clone()).unwrap();
1376 assert_eq!(progress.chunks_completed, 1);
1377 assert_eq!(progress.total_chunks, 3);
1378 assert_eq!(progress.percentage, 33); let progress = manager.process_chunk("peer", chunks[1].clone()).unwrap();
1382 assert_eq!(progress.chunks_completed, 2);
1383 assert_eq!(progress.percentage, 66); let progress = manager.process_chunk("peer", chunks[2].clone()).unwrap();
1387 assert_eq!(progress.chunks_completed, 3);
1388 assert_eq!(progress.percentage, 100); assert!(manager.is_complete("file1"));
1391 }
1392
1393 #[test]
1394 fn test_duplicate_chunk() {
1395 let mut manager = FileTransferManager::new();
1396
1397 let file_data = b"Test data".to_vec();
1398 let chunks = manager
1399 .chunk_file("file1".to_string(), "test.txt".to_string(), file_data, None)
1400 .unwrap();
1401
1402 manager.process_chunk("peer", chunks[0].clone()).unwrap();
1404 manager.process_chunk("peer", chunks[0].clone()).unwrap();
1405
1406 let progress = manager.get_progress("file1").unwrap();
1408 assert_eq!(progress.chunks_completed, 1);
1409 }
1410
1411 #[test]
1412 fn test_multiple_files() {
1413 let config = FileTransferConfig {
1414 chunk_size: 10,
1415 max_file_size: 1024,
1416 ..FileTransferConfig::default()
1417 };
1418 let mut manager = FileTransferManager::with_config(config);
1419
1420 let data1 = b"File 1 data".to_vec();
1421 let data2 = b"File 2 data".to_vec();
1422
1423 let chunks1 = manager
1424 .chunk_file(
1425 "file1".to_string(),
1426 "file1.txt".to_string(),
1427 data1.clone(),
1428 None,
1429 )
1430 .unwrap();
1431 let chunks2 = manager
1432 .chunk_file(
1433 "file2".to_string(),
1434 "file2.txt".to_string(),
1435 data2.clone(),
1436 None,
1437 )
1438 .unwrap();
1439
1440 for chunk in chunks1 {
1442 manager.process_chunk("peer", chunk).unwrap();
1443 }
1444 for chunk in chunks2 {
1445 manager.process_chunk("peer", chunk).unwrap();
1446 }
1447
1448 assert_eq!(manager.active_transfer_count(), 2);
1449 assert!(manager.is_complete("file1"));
1450 assert!(manager.is_complete("file2"));
1451 }
1452
1453 #[test]
1454 fn test_cancel_transfer() {
1455 let mut manager = FileTransferManager::new();
1456
1457 let file_data = b"Test data".to_vec();
1458 let chunks = manager
1459 .chunk_file("file1".to_string(), "test.txt".to_string(), file_data, None)
1460 .unwrap();
1461
1462 manager.process_chunk("peer", chunks[0].clone()).unwrap();
1463 assert_eq!(manager.active_transfer_count(), 1);
1464
1465 assert!(manager.cancel_transfer("file1"));
1466 assert_eq!(manager.active_transfer_count(), 0);
1467 assert!(!manager.is_complete("file1"));
1468 }
1469
1470 #[test]
1471 fn test_get_progress() {
1472 let mut manager = FileTransferManager::new();
1473
1474 let file_data = b"Test data".to_vec();
1475 let chunks = manager
1476 .chunk_file("file1".to_string(), "test.txt".to_string(), file_data, None)
1477 .unwrap();
1478
1479 manager.process_chunk("peer", chunks[0].clone()).unwrap();
1480
1481 let progress = manager.get_progress("file1").unwrap();
1482 assert_eq!(progress.file_id, "file1");
1483 assert_eq!(progress.file_name, "test.txt");
1484 assert!(progress.percentage > 0);
1485 }
1486
1487 #[test]
1488 fn test_progress_nonexistent_file() {
1489 let manager = FileTransferManager::new();
1490 assert!(manager.get_progress("nonexistent").is_none());
1491 }
1492
1493 #[test]
1494 fn test_calculate_percentage() {
1495 assert_eq!(FileProgress::calculate_percentage(0, 10), 0);
1496 assert_eq!(FileProgress::calculate_percentage(5, 10), 50);
1497 assert_eq!(FileProgress::calculate_percentage(10, 10), 100);
1498 assert_eq!(FileProgress::calculate_percentage(1, 3), 33);
1499 }
1500
1501 #[test]
1502 fn test_invalid_chunk_index_rejected() {
1503 let mut manager = FileTransferManager::new();
1504 let chunks = manager
1505 .chunk_file(
1506 "file1".to_string(),
1507 "test.txt".to_string(),
1508 b"hello world".to_vec(),
1509 None,
1510 )
1511 .unwrap();
1512
1513 let mut invalid_chunk = chunks[0].clone();
1514 invalid_chunk.chunk_index = invalid_chunk.total_chunks;
1515
1516 assert!(manager.process_chunk("peer", invalid_chunk).is_err());
1517 assert!(manager.get_progress("file1").is_none());
1518 }
1519
1520 fn attack_chunk(
1521 file_id: &str,
1522 file_size: u64,
1523 total_chunks: u32,
1524 chunk_index: u32,
1525 data: Vec<u8>,
1526 ) -> FileChunk {
1527 FileChunk {
1528 file_id: file_id.to_string(),
1529 file_name: "attack.bin".to_string(),
1530 file_size,
1531 total_chunks,
1532 chunk_index,
1533 chunk_data: data,
1534 file_checksum: "0000".to_string(),
1535 }
1536 }
1537
1538 #[test]
1539 fn test_huge_file_size_claim_rejected() {
1540 let mut manager = FileTransferManager::new();
1543 let chunk = attack_chunk("evil", u64::MAX, 1, 0, vec![0u8; 16]);
1544 assert!(manager.process_chunk("peer", chunk).is_err());
1545 assert_eq!(manager.active_transfer_count(), 0);
1546 assert!(!manager.is_complete("evil"));
1547 assert!(manager.finalize_file("evil").is_none());
1548 }
1549
1550 #[test]
1551 fn test_file_size_over_configured_max_rejected() {
1552 let config = FileTransferConfig {
1553 chunk_size: 1024,
1554 max_file_size: 1000,
1555 ..FileTransferConfig::default()
1556 };
1557 let mut manager = FileTransferManager::with_config(config);
1558 let chunk = attack_chunk("f", 2000, 2, 0, vec![0u8; 100]);
1559 assert!(manager.process_chunk("peer", chunk).is_err());
1560 assert_eq!(manager.active_transfer_count(), 0);
1561 }
1562
1563 #[test]
1564 fn test_zero_file_size_rejected() {
1565 let mut manager = FileTransferManager::new();
1566 let chunk = attack_chunk("f", 0, 1, 0, vec![0u8; 8]);
1567 assert!(manager.process_chunk("peer", chunk).is_err());
1568 }
1569
1570 #[test]
1571 fn test_total_chunks_exceeding_file_size_rejected() {
1572 let mut manager = FileTransferManager::new();
1574 let chunk = attack_chunk("f", 10_000, 20_000, 0, vec![0u8; 8]);
1575 assert!(manager.process_chunk("peer", chunk).is_err());
1576 }
1577
1578 #[test]
1579 fn test_total_chunks_exceeding_derived_cap_rejected() {
1580 let mut manager = FileTransferManager::new();
1582 let chunk = attack_chunk("f", 100 * 1024 * 1024, 200_000, 0, vec![0u8; 8]);
1583 assert!(manager.process_chunk("peer", chunk).is_err());
1584 }
1585
1586 #[test]
1587 fn test_empty_chunk_data_rejected() {
1588 let mut manager = FileTransferManager::new();
1589 let chunk = attack_chunk("f", 100, 2, 0, vec![]);
1590 assert!(manager.process_chunk("peer", chunk).is_err());
1591 }
1592
1593 #[test]
1594 fn test_chunk_data_larger_than_file_size_rejected() {
1595 let mut manager = FileTransferManager::new();
1596 let chunk = attack_chunk("f", 10, 1, 0, vec![0u8; 64]);
1597 assert!(manager.process_chunk("peer", chunk).is_err());
1598 }
1599
1600 #[test]
1601 fn test_cumulative_bytes_exceeding_claim_rejected() {
1602 let mut manager = FileTransferManager::new();
1603 let first = attack_chunk("f", 2048, 2, 0, vec![0u8; 1500]);
1605 assert!(manager.process_chunk("peer", first).is_ok());
1606 let second = attack_chunk("f", 2048, 2, 1, vec![0u8; 1500]);
1607 assert!(manager.process_chunk("peer", second).is_err());
1608 let progress = manager.get_progress("f").unwrap();
1610 assert_eq!(progress.chunks_completed, 1);
1611 }
1612
1613 #[test]
1614 fn test_duplicate_chunk_does_not_double_count_bytes() {
1615 let mut manager = FileTransferManager::new();
1616 let data0 = vec![7u8; 1500];
1619 let data1 = vec![9u8; 548];
1620 let mut whole = data0.clone();
1621 whole.extend_from_slice(&data1);
1622 let checksum = format!("{:x}", Sha256::digest(&whole));
1623
1624 let make = |index: u32, data: Vec<u8>| FileChunk {
1625 file_id: "f".to_string(),
1626 file_name: "file.bin".to_string(),
1627 file_size: 2048,
1628 total_chunks: 2,
1629 chunk_index: index,
1630 chunk_data: data,
1631 file_checksum: checksum.clone(),
1632 };
1633
1634 assert!(manager
1635 .process_chunk("peer", make(0, data0.clone()))
1636 .is_ok());
1637 assert!(manager.process_chunk("peer", make(0, data0)).is_ok());
1638 assert!(manager.process_chunk("peer", make(1, data1)).is_ok());
1639 assert!(manager.is_complete("f"));
1640 assert_eq!(manager.finalize_file("f").unwrap(), whole);
1641 }
1642
1643 #[test]
1644 fn test_concurrent_assembly_cap_enforced() {
1645 let config = FileTransferConfig {
1646 max_concurrent_assemblies: 2,
1647 ..FileTransferConfig::default()
1648 };
1649 let mut manager = FileTransferManager::with_config(config);
1650
1651 assert!(manager
1652 .process_chunk("peer", attack_chunk("f1", 100, 2, 0, vec![0u8; 10]))
1653 .is_ok());
1654 assert!(manager
1655 .process_chunk("peer", attack_chunk("f2", 100, 2, 0, vec![0u8; 10]))
1656 .is_ok());
1657 assert_eq!(
1659 manager
1660 .process_chunk("peer", attack_chunk("f3", 100, 2, 0, vec![0u8; 10]))
1661 .unwrap_err(),
1662 ChunkRejection::TooManyTransfers
1663 );
1664 assert!(manager
1666 .process_chunk("peer", attack_chunk("f1", 100, 2, 1, vec![0u8; 10]))
1667 .is_ok());
1668 assert!(manager.cancel_transfer("f2"));
1672 assert_eq!(
1673 manager
1674 .process_chunk("peer", attack_chunk("f3", 100, 2, 0, vec![0u8; 10]))
1675 .unwrap_err(),
1676 ChunkRejection::PreviouslyFailed
1677 );
1678 assert!(manager
1679 .process_chunk("peer", attack_chunk("f4", 100, 2, 0, vec![0u8; 10]))
1680 .is_ok());
1681 }
1682
1683 #[test]
1684 fn test_sender_quota_enforced() {
1685 let config = FileTransferConfig {
1686 max_concurrent_assemblies: 4,
1687 max_assemblies_per_sender: 1,
1688 ..FileTransferConfig::default()
1689 };
1690 let mut manager = FileTransferManager::with_config(config);
1691
1692 assert!(manager
1693 .process_chunk("alice", attack_chunk("f1", 100, 2, 0, vec![0u8; 10]))
1694 .is_ok());
1695 assert_eq!(
1697 manager
1698 .process_chunk("alice", attack_chunk("f2", 100, 2, 0, vec![0u8; 10]))
1699 .unwrap_err(),
1700 ChunkRejection::SenderQuotaExceeded
1701 );
1702 assert!(manager
1704 .process_chunk("bob", attack_chunk("f3", 100, 2, 0, vec![0u8; 10]))
1705 .is_ok());
1706 }
1707
1708 #[test]
1709 fn test_sender_mismatch_rejected() {
1710 let mut manager = FileTransferManager::new();
1711 assert!(manager
1712 .process_chunk("alice", attack_chunk("f", 100, 2, 0, vec![0u8; 10]))
1713 .is_ok());
1714 assert_eq!(
1717 manager
1718 .process_chunk("mallory", attack_chunk("f", 100, 2, 1, vec![0u8; 10]))
1719 .unwrap_err(),
1720 ChunkRejection::SenderMismatch
1721 );
1722 assert_eq!(manager.get_progress("f").unwrap().chunks_completed, 1);
1724 }
1725
1726 #[test]
1727 fn test_buffer_budget_rejects_new_transfer() {
1728 let config = FileTransferConfig {
1731 max_file_size: 1_000_000,
1732 max_total_buffered_bytes: 1_065_536,
1733 ..FileTransferConfig::default()
1734 };
1735 let mut manager = FileTransferManager::with_config(config);
1736
1737 assert!(manager
1738 .process_chunk(
1739 "alice",
1740 attack_chunk("f1", 1_000_000, 2, 0, vec![0u8; 600_000])
1741 )
1742 .is_ok());
1743 assert_eq!(
1746 manager
1747 .process_chunk(
1748 "bob",
1749 attack_chunk("f2", 1_000_000, 2, 0, vec![0u8; 500_000])
1750 )
1751 .unwrap_err(),
1752 ChunkRejection::BufferBudgetExhausted
1753 );
1754 assert_eq!(manager.active_transfer_count(), 1);
1755 assert!(manager.get_progress("f2").is_none());
1756 assert_eq!(manager.get_progress("f1").unwrap().chunks_completed, 1);
1758 }
1759
1760 #[test]
1761 fn test_buffer_budget_drops_existing_assembly() {
1762 let config = FileTransferConfig {
1763 max_file_size: 1_000_000,
1764 max_total_buffered_bytes: 1_065_536,
1765 ..FileTransferConfig::default()
1766 };
1767 let mut manager = FileTransferManager::with_config(config);
1768
1769 assert!(manager
1770 .process_chunk(
1771 "alice",
1772 attack_chunk("f1", 1_000_000, 3, 0, vec![0u8; 600_000])
1773 )
1774 .is_ok());
1775 assert!(manager
1776 .process_chunk("bob", attack_chunk("f2", 400_000, 2, 0, vec![0u8; 300_000]))
1777 .is_ok());
1778 assert_eq!(
1783 manager
1784 .process_chunk(
1785 "alice",
1786 attack_chunk("f1", 1_000_000, 3, 1, vec![0u8; 300_000])
1787 )
1788 .unwrap_err(),
1789 ChunkRejection::BufferBudgetExhausted
1790 );
1791 assert!(manager.get_progress("f1").is_none());
1792 assert_eq!(manager.active_transfer_count(), 1);
1793 assert!(manager
1795 .process_chunk("bob", attack_chunk("f2", 400_000, 2, 1, vec![0u8; 100_000]))
1796 .is_ok());
1797 }
1798
1799 #[test]
1800 fn test_budget_charges_per_chunk_overhead() {
1801 let config = FileTransferConfig {
1805 max_file_size: 10_000,
1806 max_total_buffered_bytes: 0,
1807 ..FileTransferConfig::default()
1808 };
1809 let mut manager = FileTransferManager::with_config(config);
1810
1811 for i in 0..1024u32 {
1813 manager
1814 .process_chunk("alice", attack_chunk("f1", 10_000, 1024, i, vec![0u8; 1]))
1815 .unwrap();
1816 }
1817
1818 let mut rejected = None;
1821 for i in 0..1024u32 {
1822 if let Err(e) =
1823 manager.process_chunk("bob", attack_chunk("f2", 10_000, 1024, i, vec![0u8; 1]))
1824 {
1825 rejected = Some((i, e));
1826 break;
1827 }
1828 }
1829 let (index, err) = rejected.expect("tiny-chunk flood must exhaust the budget");
1830 assert_eq!(err, ChunkRejection::BufferBudgetExhausted);
1831 let total_payload = 1024 + u64::from(index);
1832 assert!(
1833 total_payload < 2_000,
1834 "rejection must be overhead-driven, but {} payload bytes were buffered",
1835 total_payload
1836 );
1837 assert!(manager.get_progress("f2").is_none());
1839 }
1840
1841 #[test]
1842 fn test_manual_sender_cannot_be_a_wire_user_id() {
1843 assert!(offline_protocol_core::UserId::new(FileTransferManager::MANUAL_SENDER).is_err());
1847 }
1848
1849 #[test]
1850 fn test_resource_rejected_transfer_is_tombstoned() {
1851 let config = FileTransferConfig {
1856 max_concurrent_assemblies: 1,
1857 ..FileTransferConfig::default()
1858 };
1859 let mut manager = FileTransferManager::with_config(config);
1860
1861 assert!(manager
1862 .process_chunk("alice", attack_chunk("f1", 100, 2, 0, vec![0u8; 10]))
1863 .is_ok());
1864 assert_eq!(
1865 manager
1866 .process_chunk("bob", attack_chunk("f2", 1000, 10, 0, vec![0u8; 100]))
1867 .unwrap_err(),
1868 ChunkRejection::TooManyTransfers
1869 );
1870 for index in 1..10u32 {
1871 assert_eq!(
1872 manager
1873 .process_chunk("bob", attack_chunk("f2", 1000, 10, index, vec![0u8; 100]))
1874 .unwrap_err(),
1875 ChunkRejection::PreviouslyFailed
1876 );
1877 }
1878 assert!(manager.cancel_transfer("f1"));
1881 assert_eq!(
1882 manager
1883 .process_chunk("bob", attack_chunk("f2", 1000, 10, 5, vec![0u8; 100]))
1884 .unwrap_err(),
1885 ChunkRejection::PreviouslyFailed
1886 );
1887 assert_eq!(manager.active_transfer_count(), 0);
1888 }
1889
1890 #[test]
1891 fn test_budget_dropped_transfer_cannot_resurrect() {
1892 let config = FileTransferConfig {
1893 max_file_size: 1_000_000,
1894 max_total_buffered_bytes: 1_065_536,
1895 ..FileTransferConfig::default()
1896 };
1897 let mut manager = FileTransferManager::with_config(config);
1898
1899 assert!(manager
1900 .process_chunk(
1901 "alice",
1902 attack_chunk("f1", 1_000_000, 3, 0, vec![0u8; 600_000])
1903 )
1904 .is_ok());
1905 assert!(manager
1906 .process_chunk("bob", attack_chunk("f2", 400_000, 2, 0, vec![0u8; 300_000]))
1907 .is_ok());
1908 assert_eq!(
1909 manager
1910 .process_chunk(
1911 "alice",
1912 attack_chunk("f1", 1_000_000, 3, 1, vec![0u8; 300_000])
1913 )
1914 .unwrap_err(),
1915 ChunkRejection::BufferBudgetExhausted
1916 );
1917 assert_eq!(
1921 manager
1922 .process_chunk(
1923 "alice",
1924 attack_chunk("f1", 1_000_000, 3, 2, vec![0u8; 100_000])
1925 )
1926 .unwrap_err(),
1927 ChunkRejection::PreviouslyFailed
1928 );
1929 assert!(manager.get_progress("f1").is_none());
1930 assert_eq!(manager.active_transfer_count(), 1);
1931 }
1932
1933 #[test]
1934 fn test_finalize_integrity_failure_tombstones_id() {
1935 let mut manager = FileTransferManager::new();
1936
1937 let data = vec![3u8; 100];
1940 let checksum = format!("{:x}", Sha256::digest(&data));
1941 let chunk = FileChunk {
1942 file_id: "f".to_string(),
1943 file_name: "bad.bin".to_string(),
1944 file_size: 2048,
1945 total_chunks: 1,
1946 chunk_index: 0,
1947 chunk_data: data,
1948 file_checksum: checksum,
1949 };
1950 assert!(manager.process_chunk("peer", chunk.clone()).is_ok());
1951 assert!(manager.is_complete("f"));
1952 assert!(manager.finalize_file("f").is_none());
1953
1954 assert_eq!(
1956 manager.process_chunk("peer", chunk).unwrap_err(),
1957 ChunkRejection::PreviouslyFailed
1958 );
1959 assert_eq!(manager.active_transfer_count(), 0);
1960 }
1961
1962 #[test]
1963 fn test_stale_swept_transfer_is_tombstoned() {
1964 let mut manager = FileTransferManager::new();
1965
1966 assert!(manager
1967 .process_chunk("peer", attack_chunk("f", 200, 2, 0, vec![0u8; 100]))
1968 .is_ok());
1969 manager.backdate_transfers(Duration::from_secs(2));
1970 let swept = manager.cleanup_stale_transfers(Duration::from_secs(1));
1971 assert_eq!(swept.len(), 1);
1972
1973 assert_eq!(
1976 manager
1977 .process_chunk("peer", attack_chunk("f", 200, 2, 1, vec![0u8; 100]))
1978 .unwrap_err(),
1979 ChunkRejection::PreviouslyFailed
1980 );
1981 assert_eq!(manager.active_transfer_count(), 0);
1982 }
1983
1984 #[test]
1985 fn test_tombstone_expires_after_max_age() {
1986 let config = FileTransferConfig {
1987 max_concurrent_assemblies: 1,
1988 ..FileTransferConfig::default()
1989 };
1990 let mut manager = FileTransferManager::with_config(config);
1991
1992 assert!(manager
1993 .process_chunk("alice", attack_chunk("f1", 100, 2, 0, vec![0u8; 10]))
1994 .is_ok());
1995 assert_eq!(
1996 manager
1997 .process_chunk("bob", attack_chunk("f2", 100, 2, 0, vec![0u8; 10]))
1998 .unwrap_err(),
1999 ChunkRejection::TooManyTransfers
2000 );
2001 assert!(manager.cancel_transfer("f1"));
2002
2003 manager.backdate_transfers(Duration::from_secs(2));
2007 manager.cleanup_stale_transfers(Duration::from_secs(1));
2008 assert!(manager
2009 .process_chunk("bob", attack_chunk("f2", 100, 2, 0, vec![0u8; 10]))
2010 .is_ok());
2011 }
2012
2013 #[test]
2014 fn test_cancel_does_not_tombstone() {
2015 let mut manager = FileTransferManager::new();
2019 assert!(manager
2020 .process_chunk("peer", attack_chunk("f", 100, 2, 0, vec![0u8; 10]))
2021 .is_ok());
2022 assert!(manager.cancel_transfer("f"));
2023 assert!(manager
2024 .process_chunk("peer", attack_chunk("f", 100, 2, 0, vec![0u8; 10]))
2025 .is_ok());
2026 assert_eq!(manager.active_transfer_count(), 1);
2027 }
2028
2029 #[test]
2030 fn test_duplicate_chunk_resent_with_different_size_replaces() {
2031 let mut manager = FileTransferManager::new();
2035 let data0 = vec![7u8; 1500];
2036 let data1 = vec![9u8; 548];
2037 let mut whole = data0.clone();
2038 whole.extend_from_slice(&data1);
2039 let checksum = format!("{:x}", Sha256::digest(&whole));
2040
2041 let make = |index: u32, data: Vec<u8>| FileChunk {
2042 file_id: "f".to_string(),
2043 file_name: "file.bin".to_string(),
2044 file_size: 2048,
2045 total_chunks: 2,
2046 chunk_index: index,
2047 chunk_data: data,
2048 file_checksum: checksum.clone(),
2049 };
2050
2051 assert!(manager
2053 .process_chunk("peer", make(0, vec![0u8; 1400]))
2054 .is_ok());
2055 assert!(manager.process_chunk("peer", make(0, data0)).is_ok());
2056 assert!(manager.process_chunk("peer", make(1, data1)).is_ok());
2057 assert!(manager.is_complete("f"));
2058 assert_eq!(manager.finalize_file("f").unwrap(), whole);
2061 }
2062
2063 #[test]
2064 fn test_single_chunk_at_exactly_file_size_completes() {
2065 let mut manager = FileTransferManager::new();
2066 let data = vec![4u8; 512];
2067 let checksum = format!("{:x}", Sha256::digest(&data));
2068 let chunk = FileChunk {
2069 file_id: "f".to_string(),
2070 file_name: "exact.bin".to_string(),
2071 file_size: 512,
2072 total_chunks: 1,
2073 chunk_index: 0,
2074 chunk_data: data.clone(),
2075 file_checksum: checksum,
2076 };
2077 assert!(manager.process_chunk("peer", chunk).is_ok());
2078 assert!(manager.is_complete("f"));
2079 assert_eq!(manager.finalize_file("f").unwrap(), data);
2080 }
2081
2082 #[test]
2083 fn test_file_at_exactly_max_file_size_completes() {
2084 let config = FileTransferConfig {
2088 chunk_size: 512,
2089 max_file_size: 1000,
2090 max_total_buffered_bytes: 1000,
2091 ..FileTransferConfig::default()
2092 };
2093 let mut manager = FileTransferManager::with_config(config);
2094
2095 let file_data = vec![5u8; 1000];
2096 let chunks = manager
2097 .chunk_file(
2098 "f".to_string(),
2099 "max.bin".to_string(),
2100 file_data.clone(),
2101 None,
2102 )
2103 .unwrap();
2104 for chunk in chunks {
2105 manager.process_chunk("peer", chunk).unwrap();
2106 }
2107 assert_eq!(manager.finalize_file("f").unwrap(), file_data);
2108 }
2109
2110 #[test]
2111 fn test_chunk_file_rejects_chunking_finer_than_receive_cap() {
2112 let config = FileTransferConfig {
2117 chunk_size: 512,
2118 max_file_size: 1024 * 1024,
2119 ..FileTransferConfig::default()
2120 };
2121 let manager = FileTransferManager::with_config(config);
2122 let result = manager.chunk_file(
2123 "f".to_string(),
2124 "big.bin".to_string(),
2125 vec![0u8; 1024 * 1024],
2126 None,
2127 );
2128 assert!(result.is_err());
2129 }
2130
2131 #[test]
2132 fn test_chunk_file_small_chunk_size_round_trips_within_cap() {
2133 let config = FileTransferConfig {
2137 chunk_size: 512,
2138 max_file_size: 1024 * 1024,
2139 ..FileTransferConfig::default()
2140 };
2141 let mut manager = FileTransferManager::with_config(config);
2142 let file_data: Vec<u8> = (0..400 * 1024).map(|i| (i % 251) as u8).collect();
2143 let chunks = manager
2144 .chunk_file(
2145 "f".to_string(),
2146 "ok.bin".to_string(),
2147 file_data.clone(),
2148 None,
2149 )
2150 .unwrap();
2151 assert_eq!(chunks.len(), 800);
2152 for chunk in chunks {
2153 manager.process_chunk("peer", chunk).unwrap();
2154 }
2155 assert_eq!(manager.finalize_file("f").unwrap(), file_data);
2156 }
2157
2158 #[test]
2159 fn test_total_chunks_at_derived_cap_accepted() {
2160 let config = FileTransferConfig {
2162 max_file_size: 1024 * 1024,
2163 ..FileTransferConfig::default()
2164 };
2165 let mut manager = FileTransferManager::with_config(config);
2166 assert!(manager
2167 .process_chunk(
2168 "peer",
2169 attack_chunk("f1", 1024 * 1024, 1025, 0, vec![0u8; 8])
2170 )
2171 .is_ok());
2172 assert_eq!(
2173 manager
2174 .process_chunk(
2175 "peer",
2176 attack_chunk("f2", 1024 * 1024, 1026, 0, vec![0u8; 8])
2177 )
2178 .unwrap_err(),
2179 ChunkRejection::Invalid
2180 );
2181 }
2182
2183 #[test]
2184 fn test_duplicate_chunk_replaces_data() {
2185 let mut manager = FileTransferManager::new();
2186 let stale = vec![1u8; 100];
2189 let fresh = vec![2u8; 100];
2190 let tail = vec![3u8; 50];
2191 let mut whole = fresh.clone();
2192 whole.extend_from_slice(&tail);
2193 let checksum = format!("{:x}", Sha256::digest(&whole));
2194
2195 let make = |index: u32, data: Vec<u8>| FileChunk {
2196 file_id: "f".to_string(),
2197 file_name: "file.bin".to_string(),
2198 file_size: 150,
2199 total_chunks: 2,
2200 chunk_index: index,
2201 chunk_data: data,
2202 file_checksum: checksum.clone(),
2203 };
2204
2205 manager.process_chunk("peer", make(0, stale)).unwrap();
2206 manager.process_chunk("peer", make(0, fresh)).unwrap();
2207 manager.process_chunk("peer", make(1, tail)).unwrap();
2208 assert_eq!(manager.finalize_file("f").unwrap(), whole);
2209 }
2210
2211 #[test]
2212 fn test_finalize_rejects_size_mismatch() {
2213 let mut manager = FileTransferManager::new();
2214 let data = vec![3u8; 100];
2217 let checksum = format!("{:x}", Sha256::digest(&data));
2218 let chunk = FileChunk {
2219 file_id: "f".to_string(),
2220 file_name: "file.bin".to_string(),
2221 file_size: 2048,
2222 total_chunks: 1,
2223 chunk_index: 0,
2224 chunk_data: data,
2225 file_checksum: checksum,
2226 };
2227 assert!(manager.process_chunk("peer", chunk).is_ok());
2228 assert!(manager.is_complete("f"));
2229 assert!(manager.finalize_file("f").is_none());
2230 assert_eq!(manager.active_transfer_count(), 0);
2232 }
2233
2234 #[test]
2235 fn test_cleanup_stale_transfers() {
2236 let mut manager = FileTransferManager::new();
2237 let chunks = manager
2238 .chunk_file(
2239 "file1".to_string(),
2240 "test.txt".to_string(),
2241 vec![1u8; 64 * 1024],
2242 None,
2243 )
2244 .unwrap();
2245
2246 manager.process_chunk("peer", chunks[0].clone()).unwrap();
2247 std::thread::sleep(Duration::from_millis(3));
2248
2249 let removed = manager.cleanup_stale_transfers(Duration::from_millis(1));
2250 assert_eq!(
2251 removed,
2252 vec![StaleTransfer {
2253 file_id: "file1".to_string(),
2254 file_name: "test.txt".to_string(),
2255 sender: "peer".to_string(),
2256 }]
2257 );
2258 assert!(manager.get_progress("file1").is_none());
2259 }
2260}