sbpf_assembler/
program.rs

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

sbpf_assembler/
program.rs
