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