Skip to content

Commit

Permalink
Merge pull request #90 from Etto48:feature/custom_header
Browse files Browse the repository at this point in the history
Feature/custom_header
  • Loading branch information
Etto48 authored Jul 4, 2024
2 parents d4578a9 + 1411c3e commit f34b030
Show file tree
Hide file tree
Showing 20 changed files with 536 additions and 49 deletions.
31 changes: 31 additions & 0 deletions docs/PLUGIN_API.md
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,21 @@ The popup must be opened using `context.open_popup("FILL_POPUP_NAME")`.
|`popup_context`|`PopupContext`|The popup context.|
|`context`|`Context`|The application context.|

### Custom Headers

```lua
function HEADER_PARSER_NAME(header_context, context) end
```

This function is called whenever a new header is parsed, the first parser that returns a valid header will be used.
If during the parsing you decide that the header is not valid, you can raise an error or avoid setting values in the `header_context`.
The parser must be registered using `context.add_header_parser("HEADER_PARSER_NAME")`.

| Argument | Type | Description |
|----------|------|-------------|
|`header_context`|`HeaderContext`|The header context.|
|`context`|`Context`|The application context.|

## Types

### Context
Expand All @@ -183,6 +198,8 @@ And the following functions:
|`log`|`(level: u8, message: String)`|Logs a message in the UI.|
|`add_command`|`(command_name: String)`|Registers a command, this must be called to make the command appear in the command list.|
|`remove_command`|`(command_name: String)`|Removes a command, this removes the command from the command list.|
|`add_header_parser`|`(parser_name: String)`|Registers a header parser, this must be called to make the parser be used when a new file is opened.|
|`remove_header_parser`|`(parser_name: String)`|Removes a header parser, this removes the parser from the list of parsers.|
|`open_popup`|`(popup_handler: String)`|Opens a popup, each time the popup is drawn the handler function is called|
|`get_popup`|`() -> Option<String>`|Returns the name of the `popup_handler` of the currently open popup if there is one opened by this plugin. `nil` otherwise.|
|`close_popup`|`(popup_handler: Option<String>)`|Closes a popup opened by this plugin. If `popup_handler` is not `nil` it will also check if that is the currently open popup. If no popup is open, this plugin does not own the currently open popup, or the provided handler does not match the function will raise an error.|
Expand Down Expand Up @@ -434,6 +451,20 @@ This type has the following fields:
|`file_offset`|`u64`|The starting file offset of the section.|
|`size`|`usize`|The size of the section.|

### HeaderContext

This table contains the following fields:
| Function | Arguments | Description | Required |
|----------|-----------|-------------|----------|
|`set_bitness`|`(bitness: u8)`|Sets the bitness of the file. `bitness` must be either `32` or `64`| Yes |
|`set_entry`|`(entry_point: u64)`|Sets the virtual address of the entry point.| Yes |
|`set_endianness`|`(endianness: String)`|Sets the endianness of the file. `endianness` mut be either `little` or `big`.| Yes |
|`set_architecture`|`(architecture: String)`|Sets the architecture of the file. The possible values are listed in [Header.architecture](#headerarchitecture).| Yes |
|`add_section`|`(name: String, virtual_address: u64, file_offset: u64, size: u64)`|Adds a section to the file.| No |
|`add_symbol`|`(address: u64, name: String)`|Adds a symbol to the file.| No |

A function marked as required must be called to create a valid header. Those functions can only be called once.

### Style

This table contains the following fields:
Expand Down
4 changes: 3 additions & 1 deletion docs/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ The first key found will be used.

## Supported file formats and architectures

The following file formats are supported:
The following file formats are supported by default:

- Coff
- CoffBig
Expand All @@ -50,6 +50,8 @@ The following file formats are supported:
- Xcoff32
- Xcoff64

Other file formats can be added with [plugins](#plugins).

The following architectures are supported:

- Aarch64
Expand Down
10 changes: 9 additions & 1 deletion src/app/asm/assembly.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::{
},
asm::assembler::assemble,
get_app_context,
headers::{Header, Section},
headers::{section::Section, Header},
};

use super::{
Expand Down Expand Up @@ -420,6 +420,14 @@ impl App {
.splice(from_instruction..to_instruction, instructions);
}
}

pub(in crate::app) fn parse_header(&mut self) -> Header {
let mut app_context = get_app_context!(self);
match self.plugin_manager.try_parse_header(&mut app_context) {
Some(header) => Header::CustomHeader(header),
None => Header::parse_header(&self.data.bytes, self.filesystem.pwd(), &self.filesystem),
}
}
}

