1use objects::{
3 object::{AnnotatedTag, State},
4 store::ObjectStore,
5};
6
7use crate::{ObjectData, ObjectId, ObjectRequest, ObjectType, ProtocolError, Result};
8
9pub const MAX_RECEIVED_REDACTIONS_BLOB_SIZE: u64 = 64 * 1024 * 1024;
17
18pub const MAX_RECEIVED_STATE_VISIBILITY_BLOB_SIZE: u64 = 64 * 1024 * 1024;
25
26const PULL_DECODE_ENVELOPE_HEADROOM: u64 = 1024 * 1024;
37
38const fn max_u64(a: u64, b: u64) -> u64 {
39 if a > b { a } else { b }
40}
41
42pub const MAX_PULL_FRAME_MESSAGE_SIZE: usize = (max_u64(
58 MAX_RECEIVED_REDACTIONS_BLOB_SIZE,
59 MAX_RECEIVED_STATE_VISIBILITY_BLOB_SIZE,
60) + PULL_DECODE_ENVELOPE_HEADROOM) as usize;
61
62pub fn check_received_transfer_blob_size(
71 blob_len: usize,
72 max_bytes: u64,
73 kind: &str,
74) -> Result<()> {
75 let len = u64::try_from(blob_len).map_err(|_| {
76 ProtocolError::InvalidState(format!("{kind} blob length does not fit in u64"))
77 })?;
78 if len > max_bytes {
79 return Err(ProtocolError::InvalidState(format!(
80 "{kind} blob exceeds receive size limit: {len} bytes (max {max_bytes})"
81 )));
82 }
83 Ok(())
84}
85
86pub fn admit_declared_received_len(declared: u64, max_bytes: u64, kind: &str) -> Result<usize> {
92 if declared > max_bytes {
93 return Err(ProtocolError::InvalidState(format!(
94 "{kind} exceeds receive size limit: {declared} bytes (max {max_bytes})"
95 )));
96 }
97 usize::try_from(declared)
98 .map_err(|_| ProtocolError::InvalidState(format!("{kind} exceeds this platform")))
99}
100
101#[allow(dead_code)]
102pub fn chunk_count(object_size: usize, chunk_size: usize) -> usize {
103 if object_size == 0 || chunk_size == 0 {
104 return 0;
105 }
106 object_size.div_ceil(chunk_size)
107}
108
109#[allow(dead_code)]
110pub fn chunk_bounds(
111 object_size: usize,
112 chunk_size: usize,
113 chunk_index: usize,
114) -> Option<(usize, usize)> {
115 if chunk_size == 0 {
116 return None;
117 }
118
119 let start = chunk_index.checked_mul(chunk_size)?;
120 if start >= object_size {
121 return None;
122 }
123 let end = (start + chunk_size).min(object_size);
124 Some((start, end - start))
125}
126
127#[allow(dead_code)]
128pub fn chunk_offset(chunk_index: usize, chunk_size: usize) -> Option<usize> {
129 chunk_index.checked_mul(chunk_size)
130}
131
132pub fn load_requested_object(store: &impl ObjectStore, req: &ObjectRequest) -> Result<ObjectData> {
133 let (obj_type, data) = match &req.id {
140 ObjectId::Hash(hash) => {
141 if let Some(blob) = store.get_blob(hash)? {
142 (ObjectType::Blob, blob.content().to_vec())
143 } else if let Some(data) = store.get_tree_serialized(hash)? {
144 (ObjectType::Tree, data)
145 } else {
146 return Err(ProtocolError::ObjectNotFound(hash.to_hex()));
147 }
148 }
149 ObjectId::StateId(state_id) => {
150 let state = store
151 .get_state(state_id)?
152 .ok_or_else(|| ProtocolError::ObjectNotFound(state_id.to_string()))?;
153 (ObjectType::State, rmp_serde::to_vec_named(&state)?)
154 }
155 ObjectId::StateAttachment { state, id, kind: _ } => {
156 let attachment = store
157 .get_state_attachment(state, id)?
158 .ok_or_else(|| ProtocolError::ObjectNotFound(id.to_string()))?;
159 (
160 ObjectType::StateAttachment,
161 rmp_serde::to_vec_named(&attachment)?,
162 )
163 }
164 };
165
166 Ok(ObjectData {
167 id: req.id.clone(),
168 obj_type,
169 data,
170 is_delta: false,
171 })
172}
173
174pub fn load_object_data(
175 store: &impl ObjectStore,
176 id: &ObjectId,
177 obj_type: ObjectType,
178) -> Result<ObjectData> {
179 let data = match (id, obj_type) {
180 (ObjectId::Hash(hash), ObjectType::Blob) => store
181 .get_blob(hash)?
182 .ok_or_else(|| ProtocolError::ObjectNotFound(hash.to_hex()))?
183 .content()
184 .to_vec(),
185 (ObjectId::Hash(hash), ObjectType::Tree) => store
186 .get_tree_serialized(hash)?
187 .ok_or_else(|| ProtocolError::ObjectNotFound(hash.to_hex()))?,
188 (ObjectId::Hash(hash), ObjectType::AnnotatedTag) => store
189 .get_annotated_tag(hash)?
190 .ok_or_else(|| ProtocolError::ObjectNotFound(hash.to_hex()))?
191 .encode_current_msgpack(),
192 (ObjectId::StateId(state_id), ObjectType::State) => {
193 let state = store
194 .get_state(state_id)?
195 .ok_or_else(|| ProtocolError::ObjectNotFound(state_id.to_string()))?;
196 rmp_serde::to_vec_named(&state)?
197 }
198 (ObjectId::Hash(hash), ObjectType::Redaction) => store
199 .get_redactions_bytes_for_blob(hash)?
200 .ok_or_else(|| ProtocolError::ObjectNotFound(hash.to_hex()))?,
201 (ObjectId::Hash(hash), ObjectType::Purge) => store
202 .get_redactions_bytes_for_blob(hash)?
203 .ok_or_else(|| ProtocolError::ObjectNotFound(hash.to_hex()))?,
204 (ObjectId::StateId(state_id), ObjectType::StateVisibility) => store
205 .get_state_visibility_bytes_for_state(state_id)?
206 .ok_or_else(|| ProtocolError::ObjectNotFound(state_id.to_string_full()))?,
207 (ObjectId::StateAttachment { state, id, kind: _ }, ObjectType::StateAttachment) => {
208 let attachment = store
209 .get_state_attachment(state, id)?
210 .ok_or_else(|| ProtocolError::ObjectNotFound(id.to_string()))?;
211 rmp_serde::to_vec_named(&attachment)?
212 }
213 (ObjectId::Hash(_), ObjectType::KeyBinding) => {
214 return Err(ProtocolError::InvalidState(
215 "KeyBinding registry objects must be constructed with encode_key_binding_registry"
216 .to_string(),
217 ));
218 }
219 _ => {
220 return Err(ProtocolError::InvalidState(
221 "object id/type mismatch".to_string(),
222 ));
223 }
224 };
225
226 Ok(ObjectData {
227 id: id.clone(),
228 obj_type,
229 data,
230 is_delta: false,
231 })
232}
233
234pub fn store_received_object(store: &impl ObjectStore, data: &ObjectData) -> Result<()> {
235 match (&data.id, data.obj_type) {
236 (ObjectId::Hash(hash), ObjectType::Blob) => {
237 store.put_blob_bytes_with_hash(&data.data, *hash)?;
238 }
239 (ObjectId::Hash(hash), ObjectType::Tree) => {
240 store
241 .put_tree_serialized(&data.data, *hash)
242 .map_err(|error| match error {
243 objects::error::HeddleError::Corruption { .. }
244 | objects::error::HeddleError::TreeStream(
245 objects::object::TreeStreamError::HashMismatch { .. },
246 ) => ProtocolError::InvalidState("tree hash mismatch".to_string()),
247 error => ProtocolError::InvalidState(format!("invalid tree object: {error}")),
248 })?;
249 }
250 (ObjectId::Hash(hash), ObjectType::AnnotatedTag) => {
251 let tag = AnnotatedTag::decode_current_msgpack(&data.data)
252 .map_err(|error| ProtocolError::InvalidState(error.to_string()))?;
253 if tag.hash() != *hash {
254 return Err(ProtocolError::InvalidState(
255 "annotated tag hash mismatch".to_string(),
256 ));
257 }
258 store.put_annotated_tag(&tag)?;
259 }
260 (ObjectId::StateId(state_id), ObjectType::State) => {
261 let state: State = rmp_serde::from_slice(&data.data)?;
262 if state.id() != *state_id {
263 return Err(ProtocolError::InvalidState(format!(
264 "StateId mismatch: expected {state_id}, computed {}",
265 state.id()
266 )));
267 }
268 store.put_state_serialized(&data.data, *state_id)?;
269 }
270 (ObjectId::StateAttachment { state, id, kind }, ObjectType::StateAttachment) => {
271 let attachment: objects::object::StateAttachment = rmp_serde::from_slice(&data.data)?;
272 if attachment.state_id != *state || attachment.id() != *id {
273 return Err(ProtocolError::InvalidState(
274 "state attachment id mismatch".to_string(),
275 ));
276 }
277 let body_kind = attachment.body.kind();
282 if *kind != body_kind {
283 return Err(ProtocolError::InvalidState(format!(
284 "state attachment kind mismatch: descriptor {kind:?}, body {body_kind:?}"
285 )));
286 }
287 store.put_state_attachment(&attachment)?;
288 }
289 (_, ObjectType::Redaction) => {
290 return Err(ProtocolError::InvalidState(
295 "Redaction objects must be persisted via Repository::accept_wire_redactions, \
296 not store_received_object — signature verification is required"
297 .to_string(),
298 ));
299 }
300 (_, ObjectType::Purge) => {
301 return Err(ProtocolError::InvalidState(
302 "Purge objects must be persisted via Repository::accept_wire_purge, not store_received_object — owner authorization is required"
303 .to_string(),
304 ));
305 }
306 (_, ObjectType::StateVisibility) => {
307 return Err(ProtocolError::InvalidState(
311 "StateVisibility objects must be persisted via Repository::accept_wire_state_visibility, \
312 not store_received_object — sidecar validation is required"
313 .to_string(),
314 ));
315 }
316 (_, ObjectType::KeyBinding) => {
317 return Err(ProtocolError::InvalidState(
318 "KeyBinding registry objects must be decoded and verified with decode_key_binding_registry"
319 .to_string(),
320 ));
321 }
322 _ => {
323 return Err(ProtocolError::InvalidState(
324 "object id/type mismatch".to_string(),
325 ));
326 }
327 }
328
329 Ok(())
330}
331
332#[cfg(test)]
333mod tests {
334 use objects::{
335 object::{
336 Attribution, Blob, ContentHash, Principal, State, StateAttachment, StateAttachmentBody,
337 TREE_BLOCK_ENCODING_VERSION, TREE_HEADER_LEN, Tree, TreeEntry,
338 },
339 store::{FsStore, ObjectStore},
340 };
341 use tempfile::TempDir;
342
343 use super::*;
344
345 fn create_test_store() -> (TempDir, FsStore) {
346 let temp = TempDir::new().unwrap();
347 let store = FsStore::new(temp.path().join(".heddle"));
348 store.init().unwrap();
349 (temp, store)
350 }
351
352 fn test_attribution() -> Attribution {
353 Attribution::human(Principal::new("Wire Tester", "wire@example.com"))
354 }
355
356 #[test]
357 fn primary_objects_roundtrip_through_wire_data() {
358 let (_source_temp, source) = create_test_store();
359 let (_dest_temp, dest) = create_test_store();
360
361 let blob = Blob::from("wire transfer blob\n");
362 let blob_hash = source.put_blob(&blob).unwrap();
363 let tree = Tree::from_entries(vec![TreeEntry::file("lib.rs", blob_hash, false).unwrap()]);
364 let tree_hash = source.put_tree(&tree).unwrap();
365 let state = State::new(tree_hash, Vec::new(), test_attribution())
366 .with_intent("exercise wire transfer");
367 source.put_state(&state).unwrap();
368
369 let blob_data = load_requested_object(
370 &source,
371 &ObjectRequest {
372 id: ObjectId::Hash(blob_hash),
373 have_base: None,
374 },
375 )
376 .unwrap();
377 assert_eq!(blob_data.obj_type, ObjectType::Blob);
378 assert_eq!(blob_data.data, blob.content());
379 store_received_object(&dest, &blob_data).unwrap();
380 assert_eq!(
381 dest.get_blob(&blob_hash).unwrap().unwrap().content(),
382 blob.content()
383 );
384
385 let tree_data = load_requested_object(
386 &source,
387 &ObjectRequest {
388 id: ObjectId::Hash(tree_hash),
389 have_base: None,
390 },
391 )
392 .unwrap();
393 assert_eq!(tree_data.obj_type, ObjectType::Tree);
394 assert_eq!(
395 objects::store::codec::decode_tree_serialized_with_key(
396 &tree_data.data,
397 tree_hash,
398 None,
399 )
400 .unwrap(),
401 tree,
402 );
403 store_received_object(&dest, &tree_data).unwrap();
404 assert_eq!(dest.get_tree(&tree_hash).unwrap().unwrap(), tree);
405
406 let state_data = load_requested_object(
407 &source,
408 &ObjectRequest {
409 id: ObjectId::StateId(state.state_id),
410 have_base: None,
411 },
412 )
413 .unwrap();
414 assert_eq!(state_data.obj_type, ObjectType::State);
415 assert_eq!(
416 objects::store::codec::decode_state(&state_data.data).unwrap(),
417 state
418 );
419 store_received_object(&dest, &state_data).unwrap();
420 assert_eq!(
421 dest.get_state(&state.state_id).unwrap().unwrap().state_id,
422 state.state_id
423 );
424 }
425
426 #[test]
427 fn load_object_data_reports_missing_and_id_type_mismatch_errors() {
428 let (_temp, store) = create_test_store();
429 let missing_hash = ContentHash::from_bytes([7; 32]);
430 let missing_state = objects::object::StateId::from_bytes([9; 32]);
431
432 let missing = load_requested_object(
433 &store,
434 &ObjectRequest {
435 id: ObjectId::Hash(missing_hash),
436 have_base: None,
437 },
438 )
439 .unwrap_err();
440 assert!(
441 matches!(missing, ProtocolError::ObjectNotFound(id) if id == missing_hash.to_hex())
442 );
443
444 let missing = load_requested_object(
445 &store,
446 &ObjectRequest {
447 id: ObjectId::StateId(missing_state),
448 have_base: None,
449 },
450 )
451 .unwrap_err();
452 assert!(
453 matches!(missing, ProtocolError::ObjectNotFound(id) if id == missing_state.to_string())
454 );
455
456 let mismatch =
457 load_object_data(&store, &ObjectId::Hash(missing_hash), ObjectType::State).unwrap_err();
458 assert!(
459 matches!(mismatch, ProtocolError::InvalidState(message) if message == "object id/type mismatch")
460 );
461
462 let mismatch =
463 load_object_data(&store, &ObjectId::StateId(missing_state), ObjectType::Blob)
464 .unwrap_err();
465 assert!(
466 matches!(mismatch, ProtocolError::InvalidState(message) if message == "object id/type mismatch")
467 );
468 }
469
470 #[test]
471 fn store_received_object_rejects_mismatched_object_identity() {
472 let (_temp, store) = create_test_store();
473 let blob = Blob::from("tree leaf");
474 let blob_hash = store.put_blob(&blob).unwrap();
475 let tree = Tree::from_entries(vec![TreeEntry::file("leaf.txt", blob_hash, false).unwrap()]);
476 let tree_bytes = tree.encode_canonical().unwrap();
477 let wrong_hash = ContentHash::from_bytes([4; 32]);
478
479 let error = store_received_object(
480 &store,
481 &ObjectData {
482 id: ObjectId::Hash(wrong_hash),
483 obj_type: ObjectType::Tree,
484 data: tree_bytes,
485 is_delta: false,
486 },
487 )
488 .unwrap_err();
489 assert!(
490 matches!(error, ProtocolError::InvalidState(message) if message == "tree hash mismatch")
491 );
492
493 let state = State::new(tree.hash(), Vec::new(), test_attribution());
494 let wrong_state_id = objects::object::StateId::from_bytes([5; 32]);
495 let error = store_received_object(
496 &store,
497 &ObjectData {
498 id: ObjectId::StateId(wrong_state_id),
499 obj_type: ObjectType::State,
500 data: rmp_serde::to_vec_named(&state).unwrap(),
501 is_delta: false,
502 },
503 )
504 .unwrap_err();
505 assert!(
506 matches!(error, ProtocolError::InvalidState(message) if message.contains("StateId mismatch"))
507 );
508 }
509
510 #[test]
511 fn store_received_object_rejects_unbounded_tree_block_raw_length() {
512 const TREE_BLOCK_PREAMBLE_LEN: usize = 16;
513 const TREE_BLOCK_INDEX_RAW_LEN_OFFSET: usize = 20;
514
515 let (_temp, store) = create_test_store();
516 let tree = Tree::from_entries(
517 (0..256)
518 .map(|index| {
519 TreeEntry::file(
520 format!("crates__shared_component_{index:04}.rs"),
521 ContentHash::compute(format!("blob-{index}").as_bytes()),
522 false,
523 )
524 .expect("tree entry")
525 })
526 .collect(),
527 );
528 let mut data = tree.encode_canonical_blocked(3, 0).expect("blocked tree");
529 assert_eq!(data[4], TREE_BLOCK_ENCODING_VERSION);
530 let raw_len_offset =
531 TREE_HEADER_LEN + TREE_BLOCK_PREAMBLE_LEN + TREE_BLOCK_INDEX_RAW_LEN_OFFSET;
532 data[raw_len_offset..raw_len_offset + 4].copy_from_slice(&u32::MAX.to_le_bytes());
533
534 let error = store_received_object(
535 &store,
536 &ObjectData {
537 id: ObjectId::Hash(tree.hash()),
538 obj_type: ObjectType::Tree,
539 data,
540 is_delta: false,
541 },
542 )
543 .expect_err("received tree must reject attacker-controlled raw length");
544 assert!(
545 matches!(error, ProtocolError::InvalidState(ref message) if message.contains("raw length") && message.contains("exceeds maximum")),
546 "unexpected error: {error}"
547 );
548 }
549
550 #[test]
551 fn store_received_object_rejects_raw_sidecar_objects() {
552 let (_temp, store) = create_test_store();
553 let blob_hash = ContentHash::from_bytes([1; 32]);
554 let state_id = objects::object::StateId::from_bytes([2; 32]);
555
556 let redaction_error = store_received_object(
557 &store,
558 &ObjectData {
559 id: ObjectId::Hash(blob_hash),
560 obj_type: ObjectType::Redaction,
561 data: b"unsigned redaction bytes".to_vec(),
562 is_delta: false,
563 },
564 )
565 .unwrap_err();
566 assert!(
567 matches!(redaction_error, ProtocolError::InvalidState(message) if message.contains("signature verification is required"))
568 );
569
570 let visibility_error = store_received_object(
571 &store,
572 &ObjectData {
573 id: ObjectId::StateId(state_id),
574 obj_type: ObjectType::StateVisibility,
575 data: b"raw visibility bytes".to_vec(),
576 is_delta: false,
577 },
578 )
579 .unwrap_err();
580 assert!(
581 matches!(visibility_error, ProtocolError::InvalidState(message) if message.contains("sidecar validation is required"))
582 );
583 }
584
585 #[test]
586 fn test_chunk_count_rounds_up() {
587 assert_eq!(chunk_count(0, 64), 0);
588 assert_eq!(chunk_count(1, 64), 1);
589 assert_eq!(chunk_count(64, 64), 1);
590 assert_eq!(chunk_count(65, 64), 2);
591 }
592
593 #[test]
594 fn test_chunk_bounds_returns_ranges() {
595 assert_eq!(chunk_bounds(100, 32, 0), Some((0, 32)));
596 assert_eq!(chunk_bounds(100, 32, 2), Some((64, 32)));
597 assert_eq!(chunk_bounds(100, 32, 3), Some((96, 4)));
598 assert_eq!(chunk_bounds(100, 32, 4), None);
599 assert_eq!(chunk_bounds(100, 0, 0), None);
600 }
601
602 #[test]
603 fn test_chunk_offset_returns_position() {
604 assert_eq!(chunk_offset(0, 64), Some(0));
605 assert_eq!(chunk_offset(3, 64), Some(192));
606 assert_eq!(chunk_offset(usize::MAX, 2), None);
607 }
608
609 #[test]
610 fn received_transfer_blob_at_limit_is_accepted() {
611 check_received_transfer_blob_size(8, 8, "redactions").unwrap();
612 }
613
614 #[test]
615 fn received_transfer_blob_over_limit_is_rejected() {
616 let error = check_received_transfer_blob_size(9, 8, "redactions").unwrap_err();
617 let message = error.to_string();
618 assert!(
619 message.contains("redactions blob exceeds receive size limit"),
620 "unexpected error: {message}"
621 );
622 assert!(
623 message.contains("9 bytes (max 8)"),
624 "unexpected error: {message}"
625 );
626 }
627
628 #[test]
629 fn received_transfer_blob_caps_are_enforced_against_production_limits() {
630 check_received_transfer_blob_size(
631 MAX_RECEIVED_REDACTIONS_BLOB_SIZE as usize,
632 MAX_RECEIVED_REDACTIONS_BLOB_SIZE,
633 "redactions",
634 )
635 .unwrap();
636 check_received_transfer_blob_size(
637 MAX_RECEIVED_STATE_VISIBILITY_BLOB_SIZE as usize,
638 MAX_RECEIVED_STATE_VISIBILITY_BLOB_SIZE,
639 "state-visibility",
640 )
641 .unwrap();
642 }
643
644 #[test]
645 fn declared_receive_len_above_max_is_rejected_before_any_alloc() {
646 let error = admit_declared_received_len(9, 8, "pull raw body")
647 .expect_err("declared length above max must fail closed");
648 assert!(
649 error.to_string().contains("exceeds receive size limit"),
650 "got {error}"
651 );
652 assert!(error.to_string().contains("9 bytes (max 8)"), "got {error}");
653 }
654
655 #[test]
656 fn declared_receive_len_above_pack_cap_is_rejected_before_any_alloc() {
657 let error = admit_declared_received_len(
658 crate::MAX_RECEIVED_PACK_SIZE + 1,
659 crate::MAX_RECEIVED_PACK_SIZE,
660 "pull raw body",
661 )
662 .expect_err("attacker-chosen pack length must fail closed");
663 assert!(
664 error.to_string().contains("exceeds receive size limit"),
665 "got {error}"
666 );
667 }
668
669 #[test]
670 fn declared_receive_len_at_max_is_admitted() {
671 assert_eq!(
672 admit_declared_received_len(8, 8, "pull raw body").unwrap(),
673 8
674 );
675 }
676
677 #[test]
678 fn state_attachment_roundtrips_through_wire_data() {
679 let (_source_temp, source) = create_test_store();
680 let (_dest_temp, dest) = create_test_store();
681 let tree = source.put_tree(&Tree::new()).unwrap();
682 let state = State::new(tree, vec![], test_attribution());
683 source.put_state(&state).unwrap();
684 dest.put_state(&state).unwrap();
685 let attachment = StateAttachment {
686 state_id: state.id(),
687 body: StateAttachmentBody::RiskSignals(ContentHash::compute(b"signals")),
688 attribution: test_attribution(),
689 created_at: chrono::Utc::now(),
690 supersedes: None,
691 };
692 source.put_state_attachment(&attachment).unwrap();
693 let id = ObjectId::StateAttachment {
694 state: state.id(),
695 id: attachment.id(),
696 kind: attachment.body.kind(),
697 };
698 let data = load_object_data(&source, &id, ObjectType::StateAttachment).unwrap();
699 store_received_object(&dest, &data).unwrap();
700 assert_eq!(
701 dest.get_state_attachment(&state.id(), &attachment.id())
702 .unwrap(),
703 Some(attachment)
704 );
705 }
706}