Skip to main content

rns_core/resource/
sender.rs

1use alloc::vec;
2use alloc::vec::Vec;
3
4use rns_crypto::Rng;
5
6use super::advertisement::ResourceAdvertisement;
7use super::parts::{build_hashmap, has_collision, prepend_metadata, split_into_parts};
8use super::proof::{compute_expected_proof, compute_resource_hash, validate_proof};
9use super::types::*;
10use crate::buffer::types::Compressor;
11use crate::constants::*;
12use crate::hash::get_random_hash;
13
14/// Resource sender state machine.
15///
16/// Creates an advertisement, handles part requests, sends parts, and validates proofs.
17/// Returns `Vec<ResourceAction>` — no I/O, no callbacks.
18pub struct ResourceSender {
19    /// Current status
20    pub status: ResourceStatus,
21    /// Resource hash (SHA-256 of unencrypted data + random_hash), 32 bytes
22    pub resource_hash: [u8; 32],
23    /// Truncated hash (first 16 bytes of resource_hash)
24    pub truncated_hash: [u8; 16],
25    /// Expected proof (SHA-256 of unencrypted data + resource_hash)
26    pub expected_proof: [u8; 32],
27    /// Original hash (for multi-segment, first segment's hash)
28    pub original_hash: [u8; 32],
29    /// Random hash for map hashing (4 bytes)
30    pub random_hash: Vec<u8>,
31    /// SDU size
32    pub sdu: usize,
33    /// Encrypted parts data
34    parts: Vec<Vec<u8>>,
35    /// Part map hashes (4 bytes each)
36    pub part_hashes: Vec<[u8; RESOURCE_MAPHASH_LEN]>,
37    /// Concatenated hashmap bytes
38    hashmap: Vec<u8>,
39    /// Number of parts
40    total_parts: usize,
41    /// Number of unique parts sent
42    pub sent_parts: usize,
43    /// Tracks which part indices have been sent (for dedup)
44    sent_indices: Vec<bool>,
45    /// Flags
46    pub flags: AdvFlags,
47    /// Transfer size (encrypted data size)
48    pub transfer_size: usize,
49    /// Total uncompressed data size
50    pub data_size: usize,
51    /// Segment index (1-based)
52    pub segment_index: u64,
53    /// Total segments
54    pub total_segments: u64,
55    /// Request ID
56    pub request_id: Option<Vec<u8>>,
57    /// Retries left
58    pub retries_left: usize,
59    /// Max retries
60    pub max_retries: usize,
61    /// Max advertisement retries
62    pub max_adv_retries: usize,
63    /// RTT estimate (seconds)
64    pub rtt: Option<f64>,
65    /// Link RTT estimate (from link establishment)
66    pub link_rtt: f64,
67    /// Timeout factor
68    pub timeout_factor: f64,
69    /// Last activity timestamp
70    pub last_activity: f64,
71    /// Advertisement sent timestamp
72    pub adv_sent: f64,
73    /// Last part sent timestamp
74    pub last_part_sent: f64,
75    /// Sender grace time
76    pub sender_grace_time: f64,
77    /// Receiver min consecutive height (for search optimization)
78    receiver_min_consecutive_height: usize,
79}
80
81impl ResourceSender {
82    /// Create a new ResourceSender from unencrypted data.
83    ///
84    /// - `data`: raw application data (no metadata prefix)
85    /// - `metadata`: optional pre-serialized metadata bytes
86    /// - `sdu`: SDU size (usually RESOURCE_SDU = 464)
87    /// - `encrypt_fn`: closure to encrypt the full data blob
88    /// - `compressor`: Compressor trait for optional compression
89    /// - `rng`: random number generator
90    /// - `now`: current timestamp
91    /// - `auto_compress`: whether to attempt compression
92    /// - `is_response`: whether this is a response to a request
93    /// - `request_id`: optional request ID
94    /// - `segment_index`: 1-based segment number (1 for single-segment)
95    /// - `total_segments`: total number of segments
96    /// - `original_hash`: original hash from first segment (None for first segment)
97    /// - `link_rtt`: current link RTT estimate
98    /// - `traffic_timeout_factor`: link traffic timeout factor
99    #[allow(clippy::too_many_arguments)]
100    pub fn new(
101        data: &[u8],
102        metadata: Option<&[u8]>,
103        sdu: usize,
104        encrypt_fn: &dyn Fn(&[u8]) -> Vec<u8>,
105        compressor: &dyn Compressor,
106        rng: &mut dyn Rng,
107        now: f64,
108        auto_compress: bool,
109        is_response: bool,
110        request_id: Option<Vec<u8>>,
111        segment_index: u64,
112        total_segments: u64,
113        original_hash: Option<[u8; 32]>,
114        link_rtt: f64,
115        traffic_timeout_factor: f64,
116    ) -> Result<Self, ResourceError> {
117        // Build unencrypted data (metadata prefix + data)
118        let uncompressed_data = match metadata {
119            Some(meta) => prepend_metadata(data, meta),
120            None => data.to_vec(),
121        };
122        let has_metadata = metadata.is_some();
123
124        let data_size = uncompressed_data.len();
125
126        // Try compression
127        let (working_data, compressed) =
128            if auto_compress && uncompressed_data.len() <= RESOURCE_AUTO_COMPRESS_MAX_SIZE {
129                match compressor.compress(&uncompressed_data) {
130                    Some(compressed_data) if compressed_data.len() < uncompressed_data.len() => {
131                        (compressed_data, true)
132                    }
133                    _ => (uncompressed_data.clone(), false),
134                }
135            } else {
136                (uncompressed_data.clone(), false)
137            };
138
139        // Prepend random hash (4 bytes)
140        let random_prefix: [u8; RESOURCE_RANDOM_HASH_SIZE] = {
141            let rh = get_random_hash(rng);
142            let mut buf = [0u8; RESOURCE_RANDOM_HASH_SIZE];
143            buf.copy_from_slice(&rh[..RESOURCE_RANDOM_HASH_SIZE]);
144            buf
145        };
146        let mut data_with_random =
147            Vec::with_capacity(RESOURCE_RANDOM_HASH_SIZE + working_data.len());
148        data_with_random.extend_from_slice(&random_prefix);
149        data_with_random.extend_from_slice(&working_data);
150
151        // Encrypt
152        let encrypted_data = encrypt_fn(&data_with_random);
153        let transfer_size = encrypted_data.len();
154
155        // Keep trying until no collision in hashmap (max 100 attempts)
156        let mut resource_hash;
157        let mut truncated_resource_hash;
158        let mut expected_proof;
159        let mut final_random_hash;
160        let mut parts_data;
161        let mut part_hashes;
162        let mut collision_retries = 0;
163        const MAX_COLLISION_RETRIES: usize = 100;
164
165        loop {
166            final_random_hash = {
167                let rh = get_random_hash(rng);
168                rh[..RESOURCE_RANDOM_HASH_SIZE].to_vec()
169            };
170
171            resource_hash = compute_resource_hash(&uncompressed_data, &final_random_hash);
172            truncated_resource_hash = {
173                let mut t = [0u8; 16];
174                t.copy_from_slice(&resource_hash[..16]);
175                t
176            };
177            expected_proof = compute_expected_proof(&uncompressed_data, &resource_hash);
178
179            let (p, h) = split_into_parts(&encrypted_data, sdu, &final_random_hash);
180            parts_data = p;
181            part_hashes = h;
182
183            if !has_collision(&part_hashes) {
184                break;
185            }
186            // Collision detected, retry with new random hash
187            collision_retries += 1;
188            if collision_retries >= MAX_COLLISION_RETRIES {
189                return Err(ResourceError::CollisionDetected);
190            }
191        }
192
193        let hashmap = build_hashmap(&part_hashes);
194        let total_parts = parts_data.len();
195
196        let orig_hash = original_hash.unwrap_or(resource_hash);
197
198        let flags = AdvFlags {
199            encrypted: true,
200            compressed,
201            split: total_segments > 1,
202            is_request: request_id.is_some() && !is_response,
203            is_response: request_id.is_some() && is_response,
204            has_metadata,
205        };
206
207        Ok(ResourceSender {
208            status: ResourceStatus::Queued,
209            resource_hash,
210            truncated_hash: truncated_resource_hash,
211            expected_proof,
212            original_hash: orig_hash,
213            random_hash: final_random_hash,
214            sdu,
215            parts: parts_data,
216            part_hashes,
217            hashmap,
218            total_parts,
219            sent_parts: 0,
220            sent_indices: vec![false; total_parts],
221            flags,
222            transfer_size,
223            data_size,
224            segment_index,
225            total_segments,
226            request_id,
227            retries_left: RESOURCE_MAX_RETRIES,
228            max_retries: RESOURCE_MAX_RETRIES,
229            max_adv_retries: RESOURCE_MAX_ADV_RETRIES,
230            rtt: None,
231            link_rtt,
232            timeout_factor: traffic_timeout_factor,
233            last_activity: now,
234            adv_sent: now,
235            last_part_sent: now,
236            sender_grace_time: RESOURCE_SENDER_GRACE_TIME,
237            receiver_min_consecutive_height: 0,
238        })
239    }
240
241    /// Generate the advertisement for the given hashmap segment.
242    pub fn get_advertisement(&self, segment: usize) -> Vec<u8> {
243        let adv = ResourceAdvertisement {
244            transfer_size: self.transfer_size as u64,
245            data_size: self.data_size as u64,
246            num_parts: self.total_parts as u64,
247            resource_hash: self.resource_hash.to_vec(),
248            random_hash: self.random_hash.clone(),
249            original_hash: self.original_hash.to_vec(),
250            hashmap: self.hashmap.clone(),
251            flags: self.flags,
252            segment_index: self.segment_index,
253            total_segments: self.total_segments,
254            request_id: self.request_id.clone(),
255        };
256        adv.pack(segment)
257    }
258
259    /// Advertise the resource. Returns SendAdvertisement action.
260    pub fn advertise(&mut self, now: f64) -> Vec<ResourceAction> {
261        self.status = ResourceStatus::Advertised;
262        self.last_activity = now;
263        self.adv_sent = now;
264        self.retries_left = self.max_adv_retries;
265        let adv_data = self.get_advertisement(0);
266        vec![ResourceAction::SendAdvertisement(adv_data)]
267    }
268
269    /// Handle a request for parts (RESOURCE_REQ context).
270    ///
271    /// request_data format:
272    /// [exhausted_flag: u8][last_map_hash: 4 bytes if exhausted][resource_hash: 32 bytes][requested_hashes: N*4 bytes]
273    pub fn handle_request(&mut self, request_data: &[u8], now: f64) -> Vec<ResourceAction> {
274        if self.status == ResourceStatus::Failed {
275            return vec![];
276        }
277
278        let wants_more_hashmap = request_data.first() == Some(&RESOURCE_HASHMAP_IS_EXHAUSTED);
279        let pad = if wants_more_hashmap {
280            1 + RESOURCE_MAPHASH_LEN
281        } else {
282            1
283        };
284
285        if request_data.len() < pad + 32 || request_data[pad..pad + 32] != self.resource_hash {
286            return vec![];
287        }
288
289        // Measure RTT from advertisement
290        if self.rtt.is_none() {
291            self.rtt = Some(now - self.adv_sent);
292        }
293
294        if self.status != ResourceStatus::Transferring {
295            self.status = ResourceStatus::Transferring;
296        }
297
298        self.retries_left = self.max_retries;
299        self.last_activity = now;
300
301        let requested_hashes_data = &request_data[pad + 32..];
302        let mut actions = Vec::new();
303
304        // Parse requested map hashes
305        let num_requested = requested_hashes_data.len() / RESOURCE_MAPHASH_LEN;
306        let mut map_hashes_requested = Vec::with_capacity(num_requested);
307        for i in 0..num_requested {
308            let start = i * RESOURCE_MAPHASH_LEN;
309            let end = start + RESOURCE_MAPHASH_LEN;
310            if end <= requested_hashes_data.len() {
311                let mut h = [0u8; RESOURCE_MAPHASH_LEN];
312                h.copy_from_slice(&requested_hashes_data[start..end]);
313                map_hashes_requested.push(h);
314            }
315        }
316
317        // Search for requested parts within guard window
318        let search_start = self.receiver_min_consecutive_height;
319        let search_end = core::cmp::min(
320            search_start + RESOURCE_COLLISION_GUARD_SIZE,
321            self.total_parts,
322        );
323
324        for part_idx in search_start..search_end {
325            if map_hashes_requested.contains(&self.part_hashes[part_idx]) {
326                actions.push(ResourceAction::SendPart(self.parts[part_idx].clone()));
327                if !self.sent_indices[part_idx] {
328                    self.sent_indices[part_idx] = true;
329                    self.sent_parts += 1;
330                }
331                self.last_part_sent = now;
332            }
333        }
334
335        // Handle hashmap exhaustion
336        if wants_more_hashmap {
337            match self.build_hmu(request_data, now) {
338                Ok(hmu) => actions.push(ResourceAction::SendHmu(hmu)),
339                Err(_) => {
340                    actions.extend(self.cancel());
341                    return actions;
342                }
343            }
344        }
345
346        // Check if all parts sent
347        if self.sent_parts >= self.total_parts {
348            self.status = ResourceStatus::AwaitingProof;
349            self.retries_left = 3; // hardcoded in Python
350        }
351
352        actions
353    }
354
355    /// Build hashmap update data.
356    fn build_hmu(&mut self, request_data: &[u8], now: f64) -> Result<Vec<u8>, ResourceError> {
357        if request_data.len() < 1 + RESOURCE_MAPHASH_LEN {
358            return Err(ResourceError::InvalidState);
359        }
360
361        let last_map_hash_bytes = &request_data[1..1 + RESOURCE_MAPHASH_LEN];
362        let mut last_map_hash = [0u8; RESOURCE_MAPHASH_LEN];
363        last_map_hash.copy_from_slice(last_map_hash_bytes);
364
365        // Find the part index of the last map hash
366        let search_start = self.receiver_min_consecutive_height;
367        let search_end = core::cmp::min(
368            search_start + RESOURCE_COLLISION_GUARD_SIZE,
369            self.total_parts,
370        );
371
372        let mut part_index = search_start;
373        for idx in search_start..search_end {
374            part_index = idx + 1;
375            if self.part_hashes[idx] == last_map_hash {
376                break;
377            }
378        }
379
380        // Update receiver min consecutive height
381        self.receiver_min_consecutive_height = if part_index > RESOURCE_WINDOW_MAX {
382            part_index - 1 - RESOURCE_WINDOW_MAX
383        } else {
384            0
385        };
386
387        // Verify alignment
388        if !part_index.is_multiple_of(RESOURCE_HASHMAP_MAX_LEN) {
389            return Err(ResourceError::InvalidState);
390        }
391
392        let segment = part_index / RESOURCE_HASHMAP_MAX_LEN;
393        let hashmap_start = segment * RESOURCE_HASHMAP_MAX_LEN;
394        let hashmap_end =
395            core::cmp::min((segment + 1) * RESOURCE_HASHMAP_MAX_LEN, self.total_parts);
396
397        let mut hashmap_segment = Vec::new();
398        for i in hashmap_start..hashmap_end {
399            hashmap_segment.extend_from_slice(
400                &self.hashmap[i * RESOURCE_MAPHASH_LEN..(i + 1) * RESOURCE_MAPHASH_LEN],
401            );
402        }
403
404        if hashmap_segment.is_empty() {
405            return Err(ResourceError::InvalidState);
406        }
407
408        // Build HMU: resource_hash + msgpack([segment, hashmap])
409        let hmu_payload = crate::msgpack::pack(&crate::msgpack::Value::Array(vec![
410            crate::msgpack::Value::UInt(segment as u64),
411            crate::msgpack::Value::Bin(hashmap_segment),
412        ]));
413
414        let mut hmu = Vec::with_capacity(32 + hmu_payload.len());
415        hmu.extend_from_slice(&self.resource_hash);
416        hmu.extend_from_slice(&hmu_payload);
417
418        self.last_activity = now;
419        Ok(hmu)
420    }
421
422    /// Handle proof from receiver.
423    pub fn handle_proof(&mut self, proof_data: &[u8], _now: f64) -> Vec<ResourceAction> {
424        if self.status == ResourceStatus::Failed {
425            return vec![];
426        }
427
428        match validate_proof(proof_data, &self.resource_hash, &self.expected_proof) {
429            Ok(true) => {
430                self.status = ResourceStatus::Complete;
431                vec![ResourceAction::Completed]
432            }
433            Ok(false) => {
434                self.status = ResourceStatus::Failed;
435                vec![ResourceAction::Failed(ResourceError::InvalidProof)]
436            }
437            Err(e) => {
438                self.status = ResourceStatus::Failed;
439                vec![ResourceAction::Failed(e)]
440            }
441        }
442    }
443
444    /// Handle rejection from receiver.
445    pub fn handle_reject(&mut self) -> Vec<ResourceAction> {
446        self.status = ResourceStatus::Rejected;
447        vec![ResourceAction::Failed(ResourceError::Rejected)]
448    }
449
450    /// Cancel the transfer.
451    pub fn cancel(&mut self) -> Vec<ResourceAction> {
452        if self.status < ResourceStatus::Complete {
453            self.status = ResourceStatus::Failed;
454            vec![ResourceAction::SendCancelInitiator(
455                self.resource_hash.to_vec(),
456            )]
457        } else {
458            vec![]
459        }
460    }
461
462    /// Periodic tick. Checks for timeouts.
463    pub fn tick(&mut self, now: f64) -> Vec<ResourceAction> {
464        if self.status >= ResourceStatus::Complete {
465            return vec![];
466        }
467
468        match self.status {
469            ResourceStatus::Advertised => {
470                let timeout = self.adv_sent
471                    + self.rtt.unwrap_or(self.link_rtt * self.timeout_factor)
472                    + RESOURCE_PROCESSING_GRACE;
473                if now > timeout {
474                    if self.retries_left == 0 {
475                        self.status = ResourceStatus::Failed;
476                        return vec![ResourceAction::Failed(ResourceError::Timeout)];
477                    }
478                    self.retries_left -= 1;
479                    self.last_activity = now;
480                    self.adv_sent = now;
481                    let adv_data = self.get_advertisement(0);
482                    return vec![ResourceAction::SendAdvertisement(adv_data)];
483                }
484            }
485            ResourceStatus::Transferring => {
486                let rtt = self.rtt.unwrap_or(1.0);
487                let max_extra_wait: f64 = (0..self.max_retries)
488                    .map(|r| (r as f64 + 1.0) * RESOURCE_PER_RETRY_DELAY)
489                    .sum();
490                let max_wait = rtt * self.timeout_factor * self.max_retries as f64
491                    + self.sender_grace_time
492                    + max_extra_wait;
493                if now > self.last_activity + max_wait {
494                    self.status = ResourceStatus::Failed;
495                    return vec![ResourceAction::Failed(ResourceError::Timeout)];
496                }
497            }
498            ResourceStatus::AwaitingProof => {
499                let rtt = self.rtt.unwrap_or(1.0);
500                let timeout = self.last_part_sent
501                    + rtt * RESOURCE_PROOF_TIMEOUT_FACTOR
502                    + self.sender_grace_time;
503                if now > timeout {
504                    if self.retries_left == 0 {
505                        self.status = ResourceStatus::Failed;
506                        return vec![ResourceAction::Failed(ResourceError::Timeout)];
507                    }
508                    self.retries_left -= 1;
509                    self.last_part_sent = now;
510                    // In Python, this queries network cache. We just signal retry.
511                    return vec![];
512                }
513            }
514            _ => {}
515        }
516
517        vec![]
518    }
519
520    /// Get the total number of parts.
521    pub fn total_parts(&self) -> usize {
522        self.total_parts
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529    use crate::buffer::types::NoopCompressor;
530
531    fn identity_encrypt(data: &[u8]) -> Vec<u8> {
532        data.to_vec()
533    }
534
535    fn make_sender(data: &[u8]) -> ResourceSender {
536        let mut rng = rns_crypto::FixedRng::new(&[0x42; 64]);
537        ResourceSender::new(
538            data,
539            None,
540            RESOURCE_SDU,
541            &identity_encrypt,
542            &NoopCompressor,
543            &mut rng,
544            1000.0,
545            false,
546            false,
547            None,
548            1,
549            1,
550            None,
551            0.5,
552            6.0,
553        )
554        .unwrap()
555    }
556
557    fn varying_data(len: usize) -> Vec<u8> {
558        let mut state = 0x1234_5678u32;
559        (0..len)
560            .map(|_| {
561                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
562                (state >> 16) as u8
563            })
564            .collect()
565    }
566
567    #[test]
568    fn test_new_sender_status() {
569        let sender = make_sender(b"test data");
570        assert_eq!(sender.status, ResourceStatus::Queued);
571    }
572
573    #[test]
574    fn test_new_sender_parts() {
575        let data = vec![0xAA; 1000];
576        let sender = make_sender(&data);
577        // 4 (random) + 1000 data = 1004 encrypted bytes
578        // 1004 / 464 = 3 parts (464, 464, 76)
579        assert_eq!(sender.total_parts(), 3);
580    }
581
582    #[test]
583    fn test_advertise() {
584        let mut sender = make_sender(b"test data");
585        let actions = sender.advertise(1000.0);
586        assert_eq!(sender.status, ResourceStatus::Advertised);
587        assert_eq!(actions.len(), 1);
588        match &actions[0] {
589            ResourceAction::SendAdvertisement(data) => {
590                assert!(!data.is_empty());
591            }
592            _ => panic!("Expected SendAdvertisement"),
593        }
594    }
595
596    #[test]
597    fn test_handle_request_basic() {
598        let mut sender = make_sender(b"short");
599        sender.advertise(1000.0);
600
601        // Build a request: [not_exhausted][resource_hash][first part hash]
602        let mut request = Vec::new();
603        request.push(RESOURCE_HASHMAP_IS_NOT_EXHAUSTED);
604        request.extend_from_slice(&sender.resource_hash);
605        request.extend_from_slice(&sender.part_hashes[0]);
606
607        let actions = sender.handle_request(&request, 1001.0);
608        assert!(!actions.is_empty());
609        // Should have sent a part
610        let has_part = actions
611            .iter()
612            .any(|a| matches!(a, ResourceAction::SendPart(_)));
613        assert!(has_part);
614    }
615
616    #[test]
617    fn test_request_for_another_resource_is_ignored_without_mutating_sender() {
618        let mut sender = make_sender(b"short");
619        sender.advertise(1000.0);
620        let status = sender.status;
621        let last_activity = sender.last_activity;
622        let retries_left = sender.retries_left;
623        let mut request = Vec::new();
624        request.push(RESOURCE_HASHMAP_IS_NOT_EXHAUSTED);
625        request.extend_from_slice(&[0xEE; 32]);
626        request.extend_from_slice(&sender.part_hashes[0]);
627
628        let actions = sender.handle_request(&request, 2000.0);
629
630        assert!(actions.is_empty());
631        assert_eq!(sender.status, status);
632        assert_eq!(sender.last_activity, last_activity);
633        assert_eq!(sender.retries_left, retries_left);
634        assert_eq!(sender.sent_parts, 0);
635    }
636
637    fn exhausted_request(
638        sender: &ResourceSender,
639        last_map_hash: [u8; RESOURCE_MAPHASH_LEN],
640    ) -> Vec<u8> {
641        let mut request = Vec::new();
642        request.push(RESOURCE_HASHMAP_IS_EXHAUSTED);
643        request.extend_from_slice(&last_map_hash);
644        request.extend_from_slice(&sender.resource_hash);
645        request
646    }
647
648    #[test]
649    fn test_empty_hmu_generation_cancels_sender() {
650        let data =
651            varying_data(RESOURCE_HASHMAP_MAX_LEN * RESOURCE_SDU - RESOURCE_RANDOM_HASH_SIZE);
652        let mut sender = make_sender(&data);
653        assert_eq!(sender.total_parts(), RESOURCE_HASHMAP_MAX_LEN);
654        sender.advertise(1000.0);
655        let request = exhausted_request(&sender, sender.part_hashes[RESOURCE_HASHMAP_MAX_LEN - 1]);
656
657        let actions = sender.handle_request(&request, 1001.0);
658
659        assert_eq!(sender.status, ResourceStatus::Failed);
660        assert!(actions
661            .iter()
662            .any(|action| matches!(action, ResourceAction::SendCancelInitiator(_))));
663        assert!(!actions
664            .iter()
665            .any(|action| matches!(action, ResourceAction::SendHmu(_))));
666    }
667
668    #[test]
669    fn test_misaligned_hmu_request_cancels_sender() {
670        let data = varying_data((RESOURCE_HASHMAP_MAX_LEN + 2) * RESOURCE_SDU);
671        let mut sender = make_sender(&data);
672        sender.advertise(1000.0);
673        let request = exhausted_request(&sender, sender.part_hashes[5]);
674
675        let actions = sender.handle_request(&request, 1001.0);
676
677        assert_eq!(sender.status, ResourceStatus::Failed);
678        assert!(actions
679            .iter()
680            .any(|action| matches!(action, ResourceAction::SendCancelInitiator(_))));
681        assert!(!actions
682            .iter()
683            .any(|action| matches!(action, ResourceAction::SendHmu(_))));
684    }
685
686    #[test]
687    fn test_valid_nonempty_hmu_is_emitted() {
688        let data = varying_data((RESOURCE_HASHMAP_MAX_LEN + 2) * RESOURCE_SDU);
689        let mut sender = make_sender(&data);
690        sender.advertise(1000.0);
691        let request = exhausted_request(&sender, sender.part_hashes[RESOURCE_HASHMAP_MAX_LEN - 1]);
692
693        let actions = sender.handle_request(&request, 1001.0);
694
695        assert_ne!(sender.status, ResourceStatus::Failed);
696        let hmu = actions
697            .iter()
698            .find_map(|action| match action {
699                ResourceAction::SendHmu(data) => Some(data),
700                _ => None,
701            })
702            .expect("valid exhausted request should produce an HMU");
703        let (value, consumed) = crate::msgpack::unpack(&hmu[32..]).unwrap();
704        assert_eq!(consumed, hmu.len() - 32);
705        let array = value.as_array().unwrap();
706        assert_eq!(array[0].as_uint(), Some(1));
707        assert_eq!(
708            array[1].as_bin().map(<[u8]>::len),
709            Some((sender.total_parts() - RESOURCE_HASHMAP_MAX_LEN) * RESOURCE_MAPHASH_LEN)
710        );
711    }
712
713    #[test]
714    fn test_all_parts_sent_awaiting_proof() {
715        let mut sender = make_sender(b"hi");
716        sender.advertise(1000.0);
717
718        // Request all parts
719        let mut request = Vec::new();
720        request.push(RESOURCE_HASHMAP_IS_NOT_EXHAUSTED);
721        request.extend_from_slice(&sender.resource_hash);
722        for h in &sender.part_hashes.clone() {
723            request.extend_from_slice(h);
724        }
725
726        let _actions = sender.handle_request(&request, 1001.0);
727        assert_eq!(sender.status, ResourceStatus::AwaitingProof);
728        assert_eq!(sender.retries_left, 3);
729    }
730
731    #[test]
732    fn test_valid_proof() {
733        let mut sender = make_sender(b"data");
734        sender.advertise(1000.0);
735
736        let proof_data =
737            super::super::proof::build_proof_data(&sender.resource_hash, &sender.expected_proof);
738        let actions = sender.handle_proof(&proof_data, 1002.0);
739        assert_eq!(sender.status, ResourceStatus::Complete);
740        assert!(actions
741            .iter()
742            .any(|a| matches!(a, ResourceAction::Completed)));
743    }
744
745    #[test]
746    fn test_invalid_proof() {
747        let mut sender = make_sender(b"data");
748        sender.advertise(1000.0);
749
750        let wrong_proof = [0xFF; 32];
751        let proof_data = super::super::proof::build_proof_data(&sender.resource_hash, &wrong_proof);
752        let _actions = sender.handle_proof(&proof_data, 1002.0);
753        assert_eq!(sender.status, ResourceStatus::Failed);
754    }
755
756    #[test]
757    fn test_handle_reject() {
758        let mut sender = make_sender(b"data");
759        sender.advertise(1000.0);
760        let _actions = sender.handle_reject();
761        assert_eq!(sender.status, ResourceStatus::Rejected);
762    }
763
764    #[test]
765    fn test_cancel() {
766        let mut sender = make_sender(b"data");
767        sender.advertise(1000.0);
768        let actions = sender.cancel();
769        assert_eq!(sender.status, ResourceStatus::Failed);
770        assert!(actions
771            .iter()
772            .any(|a| matches!(a, ResourceAction::SendCancelInitiator(_))));
773    }
774
775    #[test]
776    fn test_cancel_already_complete() {
777        let mut sender = make_sender(b"data");
778        sender.status = ResourceStatus::Complete;
779        let actions = sender.cancel();
780        assert!(actions.is_empty());
781    }
782
783    #[test]
784    fn test_tick_advertised_timeout() {
785        let mut sender = make_sender(b"data");
786        sender.advertise(1000.0);
787        sender.retries_left = 0;
788
789        // Way past timeout
790        let _actions = sender.tick(2000.0);
791        assert_eq!(sender.status, ResourceStatus::Failed);
792    }
793
794    #[test]
795    fn test_tick_advertised_retry() {
796        let mut sender = make_sender(b"data");
797        sender.advertise(1000.0);
798        assert!(sender.retries_left > 0);
799
800        let actions = sender.tick(2000.0);
801        // Should retry advertisement
802        assert!(actions
803            .iter()
804            .any(|a| matches!(a, ResourceAction::SendAdvertisement(_))));
805    }
806
807    #[test]
808    fn test_resource_hash_is_32_bytes() {
809        let sender = make_sender(b"data");
810        assert_eq!(sender.resource_hash.len(), 32);
811        assert_eq!(sender.expected_proof.len(), 32);
812    }
813
814    #[test]
815    fn test_sender_with_metadata() {
816        let mut rng = rns_crypto::FixedRng::new(&[0x55; 64]);
817        let sender = ResourceSender::new(
818            b"data",
819            Some(b"metadata"),
820            RESOURCE_SDU,
821            &identity_encrypt,
822            &NoopCompressor,
823            &mut rng,
824            1000.0,
825            false,
826            false,
827            None,
828            1,
829            1,
830            None,
831            0.5,
832            6.0,
833        )
834        .unwrap();
835        assert!(sender.flags.has_metadata);
836    }
837
838    #[test]
839    fn test_multi_segment_sender() {
840        let orig_hash = [0xBB; 32];
841        let mut rng = rns_crypto::FixedRng::new(&[0x66; 64]);
842        let sender = ResourceSender::new(
843            b"segment 2 data",
844            None,
845            RESOURCE_SDU,
846            &identity_encrypt,
847            &NoopCompressor,
848            &mut rng,
849            1000.0,
850            false,
851            false,
852            None,
853            2,
854            5,
855            Some(orig_hash),
856            0.5,
857            6.0,
858        )
859        .unwrap();
860        assert_eq!(sender.segment_index, 2);
861        assert_eq!(sender.total_segments, 5);
862        assert_eq!(sender.original_hash, orig_hash);
863        assert!(sender.flags.split);
864    }
865}