Skip to main content

rns_core/resource/
sender.rs

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