rns_core/resource/
advertisement.rs

1use alloc::vec::Vec;
2
3use crate::msgpack::{self, Value};
4use crate::constants::{RESOURCE_HASHMAP_MAX_LEN, RESOURCE_MAPHASH_LEN};
5use super::types::{AdvFlags, ResourceError};
6
7/// Resource advertisement data, corresponding to Python's ResourceAdvertisement.
8#[derive(Debug, Clone)]
9pub struct ResourceAdvertisement {
10    /// Transfer size (encrypted data size)
11    pub transfer_size: u64,
12    /// Total uncompressed data size (including metadata overhead)
13    pub data_size: u64,
14    /// Number of parts
15    pub num_parts: u64,
16    /// Resource hash (full 32 bytes)
17    pub resource_hash: Vec<u8>,
18    /// Random hash (4 bytes)
19    pub random_hash: Vec<u8>,
20    /// Original hash (first segment, 32 bytes)
21    pub original_hash: Vec<u8>,
22    /// Hashmap segment (concatenated 4-byte part hashes)
23    pub hashmap: Vec<u8>,
24    /// Flags byte
25    pub flags: AdvFlags,
26    /// Segment index (1-based)
27    pub segment_index: u64,
28    /// Total segments
29    pub total_segments: u64,
30    /// Request ID (optional)
31    pub request_id: Option<Vec<u8>>,
32}
33
34impl ResourceAdvertisement {
35    /// Pack the advertisement to msgpack bytes.
36    /// `segment` controls which hashmap segment to include (0-based).
37    pub fn pack(&self, segment: usize) -> Vec<u8> {
38        let hashmap_start = segment * RESOURCE_HASHMAP_MAX_LEN * RESOURCE_MAPHASH_LEN;
39        let max_end = (segment + 1) * RESOURCE_HASHMAP_MAX_LEN * RESOURCE_MAPHASH_LEN;
40        let hashmap_end = core::cmp::min(max_end, self.hashmap.len());
41        let hashmap_segment = if hashmap_start < self.hashmap.len() {
42            &self.hashmap[hashmap_start..hashmap_end]
43        } else {
44            &[]
45        };
46
47        let q_value = match &self.request_id {
48            Some(id) => Value::Bin(id.clone()),
49            None => Value::Nil,
50        };
51
52        // Match Python's key order: t, d, n, h, r, o, i, l, q, f, m
53        let entries: Vec<(&str, Value)> = vec![
54            ("t", Value::UInt(self.transfer_size)),
55            ("d", Value::UInt(self.data_size)),
56            ("n", Value::UInt(self.num_parts)),
57            ("h", Value::Bin(self.resource_hash.clone())),
58            ("r", Value::Bin(self.random_hash.clone())),
59            ("o", Value::Bin(self.original_hash.clone())),
60            ("i", Value::UInt(self.segment_index)),
61            ("l", Value::UInt(self.total_segments)),
62            ("q", q_value),
63            ("f", Value::UInt(self.flags.to_byte() as u64)),
64            ("m", Value::Bin(hashmap_segment.to_vec())),
65        ];
66
67        msgpack::pack_str_map(&entries)
68    }
69
70    /// Unpack an advertisement from msgpack bytes.
71    pub fn unpack(data: &[u8]) -> Result<Self, ResourceError> {
72        let value = msgpack::unpack_exact(data).map_err(|_| ResourceError::InvalidAdvertisement)?;
73
74        let t = value.map_get("t")
75            .and_then(|v| v.as_uint())
76            .ok_or(ResourceError::InvalidAdvertisement)?;
77        let d = value.map_get("d")
78            .and_then(|v| v.as_uint())
79            .ok_or(ResourceError::InvalidAdvertisement)?;
80        let n = value.map_get("n")
81            .and_then(|v| v.as_uint())
82            .ok_or(ResourceError::InvalidAdvertisement)?;
83        let h = value.map_get("h")
84            .and_then(|v| v.as_bin())
85            .ok_or(ResourceError::InvalidAdvertisement)?
86            .to_vec();
87        let r = value.map_get("r")
88            .and_then(|v| v.as_bin())
89            .ok_or(ResourceError::InvalidAdvertisement)?
90            .to_vec();
91        let o = value.map_get("o")
92            .and_then(|v| v.as_bin())
93            .ok_or(ResourceError::InvalidAdvertisement)?
94            .to_vec();
95        let m = value.map_get("m")
96            .and_then(|v| v.as_bin())
97            .ok_or(ResourceError::InvalidAdvertisement)?
98            .to_vec();
99        let f = value.map_get("f")
100            .and_then(|v| v.as_uint())
101            .ok_or(ResourceError::InvalidAdvertisement)? as u8;
102        let i = value.map_get("i")
103            .and_then(|v| v.as_uint())
104            .ok_or(ResourceError::InvalidAdvertisement)?;
105        let l = value.map_get("l")
106            .and_then(|v| v.as_uint())
107            .ok_or(ResourceError::InvalidAdvertisement)?;
108
109        let q_val = value.map_get("q").ok_or(ResourceError::InvalidAdvertisement)?;
110        let request_id = if q_val.is_nil() {
111            None
112        } else {
113            Some(q_val.as_bin().ok_or(ResourceError::InvalidAdvertisement)?.to_vec())
114        };
115
116        Ok(ResourceAdvertisement {
117            transfer_size: t,
118            data_size: d,
119            num_parts: n,
120            resource_hash: h,
121            random_hash: r,
122            original_hash: o,
123            hashmap: m,
124            flags: AdvFlags::from_byte(f),
125            segment_index: i,
126            total_segments: l,
127            request_id,
128        })
129    }
130
131    /// Check if this advertisement is a request.
132    pub fn is_request(&self) -> bool {
133        self.request_id.is_some() && self.flags.is_request
134    }
135
136    /// Check if this advertisement is a response.
137    pub fn is_response(&self) -> bool {
138        self.request_id.is_some() && self.flags.is_response
139    }
140
141    /// Get the number of hashmap segments needed.
142    pub fn hashmap_segments(&self) -> usize {
143        let total_hashes = self.num_parts as usize;
144        if total_hashes == 0 {
145            return 1;
146        }
147        (total_hashes + RESOURCE_HASHMAP_MAX_LEN - 1) / RESOURCE_HASHMAP_MAX_LEN
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    fn make_adv(flags: AdvFlags) -> ResourceAdvertisement {
156        ResourceAdvertisement {
157            transfer_size: 1000,
158            data_size: 950,
159            num_parts: 3,
160            resource_hash: vec![0x11; 32],
161            random_hash: vec![0xAA, 0xBB, 0xCC, 0xDD],
162            original_hash: vec![0x22; 32],
163            hashmap: vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C],
164            flags,
165            segment_index: 1,
166            total_segments: 1,
167            request_id: None,
168        }
169    }
170
171    #[test]
172    fn test_pack_unpack_roundtrip() {
173        let flags = AdvFlags {
174            encrypted: true,
175            compressed: false,
176            split: false,
177            is_request: false,
178            is_response: false,
179            has_metadata: false,
180        };
181        let adv = make_adv(flags);
182        let packed = adv.pack(0);
183        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
184
185        assert_eq!(unpacked.transfer_size, 1000);
186        assert_eq!(unpacked.data_size, 950);
187        assert_eq!(unpacked.num_parts, 3);
188        assert_eq!(unpacked.resource_hash, vec![0x11; 32]);
189        assert_eq!(unpacked.random_hash, vec![0xAA, 0xBB, 0xCC, 0xDD]);
190        assert_eq!(unpacked.original_hash, vec![0x22; 32]);
191        assert_eq!(unpacked.flags, flags);
192        assert_eq!(unpacked.segment_index, 1);
193        assert_eq!(unpacked.total_segments, 1);
194        assert!(unpacked.request_id.is_none());
195    }
196
197    #[test]
198    fn test_flags_encrypted_compressed() {
199        let flags = AdvFlags {
200            encrypted: true,
201            compressed: true,
202            split: false,
203            is_request: false,
204            is_response: false,
205            has_metadata: false,
206        };
207        let adv = make_adv(flags);
208        let packed = adv.pack(0);
209        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
210        assert!(unpacked.flags.encrypted);
211        assert!(unpacked.flags.compressed);
212        assert!(!unpacked.flags.split);
213    }
214
215    #[test]
216    fn test_flags_with_metadata() {
217        let flags = AdvFlags {
218            encrypted: true,
219            compressed: false,
220            split: false,
221            is_request: false,
222            is_response: false,
223            has_metadata: true,
224        };
225        let adv = make_adv(flags);
226        let packed = adv.pack(0);
227        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
228        assert!(unpacked.flags.has_metadata);
229    }
230
231    #[test]
232    fn test_multi_segment() {
233        let flags = AdvFlags {
234            encrypted: true,
235            compressed: false,
236            split: true,
237            is_request: false,
238            is_response: false,
239            has_metadata: false,
240        };
241        let mut adv = make_adv(flags);
242        adv.segment_index = 2;
243        adv.total_segments = 5;
244        let packed = adv.pack(0);
245        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
246        assert!(unpacked.flags.split);
247        assert_eq!(unpacked.segment_index, 2);
248        assert_eq!(unpacked.total_segments, 5);
249    }
250
251    #[test]
252    fn test_with_request_id() {
253        let flags = AdvFlags {
254            encrypted: true,
255            compressed: false,
256            split: false,
257            is_request: true,
258            is_response: false,
259            has_metadata: false,
260        };
261        let mut adv = make_adv(flags);
262        adv.request_id = Some(vec![0xDE, 0xAD, 0xBE, 0xEF]);
263        let packed = adv.pack(0);
264        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
265        assert!(unpacked.is_request());
266        assert!(!unpacked.is_response());
267        assert_eq!(unpacked.request_id, Some(vec![0xDE, 0xAD, 0xBE, 0xEF]));
268    }
269
270    #[test]
271    fn test_is_response() {
272        let flags = AdvFlags {
273            encrypted: true,
274            compressed: false,
275            split: false,
276            is_request: false,
277            is_response: true,
278            has_metadata: false,
279        };
280        let mut adv = make_adv(flags);
281        adv.request_id = Some(vec![0x42; 16]);
282        assert!(adv.is_response());
283        assert!(!adv.is_request());
284    }
285
286    #[test]
287    fn test_nil_request_id() {
288        let flags = AdvFlags {
289            encrypted: true,
290            compressed: false,
291            split: false,
292            is_request: false,
293            is_response: false,
294            has_metadata: false,
295        };
296        let adv = make_adv(flags);
297        let packed = adv.pack(0);
298        let unpacked = ResourceAdvertisement::unpack(&packed).unwrap();
299        assert!(unpacked.request_id.is_none());
300        assert!(!unpacked.is_request());
301        assert!(!unpacked.is_response());
302    }
303
304    #[test]
305    fn test_hashmap_segmentation() {
306        // Create a large hashmap with > HASHMAP_MAX_LEN(74) hashes
307        let num_hashes = 100;
308        let hashmap: Vec<u8> = (0..num_hashes).flat_map(|i| vec![i as u8; 4]).collect();
309
310        let flags = AdvFlags {
311            encrypted: true,
312            compressed: false,
313            split: false,
314            is_request: false,
315            is_response: false,
316            has_metadata: false,
317        };
318        let adv = ResourceAdvertisement {
319            transfer_size: 50000,
320            data_size: 48000,
321            num_parts: num_hashes,
322            resource_hash: vec![0x11; 32],
323            random_hash: vec![0xAA; 4],
324            original_hash: vec![0x22; 32],
325            hashmap: hashmap.clone(),
326            flags,
327            segment_index: 1,
328            total_segments: 1,
329            request_id: None,
330        };
331
332        // Segment 0: first 74 hashes = 296 bytes
333        let packed0 = adv.pack(0);
334        let unpacked0 = ResourceAdvertisement::unpack(&packed0).unwrap();
335        assert_eq!(unpacked0.hashmap.len(), 74 * 4);
336
337        // Segment 1: remaining 26 hashes = 104 bytes
338        let packed1 = adv.pack(1);
339        let unpacked1 = ResourceAdvertisement::unpack(&packed1).unwrap();
340        assert_eq!(unpacked1.hashmap.len(), 26 * 4);
341    }
342
343    #[test]
344    fn test_hashmap_segments_count() {
345        let flags = AdvFlags {
346            encrypted: true, compressed: false, split: false,
347            is_request: false, is_response: false, has_metadata: false,
348        };
349        let mut adv = make_adv(flags);
350
351        adv.num_parts = 74; // exactly HASHMAP_MAX_LEN
352        assert_eq!(adv.hashmap_segments(), 1);
353
354        adv.num_parts = 75;
355        assert_eq!(adv.hashmap_segments(), 2);
356
357        adv.num_parts = 148;
358        assert_eq!(adv.hashmap_segments(), 2);
359
360        adv.num_parts = 149;
361        assert_eq!(adv.hashmap_segments(), 3);
362    }
363
364    #[test]
365    fn test_unpack_invalid_data() {
366        assert!(ResourceAdvertisement::unpack(&[]).is_err());
367        assert!(ResourceAdvertisement::unpack(&[0xc0]).is_err()); // nil
368        assert!(ResourceAdvertisement::unpack(&[0x01, 0x02]).is_err()); // not a map
369    }
370}
rns_core/resource/advertisement.rs

rns_core/resource/
advertisement.rs
