o1vm/
elf_loader.rs

1use crate::cannon::{Page, State, PAGE_SIZE};
2use elf::{
3    endian::{BigEndian, EndianParse, LittleEndian},
4    section::SectionHeader,
5    ElfBytes,
6};
7use log::debug;
8use std::{collections::HashMap, path::Path};
9
10pub enum Architecture {
11    Mips,
12    RiscV32,
13}
14
15pub fn make_state<T: EndianParse>(file: ElfBytes<T>) -> Result<State, String> {
16    // Checking it is RISC-V
17
18    let (shdrs_opt, strtab_opt) = file
19        .section_headers_with_strtab()
20        .expect("shdrs offsets should be valid");
21    let (shdrs, strtab) = (
22        shdrs_opt.expect("Should have shdrs"),
23        strtab_opt.expect("Should have strtab"),
24    );
25
26    // Parse the shdrs and collect them into a map keyed on their zero-copied name
27    let sections_by_name: HashMap<&str, SectionHeader> = shdrs
28        .iter()
29        .map(|shdr| {
30            (
31                strtab
32                    .get(shdr.sh_name as usize)
33                    .expect("Failed to get section name"),
34                shdr,
35            )
36        })
37        .collect();
38
39    debug!("Loading the text section, which contains the executable code.");
40    // Getting the executable code.
41    let text_section = sections_by_name
42        .get(".text")
43        .expect("Should have .text section");
44
45    let (text_section_data, _) = file
46        .section_data(text_section)
47        .expect("Failed to read data from .text section");
48
49    // address of starting instruction in code section
50    let code_section_starting_address = text_section.sh_addr as usize;
51    let code_section_size = text_section.sh_size as usize;
52    // address of last instruction in code section
53    let code_section_end_address = code_section_starting_address + code_section_size - 1;
54    debug!(
55        "The executable code starts at address {}, has size {} bytes, and ends at address {}.",
56        code_section_starting_address, code_section_size, code_section_end_address
57    );
58
59    // Building the memory pages
60    let mut memory: Vec<Page> = vec![];
61    let page_size_usize: usize = PAGE_SIZE.try_into().unwrap();
62    // Padding to get the right number of pages. We suppose that the memory
63    // index starts at 0.
64
65    // the address that the first page starts on
66    let start_page_address: usize =
67        (code_section_starting_address / page_size_usize) * page_size_usize;
68
69    // the address that the last page starts on
70    let end_page_address = (code_section_end_address / page_size_usize) * page_size_usize;
71
72    let first_page_index = start_page_address / page_size_usize;
73
74    let last_page_index = end_page_address / page_size_usize;
75
76    let mut data_offset = 0;
77    (first_page_index..=last_page_index).for_each(|page_index| {
78        let mut data = vec![0; page_size_usize];
79        // Special case where all code fits in one page
80        if first_page_index == last_page_index {
81            let data_length = code_section_end_address - code_section_starting_address;
82            let page_offset = code_section_starting_address - start_page_address;
83            data[page_offset..page_offset + data_length]
84                .copy_from_slice(&text_section_data[0..data_length]);
85            data_offset += data_length;
86        } else {
87            let data_length = if page_index == last_page_index {
88                code_section_end_address - end_page_address
89            } else {
90                page_size_usize
91            };
92            let page_offset = if page_index == first_page_index {
93                code_section_starting_address - start_page_address
94            } else {
95                0
96            };
97            data[page_offset..page_offset + data_length]
98                .copy_from_slice(&text_section_data[data_offset..data_offset + data_length]);
99
100            data_offset += data_length;
101        }
102        let page = Page {
103            index: page_index as u32,
104            data,
105        };
106        memory.push(page);
107    });
108
109    // FIXME: add data section into memory for static data saved in the binary
110
111    // FIXME: we're lucky that RISCV32i and MIPS have the same number of
112    let registers: [u32; 32] = [0; 32];
113
114    // FIXME: it is only because we share the same structure for the state.
115    let preimage_key: [u8; 32] = [0; 32];
116    // FIXME: it is only because we share the same structure for the state.
117    let preimage_offset = 0;
118
119    // Entry point of the program
120    let pc: u32 = file.ehdr.e_entry as u32;
121    assert!(pc != 0, "Entry point is 0. The documentation of the ELF library says that it means the ELF doesn't have an entry point. This is not supported. This can happen if the binary given is an object file and not an executable file. You might need to call the linker (ld) before running the binary.");
122    let next_pc: u32 = pc + 4u32;
123
124    let state = State {
125        memory,
126        // FIXME: only because Cannon related
127        preimage_key,
128        // FIXME: only because Cannon related
129        preimage_offset,
130        pc,
131        next_pc,
132        // FIXME: only because Cannon related
133        lo: 0,
134        // FIXME: only because Cannon related
135        hi: 0,
136        heap: 0,
137        exit: 0,
138        exited: false,
139        step: 0,
140        registers,
141        // FIXME: only because Cannon related
142        last_hint: None,
143        // FIXME: only because Cannon related
144        preimage: None,
145    };
146
147    Ok(state)
148}
149
150pub fn parse_elf(arch: Architecture, path: &Path) -> Result<State, String> {
151    debug!("Start parsing the ELF file to load a compatible state");
152    let file_data = std::fs::read(path).expect("Could not read file.");
153    let slice = file_data.as_slice();
154    match arch {
155        Architecture::Mips => {
156            let file = ElfBytes::<BigEndian>::minimal_parse(slice).expect("Open ELF file failed.");
157            assert_eq!(file.ehdr.e_machine, 8);
158            make_state(file)
159        }
160        Architecture::RiscV32 => {
161            let file =
162                ElfBytes::<LittleEndian>::minimal_parse(slice).expect("Open ELF file failed.");
163            assert_eq!(file.ehdr.e_machine, 243);
164            make_state(file)
165        }
166    }
167}