1#![feature(
31 const_block_items,
32 const_cmp,
33 const_convert,
34 const_default,
35 const_index,
36 const_trait_impl,
37 maybe_uninit_fill,
38 trusted_len,
39 const_try,
40 const_try_residual,
41 try_blocks,
42 widening_mul
43)]
44#![expect(incomplete_features, reason = "generic_const_exprs")]
45#![feature(generic_const_exprs)]
48#![no_std]
49
50pub mod instruction;
51
52use crate::instruction::ContractInstruction;
53use ab_contracts_common::metadata::decode::{
54 MetadataDecoder, MetadataDecodingError, MetadataItem, MethodMetadataItem,
55 MethodsMetadataDecoder,
56};
57use ab_io_type::trivial_type::TrivialType;
58use ab_riscv_primitives::prelude::*;
59use core::iter;
60use core::iter::TrustedLen;
61use core::mem::MaybeUninit;
62use replace_with::replace_with_or_abort;
63use tracing::{debug, trace};
64
65pub const CONTRACT_FILE_MAGIC: [u8; 4] = *b"ABC0";
67
68const {
70 assert!(size_of::<ContractInstruction>() == 8);
71}
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq, TrivialType)]
75#[repr(C)]
76pub struct ContractFileHeader {
77 pub magic: [u8; 4],
79 pub read_only_section_file_size: u32,
81 pub read_only_section_memory_size: u32,
85 pub metadata_offset: u32,
87 pub metadata_size: u16,
89 pub num_methods: u16,
91 pub host_call_fn_offset: u32,
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq, TrivialType)]
99#[repr(C)]
100pub struct ContractFileMethodMetadata {
101 pub offset: u32,
103 pub size: u32,
105}
106
107#[derive(Debug, Copy, Clone)]
108pub struct ContractFileMethod<'a> {
109 pub address: u32,
111 pub method_metadata_item: MethodMetadataItem<'a>,
113 pub method_metadata_bytes: &'a [u8],
119}
120
121#[derive(Debug, thiserror::Error)]
123pub enum ContractFileParseError {
124 #[error("The file is too large, must fit into `u32`: {file_size} bytes")]
126 FileTooLarge {
127 file_size: usize,
129 },
130 #[error("The file does not have a header (not enough bytes)")]
132 NoHeader,
133 #[error("The magic bytes in the header are incorrect")]
135 WrongMagicBytes,
136 #[error(
138 "The metadata section is out of bounds of the file: offset {offset}, size {size}, file \
139 size {file_size}"
140 )]
141 MetadataOutOfRange {
142 offset: u32,
144 size: u16,
146 file_size: u32,
148 },
149 #[error("Failed to decode metadata item")]
151 MetadataDecoding,
152 #[error(
154 "The file is too small: num_methods {num_methods}, read_only_section_size \
155 {read_only_section_size}, file_size {file_size}"
156 )]
157 FileTooSmall {
158 num_methods: u16,
160 read_only_section_size: u32,
162 file_size: u32,
164 },
165 #[error("Method is unaligned: file offset {file_offset}, memory address {memory_address}")]
167 MethodUnaligned {
168 file_offset: u32,
170 memory_address: u32,
172 },
173 #[error(
175 "Method offset is out of bounds of the file: offset {offset}, code section \
176 offset {code_section_offset} file_size {file_size}"
177 )]
178 MethodOutOfRange {
179 offset: u32,
181 code_section_offset: u32,
183 file_size: u32,
185 },
186 #[error(
188 "Host call function is unaligned: file offset {file_offset}, memory address \
189 {memory_address}"
190 )]
191 HostCallFnUnaligned {
192 file_offset: u32,
194 memory_address: u32,
196 },
197 #[error(
199 "The host call function offset is out of bounds of the file: offset {offset}, code section \
200 offset {code_section_offset} file_size {file_size}"
201 )]
202 HostCallFnOutOfRange {
203 offset: u32,
205 code_section_offset: u32,
207 file_size: u32,
209 },
210 #[error("The host call function doesn't have auipc + jalr tailcall pattern: {first} {second}")]
212 InvalidHostCallFnPattern {
213 first: ContractInstruction,
215 second: ContractInstruction,
217 },
218 #[error(
220 "The read-only section file size is larger than the memory size: file_size {file_size}, \
221 memory_size {memory_size}"
222 )]
223 InvalidReadOnlySizes {
224 file_size: u32,
226 memory_size: u32,
228 },
229 #[error(
232 "There are not enough methods in the header to match the number of methods in the actual \
233 metadata: header_num_methods {header_num_methods}, metadata_method_index \
234 {metadata_method_index}"
235 )]
236 InsufficientHeaderMethods {
237 header_num_methods: u16,
239 metadata_method_index: u16,
241 },
242 #[error(
245 "The number of methods in the header {header_num_methods} does not match the number of \
246 methods in the actual metadata {metadata_num_methods}"
247 )]
248 MetadataNumMethodsMismatch {
249 header_num_methods: u16,
251 metadata_num_methods: u16,
253 },
254 #[error("Invalid instruction encountered while parsing the code section: {instruction:#x}")]
256 InvalidInstruction {
257 instruction: u32,
259 },
260 #[error("The code section is empty")]
262 CodeEmpty,
263 #[error(
265 "Unexpected trailing code bytes encountered while parsing the code section: {num_bytes} \
266 trailing bytes"
267 )]
268 UnexpectedTrailingCodeBytes {
269 num_bytes: usize,
271 },
272 #[error("The last instruction in the code section must be a jump instruction: {instruction}")]
274 LastInstructionMustBeJump {
275 instruction: ContractInstruction,
277 },
278}
279
280impl From<MetadataDecodingError<'_>> for ContractFileParseError {
281 fn from(error: MetadataDecodingError<'_>) -> Self {
282 debug!(?error, "Failed to decode metadata item");
283 Self::MetadataDecoding
284 }
285}
286
287#[derive(Debug)]
289pub struct ContractFile<'a> {
290 read_only_section_file_size: u32,
291 read_only_section_memory_size: u32,
292 num_methods: u16,
293 bytes: &'a [u8],
294}
295
296impl<'a> ContractFile<'a> {
297 pub fn parse<CM>(
306 file_bytes: &'a [u8],
307 mut contract_method: CM,
308 ) -> Result<Self, ContractFileParseError>
309 where
310 CM: FnMut(ContractFileMethod<'a>) -> Result<(), ContractFileParseError>,
311 {
312 let file_size = u32::try_from(file_bytes.len()).map_err(|_error| {
313 ContractFileParseError::FileTooLarge {
314 file_size: file_bytes.len(),
315 }
316 })?;
317 let (header_bytes, after_header_bytes) = file_bytes
318 .split_at_checked(size_of::<ContractFileHeader>())
319 .ok_or(ContractFileParseError::NoHeader)?;
320 let header = unsafe { ContractFileHeader::read_unaligned_unchecked(header_bytes) };
322
323 if header.magic != CONTRACT_FILE_MAGIC {
324 return Err(ContractFileParseError::WrongMagicBytes);
325 }
326
327 if header.read_only_section_file_size > header.read_only_section_memory_size {
328 return Err(ContractFileParseError::InvalidReadOnlySizes {
329 file_size: header.read_only_section_file_size,
330 memory_size: header.read_only_section_memory_size,
331 });
332 }
333
334 let metadata_bytes = file_bytes
335 .get(header.metadata_offset as usize..)
336 .ok_or(ContractFileParseError::MetadataOutOfRange {
337 offset: header.metadata_offset,
338 size: header.metadata_size,
339 file_size,
340 })?
341 .get(..header.metadata_size as usize)
342 .ok_or(ContractFileParseError::MetadataOutOfRange {
343 offset: header.metadata_offset,
344 size: header.metadata_size,
345 file_size,
346 })?;
347
348 let read_only_padding_size =
349 header.read_only_section_memory_size - header.read_only_section_file_size;
350 let read_only_section_offset = ContractFileHeader::SIZE
351 + u32::from(header.num_methods) * ContractFileMethodMetadata::SIZE;
352 let code_section_offset =
353 read_only_section_offset.saturating_add(header.read_only_section_file_size);
354
355 {
356 let mut contract_file_methods_metadata_iter = {
357 let mut file_contract_metadata_bytes = after_header_bytes;
358
359 (0..header.num_methods).map(move |_| {
360 let contract_file_method_metadata_bytes = file_contract_metadata_bytes
361 .split_off(..size_of::<ContractFileMethodMetadata>())
362 .ok_or(ContractFileParseError::FileTooSmall {
363 num_methods: header.num_methods,
364 read_only_section_size: header.read_only_section_file_size,
365 file_size,
366 })?;
367 let contract_file_method_metadata = unsafe {
369 ContractFileMethodMetadata::read_unaligned_unchecked(
370 contract_file_method_metadata_bytes,
371 )
372 };
373
374 if (contract_file_method_metadata.offset + contract_file_method_metadata.size)
375 > file_size
376 {
377 return Err(ContractFileParseError::FileTooSmall {
378 num_methods: header.num_methods,
379 read_only_section_size: header.read_only_section_file_size,
380 file_size,
381 });
382 }
383
384 if contract_file_method_metadata.offset < code_section_offset {
385 return Err(ContractFileParseError::MethodOutOfRange {
386 offset: contract_file_method_metadata.offset,
387 code_section_offset,
388 file_size,
389 });
390 }
391
392 Ok(contract_file_method_metadata)
393 })
394 };
395
396 let mut metadata_num_methods = 0;
397 let mut remaining_metadata_bytes = metadata_bytes;
398 let mut metadata_decoder = MetadataDecoder::new(metadata_bytes);
399
400 while let Some(maybe_metadata_item) = metadata_decoder.decode_next() {
401 let metadata_item = maybe_metadata_item?;
402 trace!(?metadata_item, "Decoded metadata item");
403
404 let mut methods_metadata_decoder = metadata_item.into_decoder();
405 loop {
406 let Some(method_metadata_decoder) = methods_metadata_decoder.decode_next()
410 else {
411 break;
412 };
413
414 let before_remaining_bytes = method_metadata_decoder.remaining_metadata_bytes();
415 let (_, method_metadata_item) = method_metadata_decoder.decode_next()?;
416
417 trace!(?method_metadata_item, "Decoded method metadata item");
418 metadata_num_methods += 1;
419
420 let method_metadata_bytes = remaining_metadata_bytes
421 .split_off(
422 ..before_remaining_bytes
423 - methods_metadata_decoder.remaining_metadata_bytes(),
424 )
425 .ok_or(MetadataDecodingError::NotEnoughMetadata)?;
426
427 let contract_file_method_metadata = contract_file_methods_metadata_iter
428 .next()
429 .ok_or(ContractFileParseError::InsufficientHeaderMethods {
430 header_num_methods: header.num_methods,
431 metadata_method_index: metadata_num_methods - 1,
432 })??;
433 let address = contract_file_method_metadata.offset - read_only_section_offset
434 + read_only_padding_size;
435
436 if !address.is_multiple_of(size_of::<u16>() as u32) {
437 return Err(ContractFileParseError::MethodUnaligned {
438 file_offset: contract_file_method_metadata.offset,
439 memory_address: address,
440 });
441 }
442
443 contract_method(ContractFileMethod {
444 address,
445 method_metadata_item,
446 method_metadata_bytes,
447 })?;
448 }
449 }
450
451 if metadata_num_methods != header.num_methods {
452 return Err(ContractFileParseError::MetadataNumMethodsMismatch {
453 header_num_methods: header.num_methods,
454 metadata_num_methods,
455 });
456 }
457 }
458
459 if code_section_offset >= file_size {
460 return Err(ContractFileParseError::FileTooSmall {
461 num_methods: header.num_methods,
462 read_only_section_size: header.read_only_section_file_size,
463 file_size,
464 });
465 }
466
467 if header.host_call_fn_offset != 0 {
468 if header.host_call_fn_offset >= file_size
469 || header.host_call_fn_offset < code_section_offset
470 {
471 return Err(ContractFileParseError::HostCallFnOutOfRange {
472 offset: header.host_call_fn_offset,
473 code_section_offset,
474 file_size,
475 });
476 }
477
478 let instructions_bytes = file_bytes
479 .get(header.host_call_fn_offset as usize..)
480 .ok_or(ContractFileParseError::HostCallFnOutOfRange {
481 offset: header.host_call_fn_offset,
482 code_section_offset,
483 file_size,
484 })?
485 .get(..size_of::<[u32; 2]>())
486 .ok_or(ContractFileParseError::HostCallFnOutOfRange {
487 offset: header.host_call_fn_offset,
488 code_section_offset,
489 file_size,
490 })?;
491
492 let first_instruction = u32::from_le_bytes([
493 instructions_bytes[0],
494 instructions_bytes[1],
495 instructions_bytes[2],
496 instructions_bytes[3],
497 ]);
498 let second_instruction = u32::from_le_bytes([
499 instructions_bytes[4],
500 instructions_bytes[5],
501 instructions_bytes[6],
502 instructions_bytes[7],
503 ]);
504
505 let first = ContractInstruction::try_decode(first_instruction).ok_or(
506 ContractFileParseError::InvalidInstruction {
507 instruction: first_instruction,
508 },
509 )?;
510 let second = ContractInstruction::try_decode(second_instruction).ok_or(
511 ContractFileParseError::InvalidInstruction {
512 instruction: second_instruction,
513 },
514 )?;
515
516 let matches_expected_pattern = if let (
522 ContractInstruction::Auipc {
523 rd: auipc_rd,
524 imm: _,
525 rs1: _,
526 rs2: _,
527 },
528 ContractInstruction::Jalr {
529 rd: jalr_rd,
530 rs1: jalr_rs1,
531 imm: _,
532 rs2: _,
533 },
534 ) = (first, second)
535 {
536 auipc_rd == jalr_rs1 && jalr_rd == Register::ZERO
537 } else {
538 false
539 };
540
541 if !matches_expected_pattern {
542 return Err(ContractFileParseError::InvalidHostCallFnPattern { first, second });
543 }
544
545 let address =
546 header.host_call_fn_offset - read_only_section_offset + read_only_padding_size;
547
548 if !address.is_multiple_of(size_of::<u16>() as u32) {
549 return Err(ContractFileParseError::HostCallFnUnaligned {
550 file_offset: header.host_call_fn_offset,
551 memory_address: address,
552 });
553 }
554 }
555
556 {
558 let mut offset = code_section_offset as usize;
559
560 let mut instruction = ContractInstruction::Unimp {
561 rs1: Register::ZERO,
562 rs2: Register::ZERO,
563 };
564 while offset < file_bytes.len() {
565 let remaining = &file_bytes[offset..];
566
567 let instruction_word = if remaining.len() >= size_of::<u32>() {
568 u32::from_le_bytes([remaining[0], remaining[1], remaining[2], remaining[3]])
569 } else if remaining.len() >= size_of::<u16>() {
570 u32::from_le_bytes([remaining[0], remaining[1], 0, 0])
571 } else {
572 return Err(ContractFileParseError::UnexpectedTrailingCodeBytes {
574 num_bytes: remaining.len(),
575 });
576 };
577
578 instruction = ContractInstruction::try_decode(instruction_word).ok_or(
579 ContractFileParseError::InvalidInstruction {
580 instruction: instruction_word,
581 },
582 )?;
583
584 offset += usize::from(instruction.size());
585 }
586
587 if !instruction.is_jump() {
588 return Err(ContractFileParseError::LastInstructionMustBeJump { instruction });
589 }
590 }
591
592 Ok(Self {
593 read_only_section_file_size: header.read_only_section_file_size,
594 read_only_section_memory_size: header.read_only_section_memory_size,
595 num_methods: header.num_methods,
596 bytes: file_bytes,
597 })
598 }
599
600 pub unsafe fn parse_unchecked(file_bytes: &'a [u8]) -> Self {
608 let header = unsafe { ContractFileHeader::read_unaligned_unchecked(file_bytes) };
610
611 Self {
612 read_only_section_file_size: header.read_only_section_file_size,
613 read_only_section_memory_size: header.read_only_section_memory_size,
614 num_methods: header.num_methods,
615 bytes: file_bytes,
616 }
617 }
618
619 #[inline(always)]
621 pub fn header(&self) -> ContractFileHeader {
622 unsafe { ContractFileHeader::read_unaligned_unchecked(self.bytes) }
624 }
625
626 #[inline]
628 pub fn metadata_bytes(&self) -> &[u8] {
629 let header = self.header();
630 unsafe {
632 self.bytes
633 .get_unchecked(header.metadata_offset as usize..)
634 .get_unchecked(..header.metadata_size as usize)
635 }
636 }
637
638 #[inline]
640 pub fn contract_memory_size(&self) -> u32 {
641 let read_only_section_offset = ContractFileHeader::SIZE
642 + u32::from(self.num_methods) * ContractFileMethodMetadata::SIZE;
643 let read_only_padding_size =
644 self.read_only_section_memory_size - self.read_only_section_file_size;
645 self.bytes.len() as u32 - read_only_section_offset + read_only_padding_size
646 }
647
648 #[must_use = "Must check that contract memory was large enough"]
652 pub fn initialize_contract_memory(&self, mut contract_memory: &mut [MaybeUninit<u8>]) -> bool {
653 let contract_memory_input_size = contract_memory.len();
654 let read_only_section_offset = ContractFileHeader::SIZE
655 + u32::from(self.num_methods) * ContractFileMethodMetadata::SIZE;
656 let read_only_padding_size =
657 self.read_only_section_memory_size - self.read_only_section_file_size;
658
659 let source_bytes = unsafe {
661 self.bytes
662 .get_unchecked(read_only_section_offset as usize..)
663 };
664
665 if contract_memory.len() == source_bytes.len() {
667 contract_memory.write_copy_of_slice(source_bytes);
668 return true;
669 }
670
671 let Some(read_only_file_target_bytes) =
672 contract_memory.split_off_mut(..self.read_only_section_file_size as usize)
673 else {
674 trace!(
675 %contract_memory_input_size,
676 contract_memory_size = %self.contract_memory_size(),
677 read_only_section_file_size = self.read_only_section_file_size,
678 "Not enough bytes to write read-only section from the file"
679 );
680
681 return false;
682 };
683
684 let (read_only_file_source_bytes, code_source_bytes) =
686 unsafe { source_bytes.split_at_unchecked(self.read_only_section_file_size as usize) };
687 read_only_file_target_bytes.write_copy_of_slice(read_only_file_source_bytes);
689
690 let Some(read_only_padding_bytes) =
691 contract_memory.split_off_mut(..read_only_padding_size as usize)
692 else {
693 trace!(
694 %contract_memory_input_size,
695 contract_memory_size = %self.contract_memory_size(),
696 read_only_section_file_size = self.read_only_section_file_size,
697 read_only_section_memory_size = self.read_only_section_memory_size,
698 %read_only_padding_size,
699 "Not enough bytes to write read-only padding section"
700 );
701
702 return false;
703 };
704
705 read_only_padding_bytes.write_filled(0);
707
708 if code_source_bytes.len() != contract_memory.len() {
709 trace!(
710 %contract_memory_input_size,
711 contract_memory_size = %self.contract_memory_size(),
712 read_only_section_file_size = self.read_only_section_file_size,
713 read_only_section_memory_size = self.read_only_section_memory_size,
714 %read_only_padding_size,
715 code_size = %code_source_bytes.len(),
716 "Not enough bytes to write code section from the file"
717 );
718
719 return false;
720 }
721
722 contract_memory.write_copy_of_slice(code_source_bytes);
723
724 true
725 }
726
727 pub fn get_code(&self) -> &[u8] {
729 let read_only_section_offset = ContractFileHeader::SIZE
730 + u32::from(self.num_methods) * ContractFileMethodMetadata::SIZE;
731
732 let source_bytes = unsafe {
734 self.bytes
735 .get_unchecked(read_only_section_offset as usize..)
736 };
737
738 let (_read_only_file_source_bytes, code_source_bytes) =
740 unsafe { source_bytes.split_at_unchecked(self.read_only_section_file_size as usize) };
741
742 code_source_bytes
743 }
744
745 pub fn iterate_methods(
747 &self,
748 ) -> impl ExactSizeIterator<Item = ContractFileMethod<'_>> + TrustedLen {
749 let metadata_bytes = self.metadata_bytes();
750
751 #[ouroboros::self_referencing]
752 struct MethodsMetadataIterState<'metadata> {
753 metadata_decoder: MetadataDecoder<'metadata>,
754 #[borrows(mut metadata_decoder)]
755 #[covariant]
756 methods_metadata_decoder: Option<MethodsMetadataDecoder<'this, 'metadata>>,
757 }
758
759 let metadata_decoder = MetadataDecoder::new(metadata_bytes);
760
761 let mut methods_metadata_state =
762 MethodsMetadataIterState::new(metadata_decoder, |metadata_decoder| {
763 metadata_decoder
764 .decode_next()
765 .and_then(Result::ok)
766 .map(MetadataItem::into_decoder)
767 });
768
769 let mut metadata_methods_iter = iter::from_fn(move || {
770 loop {
771 let maybe_next_item = methods_metadata_state.with_methods_metadata_decoder_mut(
772 |maybe_methods_metadata_decoder| {
773 let methods_metadata_decoder = maybe_methods_metadata_decoder.as_mut()?;
774 let method_metadata_decoder = methods_metadata_decoder.decode_next()?;
775
776 let before_remaining_bytes =
777 method_metadata_decoder.remaining_metadata_bytes();
778
779 let (_, method_metadata_item) = method_metadata_decoder
780 .decode_next()
781 .expect("Input is valid according to function contract; qed");
782
783 let method_metadata_bytes = unsafe {
785 metadata_bytes
786 .get_unchecked(metadata_bytes.len() - before_remaining_bytes..)
787 .get_unchecked(
788 ..before_remaining_bytes
789 - methods_metadata_decoder.remaining_metadata_bytes(),
790 )
791 };
792
793 Some((method_metadata_item, method_metadata_bytes))
794 },
795 );
796
797 if let Some(next_item) = maybe_next_item {
798 return Some(next_item);
799 }
800
801 replace_with_or_abort(&mut methods_metadata_state, |methods_metadata_state| {
803 let metadata_decoder = methods_metadata_state.into_heads().metadata_decoder;
804 MethodsMetadataIterState::new(metadata_decoder, |metadata_decoder| {
805 metadata_decoder
806 .decode_next()
807 .and_then(Result::ok)
808 .map(MetadataItem::into_decoder)
809 })
810 });
811
812 if methods_metadata_state
813 .borrow_methods_metadata_decoder()
814 .is_none()
815 {
816 return None;
817 }
818 }
819 });
820
821 let read_only_padding_size =
822 self.read_only_section_memory_size - self.read_only_section_file_size;
823 let contract_file_methods_metadata_bytes =
825 unsafe { self.bytes.get_unchecked(size_of::<ContractFileHeader>()..) };
826
827 (0..self.num_methods).map(move |method_index| {
828 let contract_file_method_metadata_bytes = unsafe {
830 contract_file_methods_metadata_bytes
831 .get_unchecked(
832 method_index as usize * size_of::<ContractFileMethodMetadata>()..,
833 )
834 .get_unchecked(..size_of::<ContractFileMethodMetadata>())
835 };
836 let contract_file_method_metadata = unsafe {
838 ContractFileMethodMetadata::read_unaligned_unchecked(
839 contract_file_method_metadata_bytes,
840 )
841 };
842
843 let (method_metadata_item, method_metadata_bytes) = metadata_methods_iter
844 .next()
845 .expect("Protected internal invariant checked in constructor; qed");
846
847 ContractFileMethod {
848 address: contract_file_method_metadata.offset + read_only_padding_size,
849 method_metadata_item,
850 method_metadata_bytes,
851 }
852 })
853 }
854}