Skip to main content

offline_protocol/
file_transfer.rs

1//! File transfer with chunking and reassembly.
2
3use 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/// File transfer configuration.
32#[derive(Debug, Clone)]
33pub struct FileTransferConfig {
34    /// Size of each chunk in bytes. Chunk sizes that would split a file
35    /// into more chunks than the receive-side cap derived from
36    /// `max_file_size` (one chunk per KiB of `max_file_size`, minimum 1024)
37    /// are rejected by [`FileTransferManager::chunk_file`] — a receiver
38    /// with the same configuration would drop every chunk of such a
39    /// transfer.
40    pub chunk_size: usize,
41
42    /// Maximum file size allowed (bytes).
43    pub max_file_size: u64,
44
45    /// Maximum number of inbound files that may be reassembling at once.
46    /// Chunks that would start an additional assembly are rejected; slots
47    /// free when a transfer completes, is cancelled, or goes stale.
48    pub max_concurrent_assemblies: usize,
49
50    /// Maximum number of inbound files a single sender may have reassembling
51    /// at once, so one peer cannot occupy every assembly slot. Only as
52    /// strong as sender authentication: MLS-path senders are
53    /// cryptographically bound, but legacy plaintext senders are wire
54    /// claims a peer can rotate — there, this quota bounds only honest
55    /// peers, and the global count and byte caps are the real limits.
56    pub max_assemblies_per_sender: usize,
57
58    /// Maximum total bytes buffered across all inbound assemblies. This —
59    /// not the assembly counts — is the hard bound on receive-path memory.
60    /// Each stored chunk is charged its payload plus a small fixed
61    /// bookkeeping overhead, so floods of tiny chunks cannot pin more
62    /// memory than the budget admits. Values below what a single
63    /// `max_file_size` transfer charges are clamped up to it by
64    /// [`FileTransferManager::with_config`], or a maximum-size file could
65    /// never complete.
66    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, // 100 MB
74            max_concurrent_assemblies: 32,
75            max_assemblies_per_sender: 16,
76            max_total_buffered_bytes: 128 * 1024 * 1024, // 128 MB
77        }
78    }
79}
80
81/// Why [`FileTransferManager::process_chunk`] refused a chunk.
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum ChunkRejection {
84    /// Wire fields are malformed or out of bounds: zero or oversized size
85    /// claims, empty or oversized data, or bytes exceeding the claimed
86    /// file size.
87    Invalid,
88    /// Chunk metadata does not match the existing assembly for this file id.
89    MetadataMismatch,
90    /// Chunk sender does not match the sender that started the assembly.
91    SenderMismatch,
92    /// The global concurrent-assembly cap is reached.
93    TooManyTransfers,
94    /// The sender's concurrent-assembly quota is reached.
95    SenderQuotaExceeded,
96    /// Accepting the chunk would exceed the global buffered-bytes budget.
97    /// The affected assembly is dropped along with the chunk: a dropped
98    /// chunk is never retransmitted, so the transfer could no longer
99    /// complete and its buffer would only prolong the exhaustion.
100    BufferBudgetExhausted,
101    /// The chunk belongs to a transfer that already failed (resource drop,
102    /// integrity failure, or stale sweep). Its remaining in-flight chunks
103    /// are dropped without re-creating assembly state or re-notifying the
104    /// application — `FileReceiveFailed` fires at most once per transfer.
105    /// Retries must use a fresh `file_id`; the failed id is admitted again
106    /// only after its tombstone expires (no chunks for the stale timeout).
107    PreviouslyFailed,
108}
109
110impl ChunkRejection {
111    /// `true` for rejections caused by receiver resource limits rather than
112    /// malformed input — the cases where a well-formed transfer was lost and
113    /// the application should be told.
114    pub fn is_resource_exhaustion(self) -> bool {
115        matches!(
116            self,
117            Self::TooManyTransfers | Self::SenderQuotaExceeded | Self::BufferBudgetExhausted
118        )
119    }
120
121    /// Stable machine-readable reason string, carried by the
122    /// `FileReceiveFailed` event.
123    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/// Why [`FileChunk::from_bytes`] rejected a binary payload.
137// Adding a variant to a public error enum is a breaking change without
138// this attribute; downstream crates must carry a wildcard arm.
139#[non_exhaustive]
140#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
141pub enum ChunkDecodeError {
142    /// The payload ended before the named value could be read.
143    #[error("unexpected end of data reading {what}")]
144    UnexpectedEof {
145        /// What was being read (`"u32"`, `"u64"`, or `"bytes"`).
146        what: &'static str,
147    },
148    /// A length prefix exceeds the wire-format cap for its field.
149    #[error("{field} {len} exceeds maximum {max}")]
150    FieldTooLong {
151        /// Name of the offending length field.
152        field: &'static str,
153        /// Claimed length.
154        len: usize,
155        /// Maximum allowed length.
156        max: usize,
157    },
158    /// A string field holds invalid UTF-8.
159    #[error("invalid {field} UTF-8: {source}")]
160    InvalidUtf8 {
161        /// Name of the offending string field.
162        field: &'static str,
163        /// The underlying UTF-8 decode failure.
164        source: std::string::FromUtf8Error,
165    },
166}
167
168/// A file chunk message.
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct FileChunk {
171    /// Unique file identifier.
172    pub file_id: String,
173
174    /// Original file name.
175    pub file_name: String,
176
177    /// Total file size in bytes.
178    pub file_size: u64,
179
180    /// Total number of chunks.
181    pub total_chunks: u32,
182
183    /// Index of this chunk (0-based).
184    pub chunk_index: u32,
185
186    /// Chunk data (base64-encoded in JSON).
187    #[serde(with = "base64_bytes")]
188    pub chunk_data: Vec<u8>,
189
190    /// Checksum of the complete file (SHA256 hex string).
191    pub file_checksum: String,
192}
193
194impl FileChunk {
195    /// Serializes to JSON (kept for backward compatibility).
196    pub fn to_json(&self) -> Result<String, serde_json::Error> {
197        serde_json::to_string(self)
198    }
199
200    /// Deserializes from JSON (kept for backward compatibility).
201    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
202        serde_json::from_str(json)
203    }
204
205    /// Serializes to a compact binary format, avoiding base64/JSON overhead.
206    ///
207    /// Wire format (all multi-byte integers are little-endian):
208    /// ```text
209    /// [file_id_len:4][file_id][file_name_len:4][file_name]
210    /// [file_size:8][total_chunks:4][chunk_index:4]
211    /// [chunk_data_len:4][chunk_data]
212    /// [checksum_len:4][checksum]
213    /// ```
214    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    /// Maximum allowed length for a string field (file_id, file_name, checksum)
248    /// in the binary wire format. Prevents allocation bombs from crafted
249    /// payloads. `pub(crate)`: also bounds caller-supplied file ids at the
250    /// `send_media_with` boundary.
251    pub(crate) const MAX_STRING_FIELD_LEN: usize = 4096;
252    /// Maximum allowed chunk data length in the binary wire format.
253    const MAX_CHUNK_DATA_LEN: usize = 128 * 1024 * 1024; // 128 MB
254
255    /// Deserializes from the compact binary format produced by `to_bytes`.
256    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/// File transfer progress information.
363#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct FileProgress {
365    /// File identifier.
366    pub file_id: String,
367
368    /// File name.
369    pub file_name: String,
370
371    /// Total file size.
372    pub file_size: u64,
373
374    /// Number of chunks sent/received.
375    pub chunks_completed: u32,
376
377    /// Total number of chunks.
378    pub total_chunks: u32,
379
380    /// Progress percentage (0-100).
381    pub percentage: u8,
382}
383
384impl FileProgress {
385    /// Calculates the progress percentage.
386    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/// An incomplete transfer removed by
395/// [`FileTransferManager::cleanup_stale_transfers`].
396#[derive(Debug, Clone, PartialEq, Eq)]
397pub struct StaleTransfer {
398    /// File identifier of the dropped transfer.
399    pub file_id: String,
400    /// File name from the chunk metadata.
401    pub file_name: String,
402    /// Sender the assembly was bound to.
403    pub sender: String,
404}
405
406/// Tracks file reassembly state.
407struct FileAssembly {
408    /// Sender the first accepted chunk arrived from; all later chunks must
409    /// match it.
410    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    /// Total bytes currently held in `received_chunks`.
417    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    /// Bytes this assembly charges against the global buffer budget:
445    /// payload plus a fixed per-chunk bookkeeping charge, so many tiny
446    /// chunks cannot evade the budget.
447    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(), // Will be set by caller
463            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        // Size the buffer from bytes actually received — file_size is a
480        // sender-controlled claim and must never drive an allocation.
481        let mut file_data = Vec::with_capacity(self.received_bytes as usize);
482
483        // Reassemble chunks in order
484        for i in 0..self.total_chunks {
485            let chunk_data = self.received_chunks.get(&i)?; // Missing chunk
486            file_data.extend_from_slice(chunk_data);
487        }
488
489        Some(file_data)
490    }
491}
492
493/// File transfer manager for chunking and reassembly.
494pub struct FileTransferManager {
495    config: FileTransferConfig,
496    active_assemblies: HashMap<String, FileAssembly>,
497    /// Tombstones for failed transfers, keyed by a keyed hash of the
498    /// `file_id` (never the id itself, so attacker-length ids cannot grow
499    /// this map) mapping to the last time a chunk for that id was seen.
500    /// While tombstoned, an id's chunks are dropped as `PreviouslyFailed` —
501    /// they were already ACKed and will keep streaming in, but must neither
502    /// resurrect a partial assembly nor re-emit the failure event.
503    failed_transfers: HashMap<u64, Instant>,
504    /// Insertion order of `failed_transfers` keys, for FIFO eviction at
505    /// [`Self::TOMBSTONE_CAPACITY`].
506    failed_transfer_order: VecDeque<u64>,
507    /// Randomly keyed hasher for tombstone keys: collisions (a fresh
508    /// transfer misread as previously failed) cannot be crafted by a peer
509    /// and are a ~2⁻⁵⁴ accident at capacity.
510    tombstone_hasher: RandomState,
511}
512
513impl FileTransferManager {
514    /// Floor on the average chunk payload assumed when bounding a sender's
515    /// `total_chunks` claim, so per-assembly bookkeeping stays proportional
516    /// to `max_file_size` even when every chunk is tiny. The SDK's own
517    /// senders never chunk below `CHUNK_SIZE_BLE` (4 KiB).
518    const MIN_ACCEPTED_CHUNK_PAYLOAD: u64 = 1024;
519
520    /// Lower bound on the derived `total_chunks` cap so configurations with
521    /// a small `max_file_size` still admit fine-grained chunking.
522    const MIN_TOTAL_CHUNKS_CAP: u64 = 1024;
523
524    /// Flat per-stored-chunk charge against `max_total_buffered_bytes`,
525    /// approximating the bookkeeping memory the payload sum does not see
526    /// (hashmap entry, allocator rounding). Without it, a flood of 1-byte
527    /// chunks pins far more memory than the budget records.
528    const PER_CHUNK_OVERHEAD: u64 = 64;
529
530    /// Pseudo-sender bound to assemblies fed through the manual
531    /// `process_file_chunk` FFI path, which carries no wire sender. All
532    /// manual chunks share this identity for sender binding and quota. The
533    /// `:` is rejected by wire `UserId` validation, so no remote peer can
534    /// ever claim this identity.
535    pub const MANUAL_SENDER: &'static str = "manual:ffi";
536
537    /// Most failed-transfer tombstones retained at once. Each costs ~50
538    /// bytes (hash key, timestamp, deque slot), so the set is capped around
539    /// 50 KiB. Evicting a tombstone early only means one extra
540    /// `FileReceiveFailed` if that transfer's chunks are still arriving —
541    /// filling the set costs an attacker one rejected transfer per slot.
542    const TOMBSTONE_CAPACITY: usize = 1024;
543
544    /// Creates a new file transfer manager.
545    pub fn new() -> Self {
546        Self::with_config(FileTransferConfig::default())
547    }
548
549    /// Creates a new file transfer manager with custom configuration.
550    ///
551    /// `max_total_buffered_bytes` is clamped up to what a single
552    /// `max_file_size` transfer charges (payload plus per-chunk overhead) —
553    /// any lower value would make `max_file_size` unachievable.
554    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    /// Keyed hash under which a `file_id` is tombstoned. Hashing bounds the
574    /// tombstone set's memory regardless of how long a peer makes its ids.
575    fn tombstone_key(&self, file_id: &str) -> u64 {
576        self.tombstone_hasher.hash_one(file_id)
577    }
578
579    /// Marks a transfer as failed so its remaining in-flight chunks are
580    /// dropped as [`ChunkRejection::PreviouslyFailed`] instead of
581    /// re-creating assembly state or re-notifying the application.
582    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    /// Whether `file_id` belongs to a failed transfer. A hit refreshes the
598    /// tombstone's timestamp, so it outlives the transfer's chunk stream
599    /// and expires only once chunks stop arriving for the stale timeout.
600    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    /// Upper bound on a sender's `total_chunks` claim, derived from
612    /// `max_file_size`. Floor division plus one — always at least the exact
613    /// ceiling — without u64::div_ceil (MSRV) or overflow.
614    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    /// The smallest budget that still admits one maximum-size transfer:
621    /// its payload plus the per-chunk overhead of the most chunks such a
622    /// transfer may legally claim (bounded by both the derived cap and the
623    /// one-byte-per-chunk minimum).
624    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    /// Chunks a file for sending.
633    ///
634    /// # Arguments
635    ///
636    /// * `file_id` - Unique identifier for this file transfer
637    /// * `file_name` - Name of the file
638    /// * `file_data` - Complete file data
639    /// * `chunk_size_override` - If `Some`, uses this chunk size instead of the configured default.
640    ///   Allows callers to select a transport-appropriate chunk size.
641    ///
642    /// # Returns
643    ///
644    /// Returns a vector of FileChunk ready to send. Fails if the file is
645    /// empty or exceeds `max_file_size`, if the chunk size is zero, or if
646    /// the chunking is so fine that the chunk count would exceed the
647    /// receive-side `total_chunks` cap — [`Self::process_chunk`] on a
648    /// receiver with the same `max_file_size` would silently drop every
649    /// chunk of such a transfer.
650    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        // Mirror the receive-side cap: process_chunk rejects total_chunks
678        // above this, so a finer chunking would have every chunk silently
679        // dropped by any receiver with the same max_file_size.
680        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    /// Processes a received file chunk from `sender`.
715    ///
716    /// All wire-supplied fields (`file_size`, `total_chunks`, `chunk_data`,
717    /// `file_id`) are bounded here before any assembly state is created, so
718    /// a remote peer cannot force oversized allocations or unbounded
719    /// bookkeeping. Assemblies are bound to the sender that started them,
720    /// counted against global and per-sender caps, and the bytes buffered
721    /// across all assemblies never exceed `max_total_buffered_bytes`.
722    ///
723    /// # Arguments
724    ///
725    /// * `sender` - The authenticated (or wire-claimed) sender of the chunk
726    /// * `chunk` - The received file chunk
727    ///
728    /// # Returns
729    ///
730    /// Returns `Ok(FileProgress)` with current progress, or the
731    /// [`ChunkRejection`] explaining why the chunk was dropped.
732    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        // A failed transfer's chunks were already ACKed and keep streaming
776        // in; drop them without state or (further) events. Logged at debug
777        // — a large drained transfer produces thousands of these.
778        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        // A chunk for an existing assembly must be consistent with it; a
789        // chunk starting a new assembly must fit within the count caps.
790        // Everything is checked before any state is created or mutated.
791        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                // Duplicates replace their earlier copy, so discount it
812                // before checking the cumulative received bytes against the
813                // claimed size.
814                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                // A replaced duplicate refunds its full charge — payload and
831                // overhead — since its map entry is reused, not added.
832                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                    // The rejected chunk was already ACKed and is lost for
843                    // good, so the whole transfer is unrecoverable —
844                    // tombstone it so its remaining chunks drop silently.
845                    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        // Hard memory bound: total bytes charged across every assembly,
869        // where each stored chunk costs its payload plus PER_CHUNK_OVERHEAD.
870        // Iterating is fine — the count caps above keep this map small.
871        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            // The dropped chunk is never retransmitted, so this transfer can
887            // no longer complete — free its partial buffer now instead of
888            // letting it squat the budget until the stale sweep, and
889            // tombstone the id so later chunks cannot resurrect it as a
890            // never-completable assembly.
891            self.active_assemblies.remove(&file_id);
892            self.tombstone(&file_id);
893            return Err(ChunkRejection::BufferBudgetExhausted);
894        }
895
896        // Get or create assembly for this file
897        let assembly = self
898            .active_assemblies
899            .entry(file_id.clone())
900            .or_insert_with(|| FileAssembly::new(sender, &chunk));
901
902        // Add chunk
903        assembly.add_chunk(chunk.chunk_index, chunk.chunk_data);
904
905        // Get progress
906        let mut progress = assembly.progress();
907        progress.file_id = file_id;
908
909        Ok(progress)
910    }
911
912    /// Checks if a file transfer is complete.
913    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    /// Finalizes a file transfer and returns the complete file data.
921    ///
922    /// Verifies the SHA256 checksum of the reassembled file before returning.
923    ///
924    /// # Arguments
925    ///
926    /// * `file_id` - File identifier
927    ///
928    /// # Returns
929    ///
930    /// Returns `Some(Vec<u8>)` if file is complete, its size matches the
931    /// claimed `file_size`, and the checksum matches, `None` otherwise.
932    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        // A completed transfer must deliver exactly the bytes it claimed.
937        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        // Verify checksum before returning
950        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    /// Removes stale/incomplete transfers that have not received any chunk
968    /// updates within `max_age`, returning what was dropped so callers can
969    /// surface the loss. Swept ids are tombstoned so straggler chunks
970    /// cannot resurrect them; tombstones themselves expire by the same
971    /// rule — no chunks for `max_age` — freeing the id for reuse.
972    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    /// Gets the progress of an active file transfer.
1007    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    /// Cancels an active file transfer.
1016    pub fn cancel_transfer(&mut self, file_id: &str) -> bool {
1017        self.active_assemblies.remove(file_id).is_some()
1018    }
1019
1020    /// Refuses a transfer outright: releases anything already buffered for
1021    /// it and drops every further chunk.
1022    ///
1023    /// Stronger than [`Self::cancel_transfer`], and the difference is the
1024    /// tombstone. A caller refuses a transfer because of something it
1025    /// learned from one chunk, and chunks do not have to arrive in order, so
1026    /// cancelling alone would let the ones behind it rebuild the assembly it
1027    /// just released. The refusal has to outlive the chunk that prompted it.
1028    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    /// Gets the number of active file transfers.
1034    pub fn active_transfer_count(&self) -> usize {
1035        self.active_assemblies.len()
1036    }
1037
1038    /// Rewinds every active assembly's last-update time and every tombstone
1039    /// (test-only), so staleness paths can be exercised without sleeping.
1040    #[cfg(test)]
1041    pub(crate) fn backdate_transfers(&mut self, by: Duration) {
1042        // checked_sub instead of `-=`: Instant underflows (and panics) when
1043        // the monotonic clock's epoch is closer than `by` — a fresh CI VM.
1044        // Failing loudly beats an intermittent subtraction panic.
1045        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
1063/// Tracks the sliding-window state for a single outbound file transfer.
1064///
1065/// Instead of sending all chunks at once, this limits the number of in-flight
1066/// (unACKed) chunks to `max_in_flight`, reducing outbox pressure and providing
1067/// backpressure that adapts to the transport's throughput.
1068pub 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    /// Creates a new outbound transfer state with the given chunks and window size.
1078    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    /// Returns the next batch of chunks that can be sent without exceeding the
1089    /// in-flight window. Returned chunks are marked as in-flight.
1090    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    /// Records that a chunk was ACKed. Returns `true` if this was an in-flight
1102    /// chunk (i.e. the window has room for more).
1103    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    /// Records that a chunk failed terminally and should not be retried
1109    /// at the window level (the transfer will be aborted by the caller).
1110    pub fn on_chunk_failed(&mut self, chunk_index: u32) {
1111        self.in_flight.remove(&chunk_index);
1112    }
1113
1114    /// Returns the total number of chunks in this transfer.
1115    pub fn total_chunks(&self) -> u32 {
1116        self.chunks.len() as u32
1117    }
1118
1119    /// Returns the number of chunks that have been acknowledged.
1120    pub fn acked_count(&self) -> u32 {
1121        self.acked.len() as u32
1122    }
1123
1124    /// Returns `true` when all chunks have been acknowledged.
1125    pub fn is_fully_acked(&self) -> bool {
1126        self.acked.len() == self.chunks.len()
1127    }
1128
1129    /// Returns `true` when all chunks have been sent at least once
1130    /// (they may still be in flight awaiting ACK).
1131    pub fn all_sent(&self) -> bool {
1132        self.next_unsent >= self.chunks.len()
1133    }
1134
1135    /// Returns `true` when the window has capacity for more chunks.
1136    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        // Truncate mid-header: the first u32 (file_id_len) can't be read.
1161        let err = FileChunk::from_bytes(&bytes[..2]).unwrap_err();
1162        assert!(matches!(
1163            err,
1164            ChunkDecodeError::UnexpectedEof { what: "u32" }
1165        ));
1166        // Truncate mid-header: both string fields parse but the file_size
1167        // u64 read hits EOF.
1168        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        // Truncate mid-body: length prefixes parse but the data falls short.
1175        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        // Claim a file_id longer than the wire-format cap.
1186        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        // file_id content starts right after its 4-byte length prefix.
1202        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); // Small file = 1 chunk
1228        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, // Small chunk for testing
1238            max_file_size: 1024 * 1024,
1239            ..FileTransferConfig::default()
1240        };
1241        let manager = FileTransferManager::with_config(config);
1242
1243        let file_data = vec![0u8; 350]; // 350 bytes
1244
1245        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); // 350 bytes / 100 = 4 chunks
1250        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]; // Exceeds max
1266
1267        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        // Process all chunks
1315        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        // Reverse order
1345        chunks.reverse();
1346
1347        // Process all chunks
1348        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]; // 3 chunks
1368        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        // Process first chunk
1375        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); // 1/3 ≈ 33%
1379
1380        // Process second chunk
1381        let progress = manager.process_chunk("peer", chunks[1].clone()).unwrap();
1382        assert_eq!(progress.chunks_completed, 2);
1383        assert_eq!(progress.percentage, 66); // 2/3 ≈ 66%
1384
1385        // Process third chunk
1386        let progress = manager.process_chunk("peer", chunks[2].clone()).unwrap();
1387        assert_eq!(progress.chunks_completed, 3);
1388        assert_eq!(progress.percentage, 100); // 3/3 = 100%
1389
1390        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        // Process same chunk twice
1403        manager.process_chunk("peer", chunks[0].clone()).unwrap();
1404        manager.process_chunk("peer", chunks[0].clone()).unwrap();
1405
1406        // Should still show only 1 chunk
1407        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        // Process both files
1441        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        // SEC-H2 regression: a single chunk claiming u64::MAX bytes used to
1541        // reach reassemble() and panic on Vec::with_capacity.
1542        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        // More chunks than bytes is inherently bogus (chunks are non-empty).
1573        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        // Default config: 100 MB / 1 KiB payload floor ≈ 102_400 max chunks.
1581        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        // Claim 2048 bytes in 2 chunks, then send two 1500-byte chunks.
1604        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        // The first chunk's state survives; only the offending chunk drops.
1609        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        // 1500 + 548 = 2048 claimed bytes; resending chunk 0 must replace
1617        // (not add), so the transfer still completes.
1618        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        // Third simultaneous transfer is rejected...
1658        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        // ...but chunks for existing assemblies still flow.
1665        assert!(manager
1666            .process_chunk("peer", attack_chunk("f1", 100, 2, 1, vec![0u8; 10]))
1667            .is_ok());
1668        // Freeing a slot re-admits new transfers — but not the failed id:
1669        // f3's rejected chunk was ACKed and lost, so the transfer could
1670        // never complete. A retry must use a fresh file id.
1671        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        // A second transfer from the same sender exceeds the quota...
1696        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        // ...but other senders are unaffected.
1703        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        // Same file_id and metadata from a different sender must not be able
1715        // to poison the assembly.
1716        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        // The original assembly is untouched.
1723        assert_eq!(manager.get_progress("f").unwrap().chunks_completed, 1);
1724    }
1725
1726    #[test]
1727    fn test_buffer_budget_rejects_new_transfer() {
1728        // 1_065_536 is exactly the minimum viable budget for a 1 MB
1729        // max_file_size, so with_config leaves it unclamped.
1730        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        // 600_064 + 500_064 charged would exceed the budget; the new
1744        // transfer is rejected before any of its state is created.
1745        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        // The pre-existing transfer is untouched.
1757        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        // f1's next chunk busts the budget. The dropped chunk is never
1779        // retransmitted, so f1 can no longer complete — its partial buffer
1780        // is freed immediately rather than squatting the budget until the
1781        // stale sweep.
1782        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        // The freed budget re-admits other traffic.
1794        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        // A flood of 1-byte chunks must be stopped by its bookkeeping
1802        // charge, not its (negligible) payload sum. Budget 0 is clamped up
1803        // to the minimum viable value: 10_000 + 1024 * 64 = 75_536.
1804        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        // 1024 chunks is the derived cap for this config; f1 fills it.
1812        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        // f2's identical flood must hit the budget long before its payload
1819        // (a few hundred bytes) comes anywhere near 75_536.
1820        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        // The busting transfer is dropped with its chunk.
1838        assert!(manager.get_progress("f2").is_none());
1839    }
1840
1841    #[test]
1842    fn test_manual_sender_cannot_be_a_wire_user_id() {
1843        // Sender binding for the manual FFI path relies on this sentinel
1844        // being unrepresentable as a wire UserId (':' is rejected), so no
1845        // remote peer can alias it.
1846        assert!(offline_protocol_core::UserId::new(FileTransferManager::MANUAL_SENDER).is_err());
1847    }
1848
1849    #[test]
1850    fn test_resource_rejected_transfer_is_tombstoned() {
1851        // Once a transfer is refused for a resource reason, its remaining
1852        // in-flight chunks are dropped as previously_failed — no repeated
1853        // resource rejections (which would re-emit FileReceiveFailed
1854        // upstream) and no ghost assembly, even after capacity frees up.
1855        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        // Even with the slot freed, the failed id stays dead: its first
1879        // chunk was ACKed and lost, so the transfer could never complete.
1880        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        // f1's remaining in-flight chunk must not start a fresh assembly —
1918        // it would miss the dropped chunk forever, squatting a slot and the
1919        // budget until the stale sweep and then failing a second time.
1920        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        // Checksum matches the bytes but the claimed file_size does not, so
1938        // finalize fails its size check and drops the transfer.
1939        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        // The failed id cannot be restarted by resending its chunks.
1955        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        // A straggler chunk must not resurrect the swept transfer — the
1974        // caller already emitted its stale_timeout failure.
1975        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        // Once no chunk for the failed id has been seen for max_age, the
2004        // stale sweep releases it and the id may be reused (the manual FFI
2005        // path retries with app-chosen ids).
2006        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        // Cancellation is an app decision, not a failure — the id may be
2016        // restarted immediately (nothing was lost that a full resend
2017        // cannot supply).
2018        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        // A duplicate whose payload differs in size must fully replace the
2032        // earlier copy in both bytes and budget accounting, so the transfer
2033        // still completes and delivers exactly the claimed size.
2034        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        // Undersized garbage first, then the real 1500-byte chunk 0.
2052        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        // finalize checks reassembled length against the claimed file_size,
2059        // so this also proves the replacement byte accounting is exact.
2060        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        // The configured budget equals max_file_size, which cannot cover
2085        // the per-chunk overhead — with_config clamps it up so a
2086        // maximum-size file still completes.
2087        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        // A 1 MiB max_file_size derives a receive-side cap of 1025 chunks;
2113        // chunking a full-size file at 512 bytes would need 2048. Every
2114        // chunk of such a transfer would be ACKed then dropped by the
2115        // receiver, so the send side must refuse to produce it.
2116        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        // Small chunk sizes stay valid while the chunk count fits the
2134        // receive-side cap: 400 KiB at 512 bytes is 800 chunks (cap 1025),
2135        // and a same-config receiver accepts the whole transfer.
2136        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        // max_file_size 1 MiB → cap = 1 MiB / 1 KiB + 1 = 1025.
2161        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        // A resent chunk 0 with different bytes replaces the first copy; the
2187        // final file must contain the replacement.
2188        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        // Single 100-byte chunk with a valid checksum of those bytes, but a
2215        // claimed file_size of 2048.
2216        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        // The failed transfer is discarded.
2231        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}