1use serde::{Deserialize, Serialize};
24use std::collections::{BTreeMap, HashMap};
25
26use crate::hash::sha256;
27use crate::types::{Hash, Link, LinkType, TreeNode};
28
29#[derive(Debug, thiserror::Error)]
31pub enum CodecError {
32 #[error("Invalid node type: {0}")]
33 InvalidNodeType(u8),
34 #[error("Invalid link type: {0}")]
35 InvalidLinkType(u8),
36 #[error("Missing required field: {0}")]
37 MissingField(&'static str),
38 #[error("Invalid field type for {0}")]
39 InvalidFieldType(&'static str),
40 #[error("MessagePack encoding error: {0}")]
41 MsgpackEncode(String),
42 #[error("MessagePack decoding error: {0}")]
43 MsgpackDecode(String),
44 #[error("Invalid hash length: expected 32, got {0}")]
45 InvalidHashLength(usize),
46}
47
48#[derive(Serialize, Deserialize)]
51struct WireLink {
52 #[serde(with = "serde_bytes")]
54 h: Vec<u8>,
55 #[serde(
57 default,
58 skip_serializing_if = "Option::is_none",
59 with = "option_bytes"
60 )]
61 k: Option<Vec<u8>>,
62 #[serde(
64 skip_serializing_if = "Option::is_none",
65 serialize_with = "crate::canonical::serialize_metadata"
66 )]
67 m: Option<BTreeMap<String, serde_json::Value>>,
68 #[serde(skip_serializing_if = "Option::is_none")]
70 n: Option<String>,
71 s: u64,
73 #[serde(default)]
75 t: u8,
76}
77
78mod option_bytes {
80 use serde::{Deserialize, Deserializer, Serializer};
81
82 pub fn serialize<S>(data: &Option<Vec<u8>>, serializer: S) -> Result<S::Ok, S::Error>
83 where
84 S: Serializer,
85 {
86 match data {
87 Some(bytes) => serde_bytes::serialize(bytes, serializer),
88 None => serializer.serialize_none(),
89 }
90 }
91
92 pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error>
93 where
94 D: Deserializer<'de>,
95 {
96 Option::<serde_bytes::ByteBuf>::deserialize(deserializer)
97 .map(|opt| opt.map(|bb| bb.into_vec()))
98 }
99}
100
101#[derive(Serialize, Deserialize)]
104struct WireTreeNode {
105 l: Vec<WireLink>,
107 t: u8,
109}
110
111#[derive(Deserialize)]
113#[allow(dead_code)]
114struct WireLinkShape {
115 #[serde(with = "serde_bytes")]
116 h: Vec<u8>,
117 #[serde(default)]
118 t: u8,
119}
120
121#[derive(Deserialize)]
122#[allow(dead_code)]
123struct WireTreeShape {
124 l: Vec<WireLinkShape>,
125 t: u8,
126}
127
128pub fn encode_tree_node(node: &TreeNode) -> Result<Vec<u8>, CodecError> {
130 let links = links_for_encoding(node);
131 let wire = WireTreeNode {
132 t: node.node_type as u8,
133 l: links
134 .iter()
135 .map(|link| {
136 let sorted_meta = link
138 .meta
139 .as_ref()
140 .map(|m| m.iter().collect::<BTreeMap<_, _>>());
141 WireLink {
142 h: link.hash.to_vec(),
143 t: link.link_type as u8,
144 n: link.name.clone(),
145 s: link.size,
146 k: link.key.map(|k| k.to_vec()),
147 m: sorted_meta
148 .map(|m| m.into_iter().map(|(k, v)| (k.clone(), v.clone())).collect()),
149 }
150 })
151 .collect(),
152 };
153
154 rmp_serde::to_vec_named(&wire).map_err(|e| CodecError::MsgpackEncode(e.to_string()))
155}
156
157fn links_for_encoding(node: &TreeNode) -> Vec<&Link> {
158 let mut links = node.links.iter().collect::<Vec<_>>();
159 if node.node_type == LinkType::Dir {
160 links.sort_by(|left, right| dir_link_sort_name(left).cmp(dir_link_sort_name(right)));
161 }
162 links
163}
164
165fn dir_link_sort_name(link: &Link) -> &str {
166 link.name.as_deref().unwrap_or("")
167}
168
169pub fn decode_tree_node(data: &[u8]) -> Result<TreeNode, CodecError> {
171 let wire: WireTreeNode =
172 rmp_serde::from_slice(data).map_err(|e| CodecError::MsgpackDecode(e.to_string()))?;
173
174 let node_type = LinkType::from_u8(wire.t)
176 .filter(|t| t.is_tree())
177 .ok_or(CodecError::InvalidNodeType(wire.t))?;
178
179 let mut links = Vec::with_capacity(wire.l.len());
180 for wl in wire.l {
181 if wl.h.len() != 32 {
182 return Err(CodecError::InvalidHashLength(wl.h.len()));
183 }
184 let mut hash = [0u8; 32];
185 hash.copy_from_slice(&wl.h);
186
187 let key = match wl.k {
188 Some(k) if k.len() == 32 => {
189 let mut key = [0u8; 32];
190 key.copy_from_slice(&k);
191 Some(key)
192 }
193 _ => None,
194 };
195
196 let link_type = LinkType::from_u8(wl.t).ok_or(CodecError::InvalidLinkType(wl.t))?;
197
198 let meta = wl.m.map(|m| m.into_iter().collect::<HashMap<_, _>>());
200
201 links.push(Link {
202 hash,
203 name: wl.n,
204 size: wl.s,
205 key,
206 link_type,
207 meta,
208 });
209 }
210
211 Ok(TreeNode { node_type, links })
212}
213
214pub fn encode_and_hash(node: &TreeNode) -> Result<(Vec<u8>, Hash), CodecError> {
216 let data = encode_tree_node(node)?;
217 let hash = sha256(&data);
218 Ok((data, hash))
219}
220
221pub fn try_decode_tree_node(data: &[u8]) -> Option<TreeNode> {
225 decode_tree_node(data).ok()
226}
227
228pub fn get_node_type(data: &[u8]) -> LinkType {
231 try_decode_tree_node(data)
232 .map(|n| n.node_type)
233 .unwrap_or(LinkType::Blob)
234}
235
236pub fn is_tree_node(data: &[u8]) -> bool {
242 if try_decode_tree_node(data).is_some() {
243 return true;
244 }
245
246 rmp_serde::from_slice::<WireTreeShape>(data).is_ok()
247}
248
249pub fn is_directory_node(data: &[u8]) -> bool {
252 try_decode_tree_node(data)
253 .map(|n| n.node_type.is_directory_like())
254 .unwrap_or(false)
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260 use crate::types::to_hex;
261
262 #[test]
263 fn test_encode_decode_empty_tree() {
264 let node = TreeNode::dir(vec![]);
265
266 let encoded = encode_tree_node(&node).unwrap();
267 let decoded = decode_tree_node(&encoded).unwrap();
268
269 assert_eq!(decoded.links.len(), 0);
270 assert_eq!(decoded.node_type, LinkType::Dir);
271 }
272
273 #[test]
274 fn test_encode_decode_tree_with_links() {
275 let hash1 = [1u8; 32];
276 let hash2 = [2u8; 32];
277
278 let node = TreeNode::dir(vec![
279 Link {
280 hash: hash1,
281 name: Some("file1.txt".to_string()),
282 size: 100,
283 key: None,
284 link_type: LinkType::Blob,
285 meta: None,
286 },
287 Link {
288 hash: hash2,
289 name: Some("dir".to_string()),
290 size: 0,
291 key: None,
292 link_type: LinkType::Dir,
293 meta: None,
294 },
295 ]);
296
297 let encoded = encode_tree_node(&node).unwrap();
298 let decoded = decode_tree_node(&encoded).unwrap();
299
300 assert_eq!(decoded.links.len(), 2);
301 assert_eq!(decoded.links[0].name, Some("dir".to_string()));
302 assert_eq!(decoded.links[0].size, 0);
303 assert_eq!(decoded.links[0].link_type, LinkType::Dir);
304 assert_eq!(to_hex(&decoded.links[0].hash), to_hex(&hash2));
305 assert_eq!(decoded.links[1].name, Some("file1.txt".to_string()));
306 assert_eq!(decoded.links[1].size, 100);
307 assert_eq!(decoded.links[1].link_type, LinkType::Blob);
308 assert_eq!(to_hex(&decoded.links[1].hash), to_hex(&hash1));
309 }
310
311 #[test]
312 fn test_preserve_link_meta() {
313 let mut meta = HashMap::new();
314 meta.insert("createdAt".to_string(), serde_json::json!(1234567890));
315 meta.insert("mimeType".to_string(), serde_json::json!("image/png"));
316
317 let node = TreeNode::dir(vec![Link::new([1u8; 32])
318 .with_name("file.png")
319 .with_size(1024)
320 .with_meta(meta.clone())]);
321
322 let encoded = encode_tree_node(&node).unwrap();
323 let decoded = decode_tree_node(&encoded).unwrap();
324
325 assert!(decoded.links[0].meta.is_some());
326 let m = decoded.links[0].meta.as_ref().unwrap();
327 assert_eq!(m.get("createdAt"), Some(&serde_json::json!(1234567890)));
328 assert_eq!(m.get("mimeType"), Some(&serde_json::json!("image/png")));
329 }
330
331 #[test]
332 fn test_links_without_optional_fields() {
333 let hash = [42u8; 32];
334
335 let node = TreeNode::file(vec![Link::new(hash)]);
336
337 let encoded = encode_tree_node(&node).unwrap();
338 let decoded = decode_tree_node(&encoded).unwrap();
339
340 assert_eq!(decoded.links[0].name, None);
341 assert_eq!(decoded.links[0].size, 0);
342 assert_eq!(decoded.links[0].link_type, LinkType::Blob);
343 assert_eq!(decoded.links[0].meta, None);
344 assert_eq!(to_hex(&decoded.links[0].hash), to_hex(&hash));
345 }
346
347 #[test]
348 fn test_encode_and_hash() {
349 let node = TreeNode::dir(vec![]);
350
351 let (data, hash) = encode_and_hash(&node).unwrap();
352 let expected_hash = sha256(&data);
353
354 assert_eq!(to_hex(&hash), to_hex(&expected_hash));
355 }
356
357 #[test]
358 fn test_encode_and_hash_consistent() {
359 let node = TreeNode::dir(vec![Link {
360 hash: [1u8; 32],
361 name: Some("test".to_string()),
362 size: 100,
363 key: None,
364 link_type: LinkType::Blob,
365 meta: None,
366 }]);
367
368 let (_, hash1) = encode_and_hash(&node).unwrap();
369 let (_, hash2) = encode_and_hash(&node).unwrap();
370
371 assert_eq!(to_hex(&hash1), to_hex(&hash2));
372 }
373
374 #[test]
375 fn test_is_tree_node() {
376 let node = TreeNode::dir(vec![]);
377 let encoded = encode_tree_node(&node).unwrap();
378
379 assert!(is_tree_node(&encoded));
380 }
381
382 #[test]
383 fn test_is_tree_node_raw_blob() {
384 let blob = vec![1u8, 2, 3, 4, 5];
385 assert!(!is_tree_node(&blob));
386 }
387
388 #[test]
389 fn test_is_tree_node_invalid_msgpack() {
390 let invalid = vec![255u8, 255, 255];
391 assert!(!is_tree_node(&invalid));
392 }
393
394 #[test]
395 fn test_is_tree_node_unknown_node_type() {
396 let bytes = rmp_serde::to_vec_named(&WireTreeNode { l: vec![], t: 99 }).unwrap();
397
398 assert!(is_tree_node(&bytes));
399 assert!(matches!(
400 decode_tree_node(&bytes),
401 Err(CodecError::InvalidNodeType(99))
402 ));
403 }
404
405 #[test]
406 fn test_is_tree_node_unknown_link_type() {
407 let bytes = rmp_serde::to_vec_named(&WireTreeNode {
408 l: vec![WireLink {
409 h: vec![1u8; 32],
410 k: None,
411 m: None,
412 n: None,
413 s: 1,
414 t: 99,
415 }],
416 t: LinkType::Dir as u8,
417 })
418 .unwrap();
419
420 assert!(is_tree_node(&bytes));
421 assert!(matches!(
422 decode_tree_node(&bytes),
423 Err(CodecError::InvalidLinkType(99))
424 ));
425 }
426
427 #[test]
428 fn test_is_directory_node() {
429 let node = TreeNode::dir(vec![Link {
430 hash: [1u8; 32],
431 name: Some("file.txt".to_string()),
432 size: 100,
433 key: None,
434 link_type: LinkType::Blob,
435 meta: None,
436 }]);
437 let encoded = encode_tree_node(&node).unwrap();
438
439 assert!(is_directory_node(&encoded));
440 }
441
442 #[test]
443 fn test_is_directory_node_empty() {
444 let node = TreeNode::dir(vec![]);
445 let encoded = encode_tree_node(&node).unwrap();
446
447 assert!(is_directory_node(&encoded));
448 }
449
450 #[test]
451 fn test_is_not_directory_node() {
452 let node = TreeNode::file(vec![Link::new([1u8; 32])]);
454 let encoded = encode_tree_node(&node).unwrap();
455
456 assert!(!is_directory_node(&encoded));
457 }
458
459 #[test]
460 fn test_encrypted_link_roundtrip() {
461 let hash = [1u8; 32];
462 let key = [2u8; 32];
463
464 let node = TreeNode::dir(vec![Link {
465 hash,
466 name: Some("encrypted.dat".to_string()),
467 size: 1024,
468 key: Some(key),
469 link_type: LinkType::Blob,
470 meta: None,
471 }]);
472
473 let encoded = encode_tree_node(&node).unwrap();
474 let decoded = decode_tree_node(&encoded).unwrap();
475
476 assert_eq!(decoded.links[0].key, Some(key));
477 }
478
479 #[test]
480 fn test_encoding_determinism() {
481 let hash = [42u8; 32];
484
485 let node = TreeNode::dir(vec![Link {
486 hash,
487 name: Some("file.txt".to_string()),
488 size: 100,
489 key: None,
490 link_type: LinkType::Blob,
491 meta: None,
492 }]);
493
494 let encoded1 = encode_tree_node(&node).unwrap();
496 let encoded2 = encode_tree_node(&node).unwrap();
497 let encoded3 = encode_tree_node(&node).unwrap();
498
499 assert_eq!(encoded1, encoded2, "Encoding should be deterministic");
500 assert_eq!(encoded2, encoded3, "Encoding should be deterministic");
501 }
502
503 #[test]
504 fn test_directory_link_order_is_canonicalized() {
505 let apple = Link::new([1u8; 32]).with_name("apple").with_size(1);
506 let zebra = Link::new([2u8; 32]).with_name("zebra").with_size(2);
507
508 let sorted = TreeNode::dir(vec![apple.clone(), zebra.clone()]);
509 let reversed = TreeNode::dir(vec![zebra, apple]);
510
511 let encoded_sorted = encode_tree_node(&sorted).unwrap();
512 let encoded_reversed = encode_tree_node(&reversed).unwrap();
513 assert_eq!(encoded_sorted, encoded_reversed);
514
515 let decoded = decode_tree_node(&encoded_reversed).unwrap();
516 let names = decoded
517 .links
518 .iter()
519 .map(|link| link.name.as_deref())
520 .collect::<Vec<_>>();
521 assert_eq!(names, vec![Some("apple"), Some("zebra")]);
522 }
523
524 #[test]
525 fn test_directory_link_order_uses_utf8_bytes() {
526 let bmp_private_use = Link::new([1u8; 32]).with_name("\u{E000}").with_size(1);
527 let supplementary = Link::new([2u8; 32]).with_name("\u{10000}").with_size(2);
528
529 let node = TreeNode::dir(vec![supplementary, bmp_private_use]);
530 let decoded = decode_tree_node(&encode_tree_node(&node).unwrap()).unwrap();
531 let names = decoded
532 .links
533 .iter()
534 .map(|link| link.name.as_deref())
535 .collect::<Vec<_>>();
536
537 assert_eq!(names, vec![Some("\u{E000}"), Some("\u{10000}")]);
538 }
539
540 #[test]
541 fn test_file_link_order_is_preserved() {
542 let first = Link::new([1u8; 32]).with_size(1);
543 let second = Link::new([2u8; 32]).with_size(1);
544
545 let file = TreeNode::file(vec![first.clone(), second.clone()]);
546 let reversed_file = TreeNode::file(vec![second, first]);
547
548 assert_ne!(
549 encode_tree_node(&file).unwrap(),
550 encode_tree_node(&reversed_file).unwrap()
551 );
552
553 let decoded = decode_tree_node(&encode_tree_node(&reversed_file).unwrap()).unwrap();
554 assert_eq!(decoded.links[0].hash, [2u8; 32]);
555 assert_eq!(decoded.links[1].hash, [1u8; 32]);
556 }
557
558 #[test]
559 fn test_bud17_directory_fanout_vector() {
560 let mut first_meta = HashMap::new();
561 first_meta.insert("count".to_string(), serde_json::json!(2));
562 first_meta.insert("first".to_string(), serde_json::json!("a.txt"));
563 first_meta.insert("last".to_string(), serde_json::json!("b.txt"));
564
565 let mut second_meta = HashMap::new();
566 second_meta.insert("count".to_string(), serde_json::json!(1));
567 second_meta.insert("first".to_string(), serde_json::json!("c.txt"));
568 second_meta.insert("last".to_string(), serde_json::json!("c.txt"));
569
570 let node = TreeNode::fanout(vec![
571 Link::new([0x11u8; 32])
572 .with_size(30)
573 .with_link_type(LinkType::Dir)
574 .with_meta(first_meta),
575 Link::new([0x22u8; 32])
576 .with_size(40)
577 .with_link_type(LinkType::Dir)
578 .with_meta(second_meta),
579 ]);
580
581 let encoded = encode_tree_node(&node).unwrap();
582 assert_eq!(
583 hex::encode(&encoded),
584 "82a16c9284a168c4201111111111111111111111111111111111111111111111111111111111111111a16d83a5636f756e7402a56669727374a5612e747874a46c617374a5622e747874a1731ea1740284a168c4202222222222222222222222222222222222222222222222222222222222222222a16d83a5636f756e7401a56669727374a5632e747874a46c617374a5632e747874a17328a17402a17403"
585 );
586 assert_eq!(
587 to_hex(&sha256(&encoded)),
588 "6626ab03b5468f417d888fa25fa22b48f5bcb7dfafb88eef34c638d167afc0a3"
589 );
590
591 let decoded = decode_tree_node(&encoded).unwrap();
592 assert_eq!(decoded.node_type, LinkType::Fanout);
593 assert!(decoded.links.iter().all(|link| link.name.is_none()));
594 }
595
596 #[test]
597 fn test_link_meta_determinism() {
598 let hash = [1u8; 32];
601
602 let mut meta1 = HashMap::new();
604 meta1.insert("zebra".to_string(), serde_json::json!("last"));
605 meta1.insert("alpha".to_string(), serde_json::json!("first"));
606 meta1.insert("middle".to_string(), serde_json::json!("mid"));
607
608 let mut meta2 = HashMap::new();
609 meta2.insert("alpha".to_string(), serde_json::json!("first"));
610 meta2.insert("middle".to_string(), serde_json::json!("mid"));
611 meta2.insert("zebra".to_string(), serde_json::json!("last"));
612
613 let node1 = TreeNode::dir(vec![Link::new(hash)
614 .with_name("file")
615 .with_size(100)
616 .with_meta(meta1)]);
617 let node2 = TreeNode::dir(vec![Link::new(hash)
618 .with_name("file")
619 .with_size(100)
620 .with_meta(meta2)]);
621
622 let encoded1 = encode_tree_node(&node1).unwrap();
623 let encoded2 = encode_tree_node(&node2).unwrap();
624
625 assert_eq!(
627 encoded1, encoded2,
628 "Link meta encoding should be deterministic regardless of insertion order"
629 );
630
631 let hash1 = crate::hash::sha256(&encoded1);
633 let hash2 = crate::hash::sha256(&encoded2);
634 assert_eq!(hash1, hash2, "Hashes should match for identical content");
635 }
636
637 #[test]
638 fn test_get_node_type() {
639 let dir_node = TreeNode::dir(vec![]);
640 let dir_encoded = encode_tree_node(&dir_node).unwrap();
641 assert_eq!(get_node_type(&dir_encoded), LinkType::Dir);
642
643 let file_node = TreeNode::file(vec![]);
644 let file_encoded = encode_tree_node(&file_node).unwrap();
645 assert_eq!(get_node_type(&file_encoded), LinkType::File);
646
647 let fanout_node = TreeNode::fanout(vec![]);
648 let fanout_encoded = encode_tree_node(&fanout_node).unwrap();
649 assert_eq!(get_node_type(&fanout_encoded), LinkType::Fanout);
650
651 let blob = vec![1u8, 2, 3, 4, 5];
653 assert_eq!(get_node_type(&blob), LinkType::Blob);
654 }
655
656 #[test]
657 fn test_link_type_roundtrip() {
658 let node = TreeNode::dir(vec![
659 Link::new([1u8; 32]).with_link_type(LinkType::Blob),
660 Link::new([2u8; 32]).with_link_type(LinkType::File),
661 Link::new([3u8; 32]).with_link_type(LinkType::Dir),
662 ]);
663
664 let encoded = encode_tree_node(&node).unwrap();
665 let decoded = decode_tree_node(&encoded).unwrap();
666
667 assert_eq!(decoded.links[0].link_type, LinkType::Blob);
668 assert_eq!(decoded.links[1].link_type, LinkType::File);
669 assert_eq!(decoded.links[2].link_type, LinkType::Dir);
670 }
671}