sbpf_assembler/
program.rs

1use crate::debuginfo::DebugInfo;
2use crate::dynsym::{DynamicSymbol, RelDyn, RelocationType};
3use crate::header::ElfHeader;
4use crate::header::ProgramHeader;
5use crate::parser::ParseResult;
6use crate::section::{
7    DynStrSection, DynSymSection, DynamicSection, NullSection, RelDynSection, Section, SectionType,
8    ShStrTabSection,
9};
10use std::collections::HashMap;
11use std::fs::File;
12use std::io::Write;
13use std::path::Path;
14
15#[derive(Debug)]
16pub struct Program {
17    pub elf_header: ElfHeader,
18    pub program_headers: Option<Vec<ProgramHeader>>,
19    pub sections: Vec<SectionType>,
20}
21
22impl Program {
23    pub fn from_parse_result(
24        ParseResult {
25            code_section,
26            data_section,
27            dynamic_symbols,
28            relocation_data,
29            prog_is_static,
30        }: ParseResult,
31    ) -> Self {
32        let mut elf_header = ElfHeader::new();
33        let mut program_headers = None;
34
35        // omit program headers if static
36        let ph_count = if prog_is_static { 0 } else { 3 };
37        elf_header.e_phnum = ph_count;
38
39        // save read + execute size for program header before
40        // ownership of code/data sections is transferred
41        let text_size = code_section.size() + data_section.size();
42
43        // Calculate base offset after ELF header and program headers
44        let mut current_offset = 64 + (ph_count as u64 * 56); // 64 bytes ELF header, 56 bytes per program header
45        elf_header.e_entry = current_offset;
46
47        // Create a vector of sections
48        let mut sections = Vec::new();
49        sections.push(SectionType::Default(NullSection::new()));
50
51        let mut section_names = Vec::new();
52
53        // Code section
54        let mut text_section = SectionType::Code(code_section);
55        text_section.set_offset(current_offset);
56        current_offset += text_section.size();
57        section_names.push(text_section.name().to_string());
58        sections.push(text_section);
59
60        // Data section
61        if data_section.size() > 0 {
62            let mut rodata_section = SectionType::Data(data_section);
63            rodata_section.set_offset(current_offset);
64            current_offset += rodata_section.size();
65            section_names.push(rodata_section.name().to_string());
66            sections.push(rodata_section);
67        }
68
69        let padding = (8 - (current_offset % 8)) % 8;
70        current_offset += padding;
71
72        if !prog_is_static {
73            let mut symbol_names = Vec::new();
74            let mut dyn_syms = Vec::new();
75            let mut dyn_str_offset = 1;
76
77            dyn_syms.push(DynamicSymbol::new(0, 0, 0, 0, 0, 0));
78
79            // all symbols handled right now are all global symbols
80            for (name, _) in dynamic_symbols.get_entry_points() {
81                symbol_names.push(name.clone());
82                dyn_syms.push(DynamicSymbol::new(
83                    dyn_str_offset as u32,
84                    0x10,
85                    0,
86                    1,
87                    elf_header.e_entry,
88                    0,
89                ));
90                dyn_str_offset += name.len() + 1;
91            }
92
93            for (name, _) in dynamic_symbols.get_call_targets() {
94                symbol_names.push(name.clone());
95                dyn_syms.push(DynamicSymbol::new(dyn_str_offset as u32, 0x10, 0, 0, 0, 0));
96                dyn_str_offset += name.len() + 1;
97            }
98
99            let mut rel_count = 0;
100            let mut rel_dyns = Vec::new();
101            for (offset, rel_type, name) in relocation_data.get_rel_dyns() {
102                if rel_type == RelocationType::RSbfSyscall {
103                    if let Some(index) = symbol_names.iter().position(|n| *n == name) {
104                        rel_dyns.push(RelDyn::new(
105                            offset + elf_header.e_entry,
106                            rel_type as u64,
107                            index as u64 + 1,
108                        ));
109                    } else {
110                        panic!("Symbol {} not found in symbol_names", name);
111                    }
112                } else if rel_type == RelocationType::RSbf64Relative {
113                    rel_count += 1;
114                    rel_dyns.push(RelDyn::new(offset + elf_header.e_entry, rel_type as u64, 0));
115                }
116            }
117            let mut dynamic_section = SectionType::Dynamic(DynamicSection::new(
118                (section_names
119                    .iter()
120                    .map(|name| name.len() + 1)
121                    .sum::<usize>()
122                    + 1) as u32,
123            ));
124            dynamic_section.set_offset(current_offset);
125            if let SectionType::Dynamic(ref mut dynamic_section) = dynamic_section {
126                dynamic_section.set_rel_count(rel_count);
127            }
128            current_offset += dynamic_section.size();
129            section_names.push(dynamic_section.name().to_string());
130
131            let mut dynsym_section = SectionType::DynSym(DynSymSection::new(
132                (section_names
133                    .iter()
134                    .map(|name| name.len() + 1)
135                    .sum::<usize>()
136                    + 1) as u32,
137                dyn_syms,
138            ));
139            dynsym_section.set_offset(current_offset);
140            current_offset += dynsym_section.size();
141            section_names.push(dynsym_section.name().to_string());
142
143            let mut dynstr_section = SectionType::DynStr(DynStrSection::new(
144                (section_names
145                    .iter()
146                    .map(|name| name.len() + 1)
147                    .sum::<usize>()
148                    + 1) as u32,
149                symbol_names,
150            ));
151            dynstr_section.set_offset(current_offset);
152            current_offset += dynstr_section.size();
153            section_names.push(dynstr_section.name().to_string());
154
155            let mut rel_dyn_section = SectionType::RelDyn(RelDynSection::new(
156                (section_names
157                    .iter()
158                    .map(|name| name.len() + 1)
159                    .sum::<usize>()
160                    + 1) as u32,
161                rel_dyns,
162            ));
163            rel_dyn_section.set_offset(current_offset);
164            current_offset += rel_dyn_section.size();
165            section_names.push(rel_dyn_section.name().to_string());
166
167            if let SectionType::Dynamic(ref mut dynamic_section) = dynamic_section {
168                dynamic_section.set_rel_offset(rel_dyn_section.offset());
169                dynamic_section.set_rel_size(rel_dyn_section.size());
170                dynamic_section.set_dynsym_offset(dynsym_section.offset());
171                dynamic_section.set_dynstr_offset(dynstr_section.offset());
172                dynamic_section.set_dynstr_size(dynstr_section.size());
173            }
174
175            let mut shstrtab_section = SectionType::ShStrTab(ShStrTabSection::new(
176                (section_names
177                    .iter()
178                    .map(|name| name.len() + 1)
179                    .sum::<usize>()
180                    + 1) as u32,
181                section_names,
182            ));
183            shstrtab_section.set_offset(current_offset);
184            current_offset += shstrtab_section.size();
185
186            program_headers = Some(vec![
187                ProgramHeader::new_load(
188                    elf_header.e_entry,
189                    text_size,
190                    true, // executable
191                ),
192                ProgramHeader::new_load(
193                    dynsym_section.offset(),
194                    dynsym_section.size() + dynstr_section.size() + rel_dyn_section.size(),
195                    false,
196                ),
197                ProgramHeader::new_dynamic(dynamic_section.offset(), dynamic_section.size()),
198            ]);
199
200            sections.push(dynamic_section);
201            sections.push(dynsym_section);
202            sections.push(dynstr_section);
203            sections.push(rel_dyn_section);
204            sections.push(shstrtab_section);
205        } else {
206            // Create a vector of section names
207            let mut section_names = Vec::new();
208            for section in &sections {
209                section_names.push(section.name().to_string());
210            }
211
212            let mut shstrtab_section = ShStrTabSection::new(
213                section_names
214                    .iter()
215                    .map(|name| name.len() + 1)
216                    .sum::<usize>() as u32,
217                section_names,
218            );
219            shstrtab_section.set_offset(current_offset);
220            current_offset += shstrtab_section.size();
221            sections.push(SectionType::ShStrTab(shstrtab_section));
222        }
223
224        // Update section header offset in ELF header
225        let padding = (8 - (current_offset % 8)) % 8;
226        elf_header.e_shoff = current_offset + padding;
227        elf_header.e_shnum = sections.len() as u16;
228        elf_header.e_shstrndx = sections.len() as u16 - 1;
229
230        Self {
231            elf_header,
232            program_headers,
233            sections,
234        }
235    }
236
237    pub fn emit_bytecode(&self) -> Vec<u8> {
238        let mut bytes = Vec::new();
239
240        // Emit ELF Header bytes
241        bytes.extend(self.elf_header.bytecode());
242
243        // Emit program headers
244        if self.program_headers.is_some() {
245            for ph in self.program_headers.as_ref().unwrap() {
246                bytes.extend(ph.bytecode());
247            }
248        }
249
250        // Emit sections
251        for section in &self.sections {
252            bytes.extend(section.bytecode());
253        }
254
255        // Emit section headers
256        for section in &self.sections {
257            bytes.extend(section.section_header_bytecode());
258        }
259
260        bytes
261    }
262
263    pub fn has_rodata(&self) -> bool {
264        self.sections.iter().any(|s| s.name() == ".rodata")
265    }
266
267    pub fn parse_rodata(&self) -> Vec<(String, usize, String)> {
268        let rodata = self
269            .sections
270            .iter()
271            .find(|s| s.name() == ".rodata")
272            .unwrap();
273        if let SectionType::Data(data_section) = rodata {
274            data_section.rodata()
275        } else {
276            panic!("ROData section not found");
277        }
278    }
279
280    pub fn get_debug_map(&self) -> HashMap<u64, DebugInfo> {
281        let code = self.sections.iter().find(|s| s.name() == ".text").unwrap();
282        if let SectionType::Code(code_section) = code {
283            code_section.get_debug_map().clone()
284        } else {
285            panic!("Code section not found");
286        }
287    }
288
289    pub fn save_to_file(&self, input_path: &str) -> std::io::Result<()> {
290        // Get the file stem (name without extension) from input path
291        let path = Path::new(input_path);
292        let file_stem = path
293            .file_stem()
294            .and_then(|s| s.to_str())
295            .unwrap_or("output");
296
297        // Create the output file name with .so extension
298        let output_path = format!("{}.so", file_stem);
299
300        // Get the bytecode
301        let bytes = self.emit_bytecode();
302
303        // Write bytes to file
304        let mut file = File::create(output_path)?;
305        file.write_all(&bytes)?;
306
307        Ok(())
308    }
309}
sbpf_assembler/program.rs

sbpf_assembler/
program.rs
