From ba5e572f91def6ad36a9510e948d44e5f8511d67 Mon Sep 17 00:00:00 2001 From: Philip Craig Date: Sat, 3 Aug 2024 12:48:37 +1000 Subject: [PATCH] read: implement Iterator for more types This also changes the underlying `next` methods to fuse on error. --- src/read/elf/attributes.rs | 52 ++++++++++++++++---- src/read/elf/note.rs | 67 ++++++++++++++++---------- src/read/elf/version.rs | 87 +++++++++++++++++++++++++++++++--- src/read/macho/load_command.rs | 24 ++++++++-- src/read/pe/import.rs | 70 +++++++++++++++++++++------ src/read/pe/relocation.rs | 21 +++++++- 6 files changed, 261 insertions(+), 60 deletions(-) diff --git a/src/read/elf/attributes.rs b/src/read/elf/attributes.rs index ddddfb25..39bc2219 100644 --- a/src/read/elf/attributes.rs +++ b/src/read/elf/attributes.rs @@ -70,14 +70,14 @@ impl<'data, Elf: FileHeader> AttributesSubsectionIterator<'data, Elf> { return Ok(None); } - let result = self.parse(); + let result = self.parse().map(Some); if result.is_err() { self.data = Bytes(&[]); } result } - fn parse(&mut self) -> Result>> { + fn parse(&mut self) -> Result> { // First read the subsection length. let mut data = self.data; let length = data @@ -94,16 +94,25 @@ impl<'data, Elf: FileHeader> AttributesSubsectionIterator<'data, Elf> { data.skip(4) .read_error("Invalid ELF attributes subsection length")?; + // TODO: errors here should not prevent reading the next subsection. let vendor = data .read_string() .read_error("Invalid ELF attributes vendor")?; - Ok(Some(AttributesSubsection { + Ok(AttributesSubsection { endian: self.endian, length, vendor, data, - })) + }) + } +} + +impl<'data, Elf: FileHeader> Iterator for AttributesSubsectionIterator<'data, Elf> { + type Item = Result>; + + fn next(&mut self) -> Option { + self.next().transpose() } } @@ -153,14 +162,14 @@ impl<'data, Elf: FileHeader> AttributesSubsubsectionIterator<'data, Elf> { return Ok(None); } - let result = self.parse(); + let result = self.parse().map(Some); if result.is_err() { self.data = Bytes(&[]); } result } - fn parse(&mut self) -> Result>> { + fn parse(&mut self) -> Result> { // The format of a sub-section looks like this: // // * @@ -184,6 +193,7 @@ impl<'data, Elf: FileHeader> AttributesSubsubsectionIterator<'data, Elf> { data.skip(1 + 4) .read_error("Invalid ELF attributes sub-subsection length")?; + // TODO: errors here should not prevent reading the next sub-subsection. let indices = if tag == elf::Tag_Section || tag == elf::Tag_Symbol { data.read_string() .map(Bytes) @@ -194,12 +204,20 @@ impl<'data, Elf: FileHeader> AttributesSubsubsectionIterator<'data, Elf> { return Err(Error("Unimplemented ELF attributes sub-subsection tag")); }; - Ok(Some(AttributesSubsubsection { + Ok(AttributesSubsubsection { tag, length, indices, data, - })) + }) + } +} + +impl<'data, Elf: FileHeader> Iterator for AttributesSubsubsectionIterator<'data, Elf> { + type Item = Result>; + + fn next(&mut self) -> Option { + self.next().transpose() } } @@ -263,6 +281,15 @@ impl<'data> AttributeIndexIterator<'data> { if self.data.is_empty() { return Ok(None); } + + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse(&mut self) -> Result { let err = "Invalid ELF attribute index"; self.data .read_uleb128() @@ -270,7 +297,14 @@ impl<'data> AttributeIndexIterator<'data> { .try_into() .map_err(|_| ()) .read_error(err) - .map(Some) + } +} + +impl<'data> Iterator for AttributeIndexIterator<'data> { + type Item = Result; + + fn next(&mut self) -> Option { + self.next().transpose() } } diff --git a/src/read/elf/note.rs b/src/read/elf/note.rs index 72a8c80e..bf4a3aed 100644 --- a/src/read/elf/note.rs +++ b/src/read/elf/note.rs @@ -51,19 +51,28 @@ where /// Returns the next note. pub fn next(&mut self) -> read::Result>> { - let mut data = self.data; - if data.is_empty() { + if self.data.is_empty() { return Ok(None); } - let header = data + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse(&mut self) -> read::Result> { + let header = self + .data .read_at::(0) .read_error("ELF note is too short")?; // The name has no alignment requirement. let offset = mem::size_of::(); let namesz = header.n_namesz(self.endian) as usize; - let name = data + let name = self + .data .read_bytes_at(offset, namesz) .read_error("Invalid ELF note namesz")? .0; @@ -71,33 +80,27 @@ where // The descriptor must be aligned. let offset = util::align(offset + namesz, self.align); let descsz = header.n_descsz(self.endian) as usize; - let desc = data + let desc = self + .data .read_bytes_at(offset, descsz) .read_error("Invalid ELF note descsz")? .0; // The next note (if any) must be aligned. let offset = util::align(offset + descsz, self.align); - if data.skip(offset).is_err() { - data = Bytes(&[]); + if self.data.skip(offset).is_err() { + self.data = Bytes(&[]); } - self.data = data; - Ok(Some(Note { header, name, desc })) + Ok(Note { header, name, desc }) } } impl<'data, Elf: FileHeader> Iterator for NoteIterator<'data, Elf> { type Item = read::Result>; + fn next(&mut self) -> Option { - match self.next() { - Err(e) => { - self.data = Bytes(&[]); - Some(Err(e)) - } - Ok(Some(v)) => Some(Ok(v)), - Ok(None) => None, - } + self.next().transpose() } } @@ -238,23 +241,37 @@ pub struct GnuPropertyIterator<'data, Endian: endian::Endian> { impl<'data, Endian: endian::Endian> GnuPropertyIterator<'data, Endian> { /// Returns the next property. pub fn next(&mut self) -> read::Result>> { - let mut data = self.data; - if data.is_empty() { + if self.data.is_empty() { return Ok(None); } + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse(&mut self) -> read::Result> { (|| -> Result<_, ()> { - let pr_type = data.read_at::>(0)?.get(self.endian); - let pr_datasz = data.read_at::>(4)?.get(self.endian) as usize; - let pr_data = data.read_bytes_at(8, pr_datasz)?.0; - data.skip(util::align(8 + pr_datasz, self.align))?; - self.data = data; - Ok(Some(GnuProperty { pr_type, pr_data })) + let pr_type = self.data.read_at::>(0)?.get(self.endian); + let pr_datasz = self.data.read_at::>(4)?.get(self.endian) as usize; + let pr_data = self.data.read_bytes_at(8, pr_datasz)?.0; + self.data.skip(util::align(8 + pr_datasz, self.align))?; + Ok(GnuProperty { pr_type, pr_data }) })() .read_error("Invalid ELF GNU property") } } +impl<'data, Endian: endian::Endian> Iterator for GnuPropertyIterator<'data, Endian> { + type Item = read::Result>; + + fn next(&mut self) -> Option { + self.next().transpose() + } +} + /// A property in a [`elf::NT_GNU_PROPERTY_TYPE_0`] note. #[derive(Debug)] pub struct GnuProperty<'data> { diff --git a/src/read/elf/version.rs b/src/read/elf/version.rs index d0bd504e..2e350a8b 100644 --- a/src/read/elf/version.rs +++ b/src/read/elf/version.rs @@ -252,6 +252,14 @@ impl<'data, Elf: FileHeader> VerdefIterator<'data, Elf> { return Ok(None); } + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse(&mut self) -> Result<(&'data elf::Verdef, VerdauxIterator<'data, Elf>)> { let verdef = self .data .read_at::>(0) @@ -272,7 +280,15 @@ impl<'data, Elf: FileHeader> VerdefIterator<'data, Elf> { } else { self.data = Bytes(&[]); } - Ok(Some((verdef, verdaux))) + Ok((verdef, verdaux)) + } +} + +impl<'data, Elf: FileHeader> Iterator for VerdefIterator<'data, Elf> { + type Item = Result<(&'data elf::Verdef, VerdauxIterator<'data, Elf>)>; + + fn next(&mut self) -> Option { + self.next().transpose() } } @@ -299,6 +315,16 @@ impl<'data, Elf: FileHeader> VerdauxIterator<'data, Elf> { return Ok(None); } + let result = self.parse().map(Some); + if result.is_err() { + self.count = 0; + } else { + self.count -= 1; + } + result + } + + fn parse(&mut self) -> Result<&'data elf::Verdaux> { let verdaux = self .data .read_at::>(0) @@ -307,8 +333,15 @@ impl<'data, Elf: FileHeader> VerdauxIterator<'data, Elf> { self.data .skip(verdaux.vda_next.get(self.endian) as usize) .read_error("Invalid ELF vda_next")?; - self.count -= 1; - Ok(Some(verdaux)) + Ok(verdaux) + } +} + +impl<'data, Elf: FileHeader> Iterator for VerdauxIterator<'data, Elf> { + type Item = Result<&'data elf::Verdaux>; + + fn next(&mut self) -> Option { + self.next().transpose() } } @@ -340,6 +373,19 @@ impl<'data, Elf: FileHeader> VerneedIterator<'data, Elf> { return Ok(None); } + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse( + &mut self, + ) -> Result<( + &'data elf::Verneed, + VernauxIterator<'data, Elf>, + )> { let verneed = self .data .read_at::>(0) @@ -360,7 +406,18 @@ impl<'data, Elf: FileHeader> VerneedIterator<'data, Elf> { } else { self.data = Bytes(&[]); } - Ok(Some((verneed, vernaux))) + Ok((verneed, vernaux)) + } +} + +impl<'data, Elf: FileHeader> Iterator for VerneedIterator<'data, Elf> { + type Item = Result<( + &'data elf::Verneed, + VernauxIterator<'data, Elf>, + )>; + + fn next(&mut self) -> Option { + self.next().transpose() } } @@ -387,16 +444,32 @@ impl<'data, Elf: FileHeader> VernauxIterator<'data, Elf> { return Ok(None); } + let result = self.parse().map(Some); + if result.is_err() { + self.count = 0; + } else { + self.count -= 1; + } + result + } + + fn parse(&mut self) -> Result<&'data elf::Vernaux> { let vernaux = self .data .read_at::>(0) .read_error("ELF vernaux is too short")?; - self.data .skip(vernaux.vna_next.get(self.endian) as usize) .read_error("Invalid ELF vna_next")?; - self.count -= 1; - Ok(Some(vernaux)) + Ok(vernaux) + } +} + +impl<'data, Elf: FileHeader> Iterator for VernauxIterator<'data, Elf> { + type Item = Result<&'data elf::Vernaux>; + + fn next(&mut self) -> Option { + self.next().transpose() } } diff --git a/src/read/macho/load_command.rs b/src/read/macho/load_command.rs index bc49530d..1429e1d9 100644 --- a/src/read/macho/load_command.rs +++ b/src/read/macho/load_command.rs @@ -29,6 +29,17 @@ impl<'data, E: Endian> LoadCommandIterator<'data, E> { if self.ncmds == 0 { return Ok(None); } + + let result = self.parse().map(Some); + if result.is_err() { + self.ncmds = 0; + } else { + self.ncmds -= 1; + } + result + } + + fn parse(&mut self) -> Result> { let header = self .data .read_at::>(0) @@ -42,12 +53,19 @@ impl<'data, E: Endian> LoadCommandIterator<'data, E> { .data .read_bytes(cmdsize) .read_error("Invalid Mach-O load command size")?; - self.ncmds -= 1; - Ok(Some(LoadCommandData { + Ok(LoadCommandData { cmd, data, marker: Default::default(), - })) + }) + } +} + +impl<'data, E: Endian> Iterator for LoadCommandIterator<'data, E> { + type Item = Result>; + + fn next(&mut self) -> Option { + self.next().transpose() } } diff --git a/src/read/pe/import.rs b/src/read/pe/import.rs index f81a084d..c35dc986 100644 --- a/src/read/pe/import.rs +++ b/src/read/pe/import.rs @@ -42,7 +42,7 @@ impl<'data> ImportTable<'data> { let mut data = self.section_data; data.skip(offset as usize) .read_error("Invalid PE import descriptor address")?; - Ok(ImportDescriptorIterator { data }) + Ok(ImportDescriptorIterator { data, null: false }) } /// Return a library name given its address. @@ -101,6 +101,7 @@ impl<'data> ImportTable<'data> { #[derive(Debug, Clone)] pub struct ImportDescriptorIterator<'data> { data: Bytes<'data>, + null: bool, } impl<'data> ImportDescriptorIterator<'data> { @@ -108,18 +109,38 @@ impl<'data> ImportDescriptorIterator<'data> { /// /// Returns `Ok(None)` when a null descriptor is found. pub fn next(&mut self) -> Result> { - let import_desc = self + if self.null { + return Ok(None); + } + let result = self .data .read::() - .read_error("Missing PE null import descriptor")?; - if import_desc.is_null() { - Ok(None) - } else { - Ok(Some(import_desc)) + .read_error("Missing PE null import descriptor"); + match result { + Ok(import_desc) => { + if import_desc.is_null() { + self.null = true; + Ok(None) + } else { + Ok(Some(import_desc)) + } + } + Err(e) => { + self.null = true; + Err(e) + } } } } +impl<'data> Iterator for ImportDescriptorIterator<'data> { + type Item = Result<&'data pe::ImageImportDescriptor>; + + fn next(&mut self) -> Option { + self.next().transpose() + } +} + /// A list of import thunks. /// /// These may be in the import lookup table, or the import address table. @@ -256,7 +277,7 @@ impl<'data> DelayLoadImportTable<'data> { let mut data = self.section_data; data.skip(offset as usize) .read_error("Invalid PE delay-load import descriptor address")?; - Ok(DelayLoadDescriptorIterator { data }) + Ok(DelayLoadDescriptorIterator { data, null: false }) } /// Return a library name given its address. @@ -319,6 +340,7 @@ impl<'data> DelayLoadImportTable<'data> { #[derive(Debug, Clone)] pub struct DelayLoadDescriptorIterator<'data> { data: Bytes<'data>, + null: bool, } impl<'data> DelayLoadDescriptorIterator<'data> { @@ -326,14 +348,34 @@ impl<'data> DelayLoadDescriptorIterator<'data> { /// /// Returns `Ok(None)` when a null descriptor is found. pub fn next(&mut self) -> Result> { - let import_desc = self + if self.null { + return Ok(None); + } + let result = self .data .read::() - .read_error("Missing PE null delay-load import descriptor")?; - if import_desc.is_null() { - Ok(None) - } else { - Ok(Some(import_desc)) + .read_error("Missing PE null delay-load import descriptor"); + match result { + Ok(import_desc) => { + if import_desc.is_null() { + self.null = true; + Ok(None) + } else { + Ok(Some(import_desc)) + } + } + Err(e) => { + self.null = true; + Err(e) + } } } } + +impl<'data> Iterator for DelayLoadDescriptorIterator<'data> { + type Item = Result<&'data pe::ImageDelayloadDescriptor>; + + fn next(&mut self) -> Option { + self.next().transpose() + } +} diff --git a/src/read/pe/relocation.rs b/src/read/pe/relocation.rs index 77421b7b..92b3f088 100644 --- a/src/read/pe/relocation.rs +++ b/src/read/pe/relocation.rs @@ -23,6 +23,15 @@ impl<'data> RelocationBlockIterator<'data> { if self.data.is_empty() { return Ok(None); } + + let result = self.parse().map(Some); + if result.is_err() { + self.data = Bytes(&[]); + } + result + } + + fn parse(&mut self) -> Result> { let header = self .data .read::() @@ -38,11 +47,19 @@ impl<'data> RelocationBlockIterator<'data> { .read_slice::>(count as usize) .read_error("Invalid PE reloc block size")? .iter(); - Ok(Some(RelocationIterator { + Ok(RelocationIterator { virtual_address, size, relocs, - })) + }) + } +} + +impl<'data> Iterator for RelocationBlockIterator<'data> { + type Item = Result>; + + fn next(&mut self) -> Option { + self.next().transpose() } }