diff --git a/src/build/elf.rs b/src/build/elf.rs index 3fffd24..638b15e 100644 --- a/src/build/elf.rs +++ b/src/build/elf.rs @@ -1194,7 +1194,7 @@ impl<'data> Builder<'data> { } SectionData::DynamicString => { dynstr_addr = Some(section.sh_addr); - writer.reserve_dynstr() + writer.reserve_dynstr()? } SectionData::Hash => { hash_addr = Some(section.sh_addr); @@ -1272,7 +1272,7 @@ impl<'data> Builder<'data> { writer.reserve_symtab(); writer.reserve_symtab_shndx(); - writer.reserve_strtab(); + writer.reserve_strtab()?; // Reserve non-alloc relocations. for out_section in &mut out_sections { @@ -1287,7 +1287,7 @@ impl<'data> Builder<'data> { writer.reserve_relocations(relocations.len(), section.sh_type == elf::SHT_RELA); } - writer.reserve_shstrtab(); + writer.reserve_shstrtab()?; writer.reserve_section_headers(); // Start writing. diff --git a/src/write/coff/object.rs b/src/write/coff/object.rs index db6702a..e871724 100644 --- a/src/write/coff/object.rs +++ b/src/write/coff/object.rs @@ -594,7 +594,7 @@ impl<'a> Object<'a> { section_offsets[index].reloc_offset = writer.reserve_relocations(section.relocations.len()); } - writer.reserve_symtab_strtab(); + writer.reserve_symtab_strtab()?; // Start writing. writer.write_file_header(writer::FileHeader { diff --git a/src/write/coff/writer.rs b/src/write/coff/writer.rs index 42e18c6..013b9a8 100644 --- a/src/write/coff/writer.rs +++ b/src/write/coff/writer.rs @@ -35,7 +35,7 @@ pub struct Writer<'a> { symtab_num: u32, strtab: StringTable<'a>, - strtab_len: usize, + strtab_len: u32, strtab_offset: u32, strtab_data: Vec, } @@ -194,7 +194,6 @@ impl<'a> Writer<'a> { coff_section.name[0] = b'/'; coff_section.name[1..][..len].copy_from_slice(&name[7 - len..]); } else { - debug_assert!(str_offset as u64 <= 0xf_ffff_ffff); coff_section.name[0] = b'/'; coff_section.name[1] = b'/'; for i in 0..6 { @@ -330,7 +329,7 @@ impl<'a> Writer<'a> { Name::Short(name) => coff_symbol.name = name, Name::Long(str_id) => { let str_offset = self.strtab.get_offset(str_id); - coff_symbol.name[4..8].copy_from_slice(&u32::to_le_bytes(str_offset as u32)); + coff_symbol.name[4..8].copy_from_slice(&u32::to_le_bytes(str_offset)); } } self.buffer.write(&coff_symbol); @@ -440,22 +439,23 @@ impl<'a> Writer<'a> { /// /// This must be called after functions that reserve symbol /// indices or add strings. - pub fn reserve_symtab_strtab(&mut self) { + /// + /// Returns an error if the string table could not be finalized. + pub fn reserve_symtab_strtab(&mut self) -> Result<()> { debug_assert_eq!(self.symtab_offset, 0); self.symtab_offset = self.reserve(self.symtab_num as usize * pe::IMAGE_SIZEOF_SYMBOL, 1); debug_assert_eq!(self.strtab_offset, 0); // First 4 bytes of strtab are the length. - self.strtab.write(4, &mut self.strtab_data); - self.strtab_len = self.strtab_data.len() + 4; - self.strtab_offset = self.reserve(self.strtab_len, 1); + self.strtab_len = self.strtab.write(4, &mut self.strtab_data)?; + self.strtab_offset = self.reserve(self.strtab_len as usize, 1); + Ok(()) } /// Write the string table. pub fn write_strtab(&mut self) { debug_assert_eq!(self.strtab_offset, self.buffer.len() as u32); - self.buffer - .write_bytes(&u32::to_le_bytes(self.strtab_len as u32)); + self.buffer.write_bytes(&u32::to_le_bytes(self.strtab_len)); self.buffer.write_bytes(&self.strtab_data); } } diff --git a/src/write/elf/object.rs b/src/write/elf/object.rs index a5b8b9c..3e33af9 100644 --- a/src/write/elf/object.rs +++ b/src/write/elf/object.rs @@ -685,7 +685,7 @@ impl<'a> Object<'a> { } writer.reserve_symtab_shndx(); writer.reserve_strtab_section_index(); - writer.reserve_strtab(); + writer.reserve_strtab()?; // Calculate size of relocations. for (index, section) in self.sections.iter().enumerate() { @@ -697,7 +697,7 @@ impl<'a> Object<'a> { // Calculate size of section headers. writer.reserve_shstrtab_section_index(); - writer.reserve_shstrtab(); + writer.reserve_shstrtab()?; writer.reserve_section_headers(); // Start writing. diff --git a/src/write/elf/writer.rs b/src/write/elf/writer.rs index 5dda650..0b09cd5 100644 --- a/src/write/elf/writer.rs +++ b/src/write/elf/writer.rs @@ -515,7 +515,7 @@ impl<'a> Writer<'a> { /// Write a section header. pub fn write_section_header(&mut self, section: &SectionHeader) { let sh_name = if let Some(name) = section.name { - self.shstrtab.get_offset(name) as u32 + self.shstrtab.get_offset(name) } else { 0 }; @@ -568,15 +568,18 @@ impl<'a> Writer<'a> { /// This function does nothing if no sections were reserved. /// This must be called after [`Self::add_section_name`]. /// and other functions that reserve section names and indices. - pub fn reserve_shstrtab(&mut self) { + /// + /// Returns an error if the string table could not be finalized. + pub fn reserve_shstrtab(&mut self) -> Result<()> { debug_assert_eq!(self.shstrtab_offset, 0); if self.section_num == 0 { - return; + return Ok(()); } // Start with null section name. self.shstrtab_data = vec![0]; - self.shstrtab.write(1, &mut self.shstrtab_data); + self.shstrtab.write(1, &mut self.shstrtab_data)?; self.shstrtab_offset = self.reserve(self.shstrtab_data.len(), 1); + Ok(()) } /// Write the section header string table. @@ -657,15 +660,18 @@ impl<'a> Writer<'a> { /// /// This function does nothing if a string table is not required. /// This must be called after [`Self::add_string`]. - pub fn reserve_strtab(&mut self) { + /// + /// Returns an error if the string table could not be finalized. + pub fn reserve_strtab(&mut self) -> Result<()> { debug_assert_eq!(self.strtab_offset, 0); if !self.need_strtab { - return; + return Ok(()); } // Start with null string. self.strtab_data = vec![0]; - self.strtab.write(1, &mut self.strtab_data); + self.strtab.write(1, &mut self.strtab_data)?; self.strtab_offset = self.reserve(self.strtab_data.len(), 1); + Ok(()) } /// Write the string table. @@ -818,7 +824,7 @@ impl<'a> Writer<'a> { /// Write a symbol. pub fn write_symbol(&mut self, sym: &Sym) { let st_name = if let Some(name) = sym.name { - self.strtab.get_offset(name) as u32 + self.strtab.get_offset(name) } else { 0 }; @@ -1023,16 +1029,18 @@ impl<'a> Writer<'a> { /// /// This function does nothing if no dynamic strings were defined. /// This must be called after [`Self::add_dynamic_string`]. - pub fn reserve_dynstr(&mut self) -> usize { + /// + /// Returns an error if the string table could not be finalized. + pub fn reserve_dynstr(&mut self) -> Result { debug_assert_eq!(self.dynstr_offset, 0); if !self.need_dynstr { - return 0; + return Ok(0); } // Start with null string. self.dynstr_data = vec![0]; - self.dynstr.write(1, &mut self.dynstr_data); + self.dynstr.write(1, &mut self.dynstr_data)?; self.dynstr_offset = self.reserve(self.dynstr_data.len(), 1); - self.dynstr_offset + Ok(self.dynstr_offset) } /// Return the size of the dynamic string table. @@ -1182,7 +1190,7 @@ impl<'a> Writer<'a> { /// Write a dynamic symbol. pub fn write_dynamic_symbol(&mut self, sym: &Sym) { let st_name = if let Some(name) = sym.name { - self.dynstr.get_offset(name) as u32 + self.dynstr.get_offset(name) } else { 0 }; @@ -1295,7 +1303,7 @@ impl<'a> Writer<'a> { /// Write a dynamic string entry. pub fn write_dynamic_string(&mut self, tag: elf::DynamicTag, id: StringId) -> Result<()> { - self.write_dynamic(tag, self.dynstr.get_offset(id) as u64) + self.write_dynamic(tag, self.dynstr.get_offset(id).into()) } /// Write a dynamic value entry. @@ -1699,7 +1707,7 @@ impl<'a> Writer<'a> { mem::size_of::>() as u32 }; self.buffer.write(&elf::Verdaux { - vda_name: U32::new(self.endian, self.dynstr.get_offset(name) as u32), + vda_name: U32::new(self.endian, self.dynstr.get_offset(name)), vda_next: U32::new(self.endian, vda_next), }); } @@ -1780,7 +1788,7 @@ impl<'a> Writer<'a> { self.buffer.write(&elf::Verneed { vn_version: U16::new(self.endian, verneed.version), vn_cnt: U16::new(self.endian, verneed.aux_count), - vn_file: U32::new(self.endian, self.dynstr.get_offset(verneed.file) as u32), + vn_file: U32::new(self.endian, self.dynstr.get_offset(verneed.file)), vn_aux: U32::new(self.endian, vn_aux), vn_next: U32::new(self.endian, vn_next), }); @@ -1799,7 +1807,7 @@ impl<'a> Writer<'a> { vna_hash: U32::new(self.endian, elf::hash(self.dynstr.get_string(vernaux.name))), vna_flags: U16::new(self.endian, vernaux.flags), vna_other: U16::new(self.endian, vernaux.index), - vna_name: U32::new(self.endian, self.dynstr.get_offset(vernaux.name) as u32), + vna_name: U32::new(self.endian, self.dynstr.get_offset(vernaux.name)), vna_next: U32::new(self.endian, vna_next), }); } diff --git a/src/write/macho.rs b/src/write/macho.rs index 1696627..97ac7ef 100644 --- a/src/write/macho.rs +++ b/src/write/macho.rs @@ -603,7 +603,7 @@ impl<'a> Object<'a> { let strtab_offset = offset; // Start with null name. let mut strtab_data = vec![0]; - strtab.write(1, &mut strtab_data); + strtab.write(1, &mut strtab_data)?; write_align(&mut strtab_data, pointer_align); offset += strtab_data.len(); @@ -932,7 +932,7 @@ impl<'a> Object<'a> { macho.write_nlist( buffer, Nlist { - n_strx: n_strx as u32, + n_strx, n_type, n_sect: n_sect as u8, n_desc, diff --git a/src/write/string.rs b/src/write/string.rs index cb940cd..3e04ac4 100644 --- a/src/write/string.rs +++ b/src/write/string.rs @@ -1,5 +1,7 @@ use alloc::vec::Vec; +use crate::write::{Error, Result}; + #[cfg(feature = "write_std")] type IndexSet = indexmap::IndexSet; #[cfg(not(feature = "write_std"))] @@ -9,25 +11,30 @@ type IndexSet = indexmap::IndexSet; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct StringId(usize); +/// A string table containing null terminated byte strings. #[derive(Debug, Default)] pub(crate) struct StringTable<'a> { strings: IndexSet<&'a [u8]>, - offsets: Vec, + offsets: Vec, } impl<'a> StringTable<'a> { /// Add a string to the string table. /// - /// Panics if the string table has already been written, or - /// if the string contains a null byte. + /// Duplicate strings return the id of the existing entry. + /// + /// Must be called before [`Self::write`]. + /// + /// The string must not contain a null byte; this is asserted here for + /// debug builds, and checked by `write` for all builds. pub(crate) fn add(&mut self, string: &'a [u8]) -> StringId { - assert!(self.offsets.is_empty()); - assert!(!string.contains(&0)); + debug_assert!(self.offsets.is_empty()); + debug_assert!(!string.contains(&0)); let id = self.strings.insert_full(string).0; StringId(id) } - /// Return the id of the given string. + /// Return the id of a previously added string. /// /// Panics if the string is not in the string table. #[allow(dead_code)] @@ -38,7 +45,7 @@ impl<'a> StringTable<'a> { /// Return the string for the given id. /// - /// Panics if the string is not in the string table. + /// Panics if `id` is invalid. #[allow(dead_code)] pub(crate) fn get_string(&self, id: StringId) -> &'a [u8] { self.strings.get_index(id.0).unwrap() @@ -46,9 +53,10 @@ impl<'a> StringTable<'a> { /// Return the offset of the given string. /// - /// Panics if the string table has not been written, or - /// if the string is not in the string table. - pub(crate) fn get_offset(&self, id: StringId) -> usize { + /// Must be called after [`Self::write`]. + /// + /// Panics if `id` is invalid or `write` was not called. + pub(crate) fn get_offset(&self, id: StringId) -> u32 { self.offsets[id.0] } @@ -59,28 +67,40 @@ impl<'a> StringTable<'a> { /// this should be 1 for ELF, to account for the initial /// null byte (which must have been written by the caller). /// - /// Panics if the string table has already been written. - pub(crate) fn write(&mut self, base: usize, w: &mut Vec) { - assert!(self.offsets.is_empty()); + /// Returns the total size, including base. + /// + /// Returns an error if: + /// - `write` has already been called + /// - any string contains a null byte + /// - the string table size is > `u32::MAX` + pub(crate) fn write(&mut self, base: u32, w: &mut Vec) -> Result { + if !self.offsets.is_empty() { + return Err(Error("string table already written".into())); + } + if self.strings.iter().any(|s| s.contains(&0)) { + return Err(Error("string table entry contains null byte".into())); + } let mut ids: Vec<_> = (0..self.strings.len()).collect(); sort(&mut ids, 1, &self.strings); self.offsets = vec![0; ids.len()]; - let mut offset = base; + let mut offset = u64::from(base); let mut previous = &[][..]; for id in ids { let string = self.strings.get_index(id).unwrap(); + let len = string.len() as u64 + 1; if previous.ends_with(string) { - self.offsets[id] = offset - string.len() - 1; + self.offsets[id] = (offset - len) as u32; } else { - self.offsets[id] = offset; + self.offsets[id] = offset as u32; w.extend_from_slice(string); w.push(0); - offset += string.len() + 1; + offset += len; previous = string; } } + u32::try_from(offset).map_err(|_| Error("string table size overflow".into())) } /// Calculate the size in bytes of the string table. @@ -175,12 +195,15 @@ mod tests { let mut data = Vec::new(); data.push(0); - table.write(1, &mut data); + assert_eq!(table.write(1, &mut data), Ok(12)); assert_eq!(data, b"\0foobar\0foo\0"); assert_eq!(table.get_offset(id0), 11); assert_eq!(table.get_offset(id1), 8); assert_eq!(table.get_offset(id2), 4); assert_eq!(table.get_offset(id3), 1); + + let mut data = Vec::new(); + assert!(table.write(1, &mut data).is_err()); } } diff --git a/src/write/xcoff.rs b/src/write/xcoff.rs index c2ff1c1..28509b1 100644 --- a/src/write/xcoff.rs +++ b/src/write/xcoff.rs @@ -349,9 +349,8 @@ impl<'a> Object<'a> { let strtab_offset = offset; let mut strtab_data = Vec::new(); // First 4 bytes of strtab are the length. - strtab.write(4, &mut strtab_data); - let strtab_len = strtab_data.len() + 4; - offset += strtab_len; + let strtab_len = strtab.write(4, &mut strtab_data)?; + offset += strtab_len as usize; // Start writing. buffer @@ -513,7 +512,7 @@ impl<'a> Object<'a> { }; let xcoff_sym = xcoff::Symbol64 { n_value: n_value.into(), - n_offset: (strtab.get_offset(str_id) as u32).into(), + n_offset: strtab.get_offset(str_id).into(), n_scnum: n_scnum.into(), n_type: n_type.into(), n_sclass, @@ -528,7 +527,7 @@ impl<'a> Object<'a> { sym_name[..symbol.name.len()].copy_from_slice(&symbol.name[..]); } else { let str_offset = strtab.get_offset(symbol_offsets[index].str_id.unwrap()); - sym_name[4..8].copy_from_slice(&u32::to_be_bytes(str_offset as u32)); + sym_name[4..8].copy_from_slice(&u32::to_be_bytes(str_offset)); } let xcoff_sym = xcoff::Symbol32 { n_name: sym_name, @@ -548,7 +547,7 @@ impl<'a> Object<'a> { x_fname[..symbol.name.len()].copy_from_slice(&symbol.name[..]); } else { let str_offset = strtab.get_offset(symbol_offsets[index].str_id.unwrap()); - x_fname[4..8].copy_from_slice(&u32::to_be_bytes(str_offset as u32)); + x_fname[4..8].copy_from_slice(&u32::to_be_bytes(str_offset)); } if is_64 { let file_aux = xcoff::FileAux64 { @@ -610,7 +609,7 @@ impl<'a> Object<'a> { // Write string table. debug_assert_eq!(strtab_offset, buffer.len()); - buffer.write_bytes(&u32::to_be_bytes(strtab_len as u32)); + buffer.write_bytes(&u32::to_be_bytes(strtab_len)); buffer.write_bytes(&strtab_data); debug_assert_eq!(offset, buffer.len());