#[cfg(test)]
Expand Down
4 changes: 3 additions & 1 deletion src/app/files/files.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,8 @@ impl App {
NotificationLevel::Info,
&format!("File type: {:?}", header.file_type()),
),
// TODO: maybe add info for a more detailed log
Header::CustomHeader(_) => self.log(NotificationLevel::Info, "File type: Custom"),
Header::None => unreachable!(),
}
self.log(
Expand Down Expand Up @@ -195,7 +197,7 @@ impl App {
} else {
None
};
self.header = Header::parse_header(&self.data.bytes, path, &self.filesystem);
self.header = self.parse_header();

terminal = if let Some(terminal) = terminal {
Self::print_loading_status(
Expand Down
67 changes: 66 additions & 1 deletion src/app/plugins/app_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,10 @@ use crate::{
headers::Header,
};

use super::{exported_commands::ExportedCommands, instruction_info::InstructionInfo};
use super::{
exported_commands::ExportedCommands, exported_header_parsers::ExportedHeaderParsers,
instruction_info::InstructionInfo,
};

#[macro_export]
macro_rules! get_app_context {
Expand All @@ -33,6 +36,7 @@ macro_rules! get_app_context {

pub struct AppContext<'app> {
pub exported_commands: Arc<Mutex<ExportedCommands>>,
pub exported_header_parsers: Arc<Mutex<ExportedHeaderParsers>>,
pub plugin_index: Option<usize>,

pub screen_height: u16,
Expand Down Expand Up @@ -61,6 +65,7 @@ impl<'app> AppContext<'app> {
) -> Self {
Self {
exported_commands: Arc::new(Mutex::new(ExportedCommands::default())),
exported_header_parsers: Arc::new(Mutex::new(ExportedHeaderParsers::default())),
plugin_index: None,
screen_height,
screen_width,
Expand All @@ -78,14 +83,26 @@ impl<'app> AppContext<'app> {
self.exported_commands = Arc::new(Mutex::new(ExportedCommands::default()));
}

pub fn reset_exported_header_parsers(&mut self) {
self.exported_header_parsers = Arc::new(Mutex::new(ExportedHeaderParsers::default()));
}

pub fn set_exported_commands(&mut self, exported_commands: ExportedCommands) {
self.exported_commands = Arc::new(Mutex::new(exported_commands));
}

pub fn set_exported_header_parsers(&mut self, exported_header_parsers: ExportedHeaderParsers) {
self.exported_header_parsers = Arc::new(Mutex::new(exported_header_parsers));
}

pub fn take_exported_commands(&mut self) -> ExportedCommands {
self.exported_commands.lock().unwrap().take()
}

pub fn take_exported_header_parsers(&mut self) -> ExportedHeaderParsers {
self.exported_header_parsers.lock().unwrap().take()
}

pub fn to_lua<'lua>(
&'lua mut self,
lua: &'lua Lua,
Expand Down Expand Up @@ -146,6 +163,54 @@ impl<'app> AppContext<'app> {
)
.unwrap();

let exported_header_parsers = self.exported_header_parsers.clone();
context
.set(
"add_header_parser",
scope
.create_function_mut(move |lua, callback: String| {
if let Ok(_header_parser_fn) =
lua.globals().get::<_, Function>(callback.clone())
{
exported_header_parsers
.lock()
.unwrap()
.add_header_parser(callback);
Ok(())
} else {
Err(mlua::Error::external(format!(
"Function '{}' not found but needed to export the header parser",
callback
)))
}
})
.unwrap(),
)
.unwrap();

let exported_header_parsers = self.exported_header_parsers.clone();
context
.set(
"remove_header_parser",
scope
.create_function_mut(move |_, callback: String| {
if exported_header_parsers
.lock()
.unwrap()
.remove_header_parser(&callback)
{
Ok(())
} else {
Err(mlua::Error::external(format!(
"Header parser '{}' not found",
callback
)))
}
})
.unwrap(),
)
.unwrap();

context
.set(
"open_popup",
Expand Down
25 changes: 25 additions & 0 deletions src/app/plugins/exported_header_parsers.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
use super::header_parser_info::HeaderParserInfo;

#[derive(Debug, Clone, Default)]
pub struct ExportedHeaderParsers {
pub parsers: Vec<HeaderParserInfo>,
}

impl ExportedHeaderParsers {
pub fn add_header_parser(&mut self, parser: String) {
self.parsers.push(HeaderParserInfo { parser });
}

pub fn remove_header_parser(&mut self, parser: &str) -> bool {
if let Some(index) = self.parsers.iter().position(|c| c.parser == parser) {
self.parsers.remove(index);
true
} else {
false
}
}

pub fn take(&mut self) -> Self {
std::mem::take(self)
}
}
144 changes: 144 additions & 0 deletions src/app/plugins/header_context.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
use std::collections::HashMap;

use mlua::UserData;
use object::{Architecture, Endianness};

use crate::headers::{bitness::Bitness, custom_header::CustomHeader, section::Section};

#[derive(Debug, Clone, Default)]
pub struct HeaderContext {
pub bitness: Option<Bitness>,
pub endianness: Option<Endianness>,
pub entry: Option<u64>,
pub architecture: Option<Architecture>,
pub sections: Vec<Section>,
pub symbols: HashMap<u64, String>,
}

impl HeaderContext {
pub fn try_into_custom_header(self) -> Option<CustomHeader> {
if let (Some(bitness), Some(endianness), Some(entry), Some(architecture)) =
(self.bitness, self.endianness, self.entry, self.architecture)
{
let symbols_by_name = self.symbols.iter().map(|(k, v)| (v.clone(), *k)).collect();
Some(CustomHeader {
bitness,
entry,
endianness,
architecture,
sections: self.sections,
symbols: self.symbols,
symbols_by_name,
})
} else {
None
}
}
}

impl UserData for HeaderContext {
fn add_methods<'lua, M: mlua::UserDataMethods<'lua, Self>>(methods: &mut M) {
methods.add_method_mut("set_bitness", |_, this, bitness: u8| {
if this.bitness.is_some() {
Err(mlua::Error::external("bitness already set"))
} else {
match bitness {
32 => {
this.bitness = Some(Bitness::Bit32);
Ok(())
}
64 => {
this.bitness = Some(Bitness::Bit64);
Ok(())
}
_ => Err(mlua::Error::external("invalid bitness")),
}
}
});

methods.add_method_mut("set_entry", |_, this, entry_point: u64| {
if this.entry.is_some() {
Err(mlua::Error::external("entry point already set"))
} else {
this.entry = Some(entry_point);
Ok(())
}
});

methods.add_method_mut("set_endianness", |_, this, endianness: String| {
if this.endianness.is_some() {
Err(mlua::Error::external("endianness already set"))
} else {
match endianness.as_str() {
"little" => {
this.endianness = Some(Endianness::Little);
Ok(())
}
"big" => {
this.endianness = Some(Endianness::Big);
Ok(())
}
_ => Err(mlua::Error::external("invalid endianness")),
}
}
});

methods.add_method_mut("set_architecture", |_, this, architecture: String| {
if this.architecture.is_some() {
Err(mlua::Error::external("architecture already set"))
} else {
let architecture = match architecture.as_str() {
"Unknown" => Architecture::Unknown,
"Aarch64" => Architecture::Aarch64,
"Aarch64_Ilp32" => Architecture::Aarch64_Ilp32,
"Arm" => Architecture::Arm,
"Avr" => Architecture::Avr,
"Bpf" => Architecture::Bpf,
"Csky" => Architecture::Csky,
"I386" => Architecture::I386,
"X86_64" => Architecture::X86_64,
"X86_64_X32" => Architecture::X86_64_X32,
"Hexagon" => Architecture::Hexagon,
"LoongArch64" => Architecture::LoongArch64,
"Mips" => Architecture::Mips,
"Mips64" => Architecture::Mips64,
"Msp430" => Architecture::Msp430,
"PowerPc" => Architecture::PowerPc,
"PowerPc64" => Architecture::PowerPc64,
"Riscv32" => Architecture::Riscv32,
"Riscv64" => Architecture::Riscv64,
"S390x" => Architecture::S390x,
"Sbf" => Architecture::Sbf,
"Sharc" => Architecture::Sharc,
"Sparc" => Architecture::Sparc,
"Sparc32Plus" => Architecture::Sparc32Plus,
"Sparc64" => Architecture::Sparc64,
"Wasm32" => Architecture::Wasm32,
"Wasm64" => Architecture::Wasm64,
"Xtensa" => Architecture::Xtensa,
_ => return Err(mlua::Error::external("invalid architecture")),
};
this.architecture = Some(architecture);
Ok(())
}
});

methods.add_method_mut(
"add_section",
|_, this, (name, virtual_address, file_offset, size): (String, u64, u64, u64)| {
this.sections.push(Section {
name,
virtual_address,
file_offset,
size,
});
Ok(())
},
);

methods.add_method_mut("add_symbol", |_, this, (address, name): (u64, String)| {
this.symbols.insert(address, name);
Ok(())
});
}
}
4 changes: 4 additions & 0 deletions src/app/plugins/header_parser_info.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
#[derive(Debug, Clone)]
pub struct HeaderParserInfo {
pub parser: String,
}
Loading

0 comments on commit f34b030

Please sign in to comment.