object_rewrite/
rewriter.rs

1use std::collections::{HashMap, HashSet};
2
3use object::build;
4
5use super::{Error, Result};
6
7/// Options for modifying a file.
8///
9/// This is used as an argument to the [`Rewriter::modify`] method.
10///
11/// The options are listed in the order they are processed.
12#[derive(Debug, Default)]
13#[non_exhaustive]
14pub struct Options {
15    /// Delete symbols from the symbol table.
16    ///
17    /// See [`Rewriter::delete_symbols`].
18    pub delete_symbols: HashSet<Vec<u8>>,
19    /// Rename symbols in the symbol table.
20    ///
21    /// See [`Rewriter::rename_symbols`].
22    pub rename_symbols: HashMap<Vec<u8>, Vec<u8>>,
23    /// Delete sections from the file.
24    ///
25    /// See [`Rewriter::delete_sections`].
26    pub delete_sections: HashSet<Vec<u8>>,
27    /// Rename sections in the file.
28    ///
29    /// See [`Rewriter::rename_sections`].
30    pub rename_sections: HashMap<Vec<u8>, Vec<u8>>,
31    /// Options that are specific to ELF files.
32    pub elf: super::ElfOptions,
33}
34
35/// A rewriter for object and executable files.
36///
37/// This struct provides a way to read a file, modify it, and write it back.
38#[derive(Debug)]
39pub struct Rewriter<'data> {
40    pub(crate) builder: build::elf::Builder<'data>,
41    pub(crate) modified: bool,
42}
43
44impl<'data> Rewriter<'data> {
45    /// Read a file and create a new rewriter.
46    pub fn read(data: &'data [u8]) -> Result<Self> {
47        let builder = build::elf::Builder::read(data).map_err(Error::parse)?;
48        Ok(Self {
49            builder,
50            modified: false,
51        })
52    }
53
54    /// Write the file to an output stream.
55    pub fn write<W: std::io::Write>(mut self, w: W) -> Result<()> {
56        self.elf_finalize()?;
57        let mut buffer = object::write::StreamingBuffer::new(w);
58        self.builder.write(&mut buffer).map_err(Error::write)?;
59        buffer.result().map_err(Error::io)
60    }
61
62    /// Modify the file according to the given options.
63    pub fn modify(&mut self, options: Options) -> Result<()> {
64        if !options.delete_symbols.is_empty() {
65            self.delete_symbols(&options.delete_symbols);
66        }
67        if !options.rename_symbols.is_empty() {
68            self.rename_symbols(&options.rename_symbols);
69        }
70        if !options.delete_sections.is_empty() {
71            self.delete_sections(&options.delete_sections);
72        }
73        if !options.rename_sections.is_empty() {
74            self.rename_sections(&options.rename_sections);
75        }
76        self.elf_modify(options.elf)?;
77        Ok(())
78    }
79
80    /// Delete symbols from the symbol table.
81    ///
82    /// For ELF files, this deletes symbols from both the symbol table and the
83    /// dynamic symbol table.
84    pub fn delete_symbols(&mut self, names: &HashSet<Vec<u8>>) {
85        self.elf_delete_symbols(names);
86        self.elf_delete_dynamic_symbols(names);
87    }
88
89    /// Rename symbols in the symbol table.
90    ///
91    /// For ELF files, this renames symbols in both the symbol table and the
92    /// dynamic symbol table.
93    ///
94    /// The `names` map is from old names to new names.
95    pub fn rename_symbols(&mut self, names: &HashMap<Vec<u8>, Vec<u8>>) {
96        self.elf_rename_symbols(names);
97        self.elf_rename_dynamic_symbols(names);
98    }
99
100    /// Delete sections from the file.
101    pub fn delete_sections(&mut self, names: &HashSet<Vec<u8>>) {
102        self.elf_delete_sections(names);
103    }
104
105    /// Rename sections in the file.
106    ///
107    /// The `names` map is from old names to new names.
108    pub fn rename_sections(&mut self, names: &HashMap<Vec<u8>, Vec<u8>>) {
109        self.elf_rename_sections(names);
110    }
111}