use std::slice::Iter;
use crate::{BufferIterator, FromBytes, ToBytes, WritableBuffer};
use self::sections::{DataSection, FuncSection, SectionHeader, StringTable, SymbolTable};
use self::sections::{ReldSection, SectionKind};
pub mod errors;
pub mod instructions;
pub mod sections;
pub mod symbols;
use crate::ko::errors::{HeaderParseError, KOParseError, ValidationError};
use crate::ko::sections::StringIdx;
pub use instructions::Instr;
const FILE_VERSION: u8 = 4;
const MAGIC_NUMBER: u32 = 0x666f016b;
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct SectionIdx(u16);
impl From<u8> for SectionIdx {
fn from(i: u8) -> Self {
Self(i as u16)
}
}
impl From<u16> for SectionIdx {
fn from(i: u16) -> Self {
Self(i)
}
}
impl From<SectionIdx> for u16 {
fn from(idx: SectionIdx) -> Self {
idx.0
}
}
impl From<SectionIdx> for usize {
fn from(idx: SectionIdx) -> Self {
idx.0 as usize
}
}
impl SectionIdx {
pub const NULL: SectionIdx = SectionIdx(0u16);
}
#[derive(Debug)]
pub struct KOFile {
header: KOHeader,
shstrtab: StringTable,
section_headers: Vec<SectionHeader>,
str_tabs: Vec<StringTable>,
sym_tabs: Vec<SymbolTable>,
data_sections: Vec<DataSection>,
func_sections: Vec<FuncSection>,
reld_sections: Vec<ReldSection>,
}
impl KOFile {
pub fn new() -> Self {
let shstrtab_idx = SectionIdx::from(1u16);
let mut kofile = Self {
header: KOHeader::new(2, shstrtab_idx),
shstrtab: StringTable::with_capacity(256, shstrtab_idx),
section_headers: Vec::with_capacity(4),
str_tabs: Vec::with_capacity(2),
sym_tabs: Vec::with_capacity(1),
data_sections: Vec::with_capacity(1),
func_sections: Vec::with_capacity(1),
reld_sections: Vec::with_capacity(1),
};
kofile.add_section_header(SectionHeader::null());
let name_idx = kofile.shstrtab.add(".shstrtab");
kofile.add_section_header(SectionHeader::new(name_idx, SectionKind::StrTab));
kofile
}
pub fn add_section_name(&mut self, name: impl Into<String>) -> StringIdx {
self.shstrtab.add(name)
}
pub fn get_section_name(&self, index: StringIdx) -> Option<&String> {
self.shstrtab.get(index)
}
pub fn get_section_name_by_index(&self, index: SectionIdx) -> Option<&String> {
let header = self.section_headers.get(usize::from(index))?;
self.shstrtab.get(header.name_idx)
}
pub fn add_section_header(&mut self, header: SectionHeader) -> SectionIdx {
let index = self.section_headers.len();
if index > u16::MAX as usize {
panic!(
"Attempted to insert {} section headers. Section indices have a maximum of {}",
u16::MAX as usize + 1,
u16::MAX
);
}
self.section_headers.push(header);
SectionIdx::from(index as u16)
}
pub fn get_section_header(&self, index: SectionIdx) -> Option<&SectionHeader> {
self.section_headers.get(usize::from(index))
}
pub fn get_header_name(&self, header: &SectionHeader) -> Option<&String> {
self.shstrtab.get(header.name_idx)
}
pub fn add_str_tab(&mut self, str_tab: StringTable) {
self.str_tabs.push(str_tab);
}
pub fn add_sym_tab(&mut self, sym_tab: SymbolTable) {
self.sym_tabs.push(sym_tab);
}
pub fn add_data_section(&mut self, data_section: DataSection) {
self.data_sections.push(data_section);
}
pub fn add_func_section(&mut self, func_section: FuncSection) {
self.func_sections.push(func_section);
}
pub fn add_reld_section(&mut self, reld_section: ReldSection) {
self.reld_sections.push(reld_section);
}
pub fn str_tabs(&self) -> Iter<StringTable> {
self.str_tabs.iter()
}
pub fn sym_tabs(&self) -> Iter<SymbolTable> {
self.sym_tabs.iter()
}
pub fn data_sections(&self) -> Iter<DataSection> {
self.data_sections.iter()
}
pub fn func_sections(&self) -> Iter<FuncSection> {
self.func_sections.iter()
}
pub fn reld_sections(&self) -> Iter<ReldSection> {
self.reld_sections.iter()
}
pub fn new_section_header(&mut self, name: impl Into<String>, kind: SectionKind) -> SectionIdx {
let name_idx = self.add_section_name(name);
let header = SectionHeader::new(name_idx, kind);
self.add_section_header(header)
}
pub fn get_section_index_by_name(&self, name: impl AsRef<str>) -> Option<SectionIdx> {
for (index, header) in self.section_headers.iter().enumerate() {
let name_idx = header.name_idx;
let sh_name = match self.shstrtab.get(name_idx) {
Some(name) => name,
None => {
continue;
}
};
if name.as_ref() == sh_name {
return Some(SectionIdx::from(index as u16));
}
}
None
}
pub fn header(&self) -> KOHeader {
self.header
}
pub fn section_headers(&self) -> Iter<SectionHeader> {
self.section_headers.iter()
}
pub fn section_header_count(&self) -> usize {
self.section_headers.len()
}
pub fn shstrtab_index(&self) -> SectionIdx {
self.shstrtab.section_index()
}
fn update_section_header(
&mut self,
section_kind: SectionKind,
section_index: SectionIdx,
section_size: u32,
) -> Result<(), ValidationError> {
match self
.section_headers
.get_mut(u16::from(section_index) as usize)
{
Some(header) => {
if header.section_kind != section_kind {
Err(ValidationError::SectionKindMismatchError(
section_kind,
u16::from(section_index),
header.section_kind,
))
} else {
header.size = section_size;
Ok(())
}
}
None => Err(ValidationError::InvalidSectionIndexError(
section_kind,
u16::from(section_index),
)),
}
}
fn update_section_headers(&mut self) -> Result<(), ValidationError> {
let mut header_set = vec![self.shstrtab.section_index()];
self.section_headers
.get_mut(usize::from(self.shstrtab.section_index()))
.unwrap()
.size = self.shstrtab.size();
for i in 0..self.str_tabs.len() {
let section = self.str_tabs.get(i).unwrap();
let idx = section.section_index();
let size = section.size();
header_set.push(idx);
self.update_section_header(SectionKind::StrTab, idx, size)?;
}
for i in 0..self.sym_tabs.len() {
let section = self.sym_tabs.get(i).unwrap();
let idx = section.section_index();
let size = section.size();
header_set.push(idx);
self.update_section_header(SectionKind::SymTab, idx, size)?;
}
for i in 0..self.data_sections.len() {
let section = self.data_sections.get(i).unwrap();
let idx = section.section_index();
let size = section.size();
header_set.push(idx);
self.update_section_header(SectionKind::Data, idx, size)?;
}
for i in 0..self.func_sections.len() {
let section = self.func_sections.get(i).unwrap();
let idx = section.section_index();
let size = section.size();
header_set.push(idx);
self.update_section_header(SectionKind::Func, idx, size)?;
}
for i in 0..self.reld_sections.len() {
let section = self.reld_sections.get(i).unwrap();
let idx = section.section_index();
let size = section.size();
header_set.push(idx);
self.update_section_header(SectionKind::Reld, idx, size)?;
}
for (i, header) in self
.section_headers
.iter()
.enumerate()
.skip(1)
.map(|(i, h)| (SectionIdx::from(i as u16), h))
{
let name = self
.shstrtab
.get(header.name_idx)
.ok_or_else(|| {
ValidationError::InvalidSectionHeaderNameIndexError(
u16::from(i),
header.section_kind,
usize::from(header.name_idx),
)
})?
.clone();
if !header_set.contains(&i) {
return Err(ValidationError::InvalidSectionHeaderError(
u16::from(i),
header.section_kind,
name,
));
}
}
Ok(())
}
pub fn validate(mut self) -> Result<WritableKOFile, (Self, ValidationError)> {
self.header = KOHeader::new(
self.section_headers.len() as u16,
self.shstrtab.section_index(),
);
if let Err(e) = self.update_section_headers() {
Err((self, e))
} else {
Ok(WritableKOFile(self))
}
}
pub fn parse(source: &mut BufferIterator) -> Result<Self, KOParseError> {
let header = KOHeader::parse(source).map_err(KOParseError::HeaderError)?;
let mut section_headers = Vec::with_capacity(header.num_headers as usize);
let mut str_tabs;
let mut sym_tabs;
let mut data_sections;
let mut func_sections;
let mut reld_sections;
let mut num_str_tabs = 0;
let mut num_sym_tabs = 0;
let mut num_data_sections = 0;
let mut num_func_sections = 0;
let mut num_reld_sections = 0;
let null_header =
SectionHeader::parse(source).map_err(KOParseError::SectionHeaderParseError)?;
if null_header.section_kind != SectionKind::Null {
return Err(KOParseError::MissingNullSectionHeader(
null_header.section_kind,
));
}
section_headers.push(null_header);
for i in 1..header.num_headers {
let header =
SectionHeader::parse(source).map_err(KOParseError::SectionHeaderParseError)?;
match header.section_kind {
SectionKind::StrTab => {
num_str_tabs += 1;
}
SectionKind::SymTab => {
num_sym_tabs += 1;
}
SectionKind::Data => {
num_data_sections += 1;
}
SectionKind::Func => {
num_func_sections += 1;
}
SectionKind::Reld => {
num_reld_sections += 1;
}
SectionKind::Debug => {
return Err(KOParseError::DebugSectionUnsupportedError(i));
}
SectionKind::Null => {
return Err(KOParseError::StrayNullSectionHeader(i));
}
}
section_headers.push(header);
}
if header.shstrtab_idx == SectionIdx::from(0u16) {
return Err(KOParseError::NullShStrTabIndexError);
}
str_tabs = Vec::with_capacity(num_str_tabs);
sym_tabs = Vec::with_capacity(num_sym_tabs);
data_sections = Vec::with_capacity(num_data_sections);
func_sections = Vec::with_capacity(num_func_sections);
reld_sections = Vec::with_capacity(num_reld_sections);
let shstrtab_size = section_headers
.get(usize::from(header.shstrtab_idx))
.ok_or_else(|| {
KOParseError::InvalidShStrTabIndexError(
header.shstrtab_idx.into(),
section_headers.len() as u16,
)
})?
.size;
let shstrtab = StringTable::parse(source, shstrtab_size, header.shstrtab_idx)
.map_err(KOParseError::StringTableParseError)?;
for section_idx in (1..section_headers.len() as u16).map(SectionIdx::from) {
if section_idx == header.shstrtab_idx {
continue;
}
let header = section_headers.get(usize::from(section_idx)).unwrap();
match header.section_kind {
SectionKind::StrTab => {
str_tabs.push(
StringTable::parse(source, header.size, section_idx)
.map_err(KOParseError::StringTableParseError)?,
);
}
SectionKind::SymTab => {
sym_tabs.push(
SymbolTable::parse(source, header.size, section_idx)
.map_err(KOParseError::SymbolTableParseError)?,
);
}
SectionKind::Data => {
data_sections.push(
DataSection::parse(source, header.size, section_idx)
.map_err(KOParseError::DataSectionParseError)?,
);
}
SectionKind::Func => {
func_sections.push(
FuncSection::parse(source, header.size, section_idx)
.map_err(KOParseError::FunctionSectionParseError)?,
);
}
SectionKind::Reld => {
reld_sections.push(
ReldSection::parse(source, header.size, section_idx)
.map_err(KOParseError::ReldSectionParseError)?,
);
}
SectionKind::Debug => unimplemented!(),
SectionKind::Null => {
panic!("Internal library error. Attempted to parse \"null\" section, this should be unreachable.");
}
}
}
Ok(Self {
header,
shstrtab,
section_headers,
str_tabs,
sym_tabs,
data_sections,
func_sections,
reld_sections,
})
}
}
impl Default for KOFile {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct WritableKOFile(KOFile);
impl WritableKOFile {
pub fn write(&self, buf: &mut impl WritableBuffer) {
let ko = &self.0;
ko.header.write(buf);
for header in &ko.section_headers {
header.write(buf);
}
for section_index in (0..ko.section_header_count()).map(|i| SectionIdx::from(i as u16)) {
if self.0.shstrtab.section_index() == section_index {
self.0.shstrtab.write(buf);
continue;
}
if let Some(section) = ko
.str_tabs
.iter()
.find(|s| s.section_index() == section_index)
{
section.write(buf);
continue;
}
if let Some(section) = ko
.sym_tabs
.iter()
.find(|s| s.section_index() == section_index)
{
section.write(buf);
continue;
}
if let Some(section) = ko
.data_sections
.iter()
.find(|s| s.section_index() == section_index)
{
section.write(buf);
continue;
}
if let Some(section) = ko
.func_sections
.iter()
.find(|s| s.section_index() == section_index)
{
section.write(buf);
continue;
}
if let Some(section) = ko
.reld_sections
.iter()
.find(|s| s.section_index() == section_index)
{
section.write(buf);
continue;
}
}
}
pub fn get(self) -> KOFile {
self.0
}
}
impl From<WritableKOFile> for KOFile {
fn from(w_kofile: WritableKOFile) -> Self {
w_kofile.0
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct KOHeader {
pub magic: u32,
pub version: u8,
pub num_headers: u16,
pub shstrtab_idx: SectionIdx,
}
impl KOHeader {
const HEADER_SIZE: usize = 9;
pub const fn new(num_headers: u16, shstrtab_idx: SectionIdx) -> Self {
Self {
magic: MAGIC_NUMBER,
version: FILE_VERSION,
num_headers,
shstrtab_idx,
}
}
pub const fn size() -> usize {
Self::HEADER_SIZE
}
fn write(&self, buf: &mut impl WritableBuffer) {
self.magic.to_bytes(buf);
self.version.to_bytes(buf);
self.num_headers.to_bytes(buf);
u16::from(self.shstrtab_idx).to_bytes(buf);
}
pub fn parse(source: &mut BufferIterator) -> Result<Self, HeaderParseError> {
let magic = u32::from_bytes(source).map_err(|_| HeaderParseError::MissingMagicError)?;
let version = u8::from_bytes(source).map_err(|_| HeaderParseError::MissingVersionError)?;
let num_headers =
u16::from_bytes(source).map_err(|_| HeaderParseError::MissingNumHeaders)?;
let shstrtab_idx = SectionIdx(
u16::from_bytes(source).map_err(|_| HeaderParseError::MissingShStrTabIndex)?,
);
if magic != MAGIC_NUMBER {
return Err(HeaderParseError::InvalidMagicError(magic, MAGIC_NUMBER));
}
if version != FILE_VERSION {
return Err(HeaderParseError::UnsupportedVersionError(
version,
FILE_VERSION,
));
}
Ok(Self {
magic,
version,
num_headers,
shstrtab_idx,
})
}
}
macro_rules! gen_get_by_name {
($(#[$ref_attr:meta])* => $func_name: ident, $(#[$mut_attr:meta])* => $mut_func_name: ident, $section_name: ident, $section_type: ty) => {
$(#[$ref_attr])*
pub fn $func_name(&self, name: impl AsRef<str>) -> Option<&$section_type> {
self.$section_name.iter().find(|section| {
match self.get_section_name_by_index(section.section_index()) {
Some(s) => s == name.as_ref(),
None => false,
}
})
}
$(#[$mut_attr])*
pub fn $mut_func_name(&mut self, name: impl AsRef<str>) -> Option<&mut $section_type> {
let mut index_opt = Some(0);
for (index, section) in self.$section_name.iter().enumerate() {
match self.get_section_name_by_index(section.section_index()) {
Some(s) => {
if s == name.as_ref() {
index_opt = Some(index);
break;
}
}
None => continue,
}
}
let index = index_opt?;
self.$section_name.iter_mut().nth(index)
}
};
}
macro_rules! gen_new_section {
($(#[$attr:meta])* => $func_name: ident, $section_type: ty, $section_kind: expr) => {
$(#[$attr])*
pub fn $func_name(&mut self, name: impl Into<String>) -> $section_type {
let sh_index = self.new_section_header(name, $section_kind);
<$section_type>::with_capacity(4, sh_index)
}
};
}
impl KOFile {
gen_new_section! {
=>
new_strtab,
StringTable,
SectionKind::StrTab
}
gen_new_section! {
=>
new_symtab,
SymbolTable,
SectionKind::SymTab
}
gen_new_section! {
=>
new_data_section,
DataSection,
SectionKind::Data
}
gen_new_section! {
=>
new_func_section,
FuncSection,
SectionKind::Func
}
gen_new_section! {
=>
new_reld_section,
ReldSection,
SectionKind::Reld
}
gen_get_by_name! {
=>
str_tab_by_name,
=>
str_tab_by_name_mut, str_tabs, StringTable}
gen_get_by_name! {
=>
sym_tab_by_name,
=>
sym_tab_by_name_mut,
sym_tabs,
SymbolTable
}
gen_get_by_name! {
=>
data_section_by_name,
=>
data_section_by_name_mut,
data_sections,
DataSection
}
gen_get_by_name! {
=>
func_section_by_name,
=>
func_section_by_name_mut,
func_sections,
FuncSection
}
gen_get_by_name!(
=>
reld_section_by_name,
=>
reld_section_by_name_mut,
reld_sections,
ReldSection
);
}