diff --git a/packages/fuels-programs/src/assembly/cursor.rs b/packages/fuels-programs/src/assembly/cursor.rs index c5cb4cfa7..327e9697c 100644 --- a/packages/fuels-programs/src/assembly/cursor.rs +++ b/packages/fuels-programs/src/assembly/cursor.rs @@ -28,14 +28,20 @@ impl<'a> WasmFriendlyCursor<'a> { &mut self, ctx: &'static str, ) -> Result<[u8; AMOUNT]> { - Ok(self + let data = self .consume(AMOUNT, ctx)? .try_into() - .expect("should have failed if not enough data")) + .expect("should have failed if not enough data"); + + Ok(data) } - pub fn consume_all(&self) -> &'a [u8] { - self.data + pub fn consume_all(&mut self) -> &'a [u8] { + let data = self.data; + + self.data = &[]; + + data } pub fn unconsumed(&self) -> usize { diff --git a/packages/fuels-programs/src/assembly/script_and_predicate_loader.rs b/packages/fuels-programs/src/assembly/script_and_predicate_loader.rs index 23223ca29..f491d025a 100644 --- a/packages/fuels-programs/src/assembly/script_and_predicate_loader.rs +++ b/packages/fuels-programs/src/assembly/script_and_predicate_loader.rs @@ -11,16 +11,22 @@ use itertools::Itertools; use crate::assembly::cursor::WasmFriendlyCursor; pub struct LoaderCode { - data_section: Vec, blob_id: BlobId, + code: Vec, + data_offset: usize, } impl LoaderCode { pub fn from_normal_binary(binary: Vec) -> Result { - let (code, data_section) = split_at_data_offset(&binary)?; + let (original_code, data_section) = split_at_data_offset(&binary)?; + + let blob_id = Blob::from(original_code.to_vec()).id(); + let (loader_code, data_offset) = Self::to_bytes_w_offset(blob_id, data_section); + Ok(Self { - data_section: data_section.to_owned(), - blob_id: Blob::from(code.to_vec()).id(), + blob_id, + code: loader_code, + data_offset, }) } @@ -57,9 +63,14 @@ impl LoaderCode { let blob_id = script_cursor.consume_fixed("blob id")?; let _data_section_len = script_cursor.consume(WORD_SIZE, "data section len")?; + let data_section_offset = binary + .len() + .checked_sub(script_cursor.unconsumed()) + .expect("must be less or eq"); Ok(Some(Self { - data_section: script_cursor.consume_all().to_vec(), + data_offset: data_section_offset, + code: binary.to_vec(), blob_id, })) } @@ -69,22 +80,30 @@ impl LoaderCode { Ok(Blob::from(code.to_vec())) } - pub fn to_bytes_w_offset(&self) -> (Vec, usize) { + pub fn as_bytes(&self) -> &[u8] { + &self.code + } + + pub fn data_section_offset(&self) -> usize { + self.data_offset + } + + fn to_bytes_w_offset(blob_id: BlobId, data_section: &[u8]) -> (Vec, usize) { // The final code is going to have this structure (if the data section is non-empty): // 1. loader instructions // 2. blob id // 3. length_of_data_section // 4. the data_section (updated with configurables as needed) - if !self.data_section.is_empty() { + if !data_section.is_empty() { let instruction_bytes = loader_instructions() .into_iter() .flat_map(|instruction| instruction.to_bytes()) .collect_vec(); - let blob_bytes = self.blob_id.iter().copied().collect_vec(); + let blob_bytes = blob_id.iter().copied().collect_vec(); - let original_data_section_len_encoded = u64::try_from(self.data_section.len()) + let original_data_section_len_encoded = u64::try_from(data_section.len()) .expect("data section to be less than u64::MAX") .to_be_bytes(); @@ -98,7 +117,7 @@ impl LoaderCode { .into_iter() .chain(blob_bytes) .chain(original_data_section_len_encoded) - .chain(self.data_section.clone()) + .chain(data_section.to_vec()) .collect(); (code, new_data_section_offset) @@ -108,7 +127,7 @@ impl LoaderCode { .flat_map(|instruction| instruction.to_bytes()); let code = instruction_bytes - .chain(self.blob_id.iter().copied()) + .chain(blob_id.iter().copied()) .collect_vec(); // there is no data section, so we point the offset to the end of the file let new_data_section_offset = code.len(); diff --git a/packages/fuels-programs/src/debug.rs b/packages/fuels-programs/src/debug.rs index 912d5ba0f..bcbc78cca 100644 --- a/packages/fuels-programs/src/debug.rs +++ b/packages/fuels-programs/src/debug.rs @@ -148,7 +148,7 @@ fn parse_loader_script(script: &[u8], data: &[u8]) -> Result { } pub fn data_offset_in_code(&self) -> usize { - self.loader_code().to_bytes_w_offset().1 + self.loader_code().data_section_offset() } fn loader_code(&self) -> LoaderCode { @@ -131,7 +131,7 @@ impl Executable { /// Returns the code of the loader executable with configurables applied. pub fn code(&self) -> Vec { - self.loader_code().to_bytes_w_offset().0 + self.loader_code().as_bytes().to_vec() } /// A Blob containing the original executable code minus the data section. @@ -285,10 +285,7 @@ mod tests { let loader_code = loader.code(); assert_eq!( loader_code, - LoaderCode::from_normal_binary(code) - .unwrap() - .to_bytes_w_offset() - .0 + LoaderCode::from_normal_binary(code).unwrap().as_bytes() ) }