sbpf_assembler/
section.rs

1use crate::astnode::ASTNode;
2use crate::header::SectionHeader;
3use crate::dynsym::DynamicSymbol;
4use crate::dynsym::RelDyn;
5use crate::lexer::Token;
6use crate::debuginfo::DebugInfo;
7use std::collections::HashMap;
8use crate::astnode::ROData;
9
10// Base Section trait
11pub trait Section {
12    fn name(&self) -> &str {
13        ".unknown"  // Default section name
14    }
15    
16    fn bytecode(&self) -> Vec<u8> {
17        Vec::new()  // Default empty bytecode
18    }
19    
20    // fn get_size(&self) -> u64
21    fn size(&self) -> u64 {
22        self.bytecode().len() as u64
23    }
24
25    // fn get_aligned_size(&self) -> u64
26
27    // fn section_header_bytecode(&self) -> Vec<u8>
28}
29
30// Code Section implementation
31#[derive(Debug)]
32pub struct CodeSection {
33    name: String,
34    nodes: Vec<ASTNode>,
35    size: u64,
36    offset: u64,
37    debug_map: HashMap<u64, DebugInfo>,
38}
39
40impl CodeSection {
41    pub fn new(nodes: Vec<ASTNode>, size: u64) -> Self {
42        let mut debug_map = HashMap::new();
43        for node in &nodes {
44            if let Some((_, node_debug_map)) = node.bytecode_with_debug_map() {
45                debug_map.extend(node_debug_map);
46            }
47        }
48        Self {
49            name: String::from(".text"),
50            nodes,
51            size,
52            offset: 0,
53            debug_map,
54        }
55    }
56
57    pub fn get_nodes(&self) -> &Vec<ASTNode> {
58        &self.nodes
59    }
60
61    pub fn get_size(&self) -> u64 {
62        self.size
63    }
64
65    pub fn get_debug_map(&self) -> &HashMap<u64, DebugInfo> {
66        &self.debug_map
67    }
68
69    pub fn set_offset(&mut self, offset: u64) {
70        self.offset = offset;
71    }
72
73
74    pub fn section_header_bytecode(&self) -> Vec<u8> {
75        let flags = SectionHeader::SHF_ALLOC | SectionHeader::SHF_EXECINSTR;
76        SectionHeader::new(
77            1,
78            SectionHeader::SHT_PROGBITS,
79            flags,
80            self.offset,
81            self.offset,
82            self.size,
83            0,
84            0,
85            4,
86            0
87        ).bytecode()
88    }
89}
90
91impl Section for CodeSection {
92    fn name(&self) -> &str {
93        &self.name
94    }
95
96    fn bytecode(&self) -> Vec<u8> {
97        let mut bytecode = Vec::new();
98        for node in &self.nodes {
99            if let Some(node_bytes) = node.bytecode() {
100                bytecode.extend(node_bytes);
101            }
102        }
103        bytecode
104    }
105
106    fn size(&self) -> u64 {
107        self.size
108    }
109}
110
111// Data Section implementation
112#[derive(Debug)]
113pub struct DataSection {
114    name: String,
115    nodes: Vec<ASTNode>,
116    size: u64,
117    offset: u64,
118}
119
120impl DataSection {
121    pub fn new(nodes: Vec<ASTNode>, size: u64) -> Self {
122        Self {
123            name: String::from(".rodata"),
124            nodes,
125            size,
126            offset: 0,
127        }
128    }
129
130    pub fn get_nodes(&self) -> &Vec<ASTNode> {
131        &self.nodes
132    }
133
134    pub fn get_size(&self) -> u64 {
135        self.size
136    }
137
138    pub fn set_offset(&mut self, offset: u64) {
139        self.offset = offset;
140    }
141
142    pub fn rodata(&self) -> Vec<(String, usize, String)> {
143        let mut ro_data_labels = Vec::new();
144        for node in &self.nodes {    
145            if let ASTNode::ROData { rodata: ROData { name, args, .. }, offset } = node {
146                if let Some(Token::StringLiteral(str_literal, _)) = args.get(1) {
147                    ro_data_labels.push((name.clone(), offset.clone() as usize, str_literal.clone()));
148                }
149            }
150        }
151        ro_data_labels
152    }
153
154    pub fn section_header_bytecode(&self) -> Vec<u8> {
155        let flags = SectionHeader::SHF_ALLOC;  // Read-only data
156        SectionHeader::new(
157            7,
158            SectionHeader::SHT_PROGBITS,
159            flags,
160            self.offset,
161            self.offset,
162            self.size,
163            0,
164            0,
165            1,
166            0
167        ).bytecode()
168    }
169}
170
171impl Section for DataSection {
172    fn name(&self) -> &str {
173        &self.name
174    }
175
176    fn size(&self) -> u64 {
177        self.size
178    }
179
180    fn bytecode(&self) -> Vec<u8> {
181        let mut bytecode = Vec::new();
182        for node in &self.nodes {
183            if let Some(node_bytes) = node.bytecode() {
184                bytecode.extend(node_bytes);
185            }
186        }
187        // Add padding to make size multiple of 8
188        while bytecode.len() % 8 != 0 {
189            bytecode.push(0);
190        }
191
192        bytecode
193    }
194}
195
196#[derive(Debug)]
197pub struct NullSection {
198    name: String,
199    offset: u64,
200}
201
202impl NullSection {
203    pub fn new() -> Self {
204        Self {
205            name: String::from(""),
206            offset: 0,
207        }
208    }
209
210    pub fn section_header_bytecode(&self) -> Vec<u8> {
211        SectionHeader::new(
212            0,
213            SectionHeader::SHT_NULL,
214            0,
215            0,
216            0,
217            0,
218            0,
219            0,
220            0,
221            0
222        ).bytecode()
223    }
224
225}
226
227impl Section for NullSection {
228    // We can use all default implementations from the Section trait
229}
230
231#[derive(Debug)]
232pub struct ShStrTabSection {
233    name: String,
234    name_offset: u32,
235    section_names: Vec<String>,
236    offset: u64,
237}
238
239impl ShStrTabSection {
240    pub fn new(name_offset: u32, section_names: Vec<String>) -> Self {
241        Self {
242            name: String::from(".s"),
243            name_offset,
244            section_names: {
245                let mut names = section_names;
246                names.push(".s".to_string());
247                names
248            },
249            offset: 0,
250        }
251    }
252
253    pub fn set_offset(&mut self, offset: u64) {
254        self.offset = offset;
255    }
256
257    pub fn section_header_bytecode(&self) -> Vec<u8> {
258        SectionHeader::new(
259            self.name_offset,
260            SectionHeader::SHT_STRTAB,
261            0,
262            0,
263            self.offset,
264            self.size(),
265            0,
266            0,
267            1,
268            0
269        ).bytecode()
270    }
271
272}
273
274impl Section for ShStrTabSection {
275    fn name(&self) -> &str {
276        &self.name
277    }
278
279    fn bytecode(&self) -> Vec<u8> {
280        let mut bytes = Vec::new();
281        // First byte is null
282        bytes.push(0);
283        
284        // Add each non-empty section name with null terminator
285        for name in &self.section_names {
286            if !name.is_empty() {
287                bytes.extend(name.as_bytes());
288                bytes.push(0); // null terminator
289            }
290        }
291
292        // Add padding to make size multiple of 8
293        while bytes.len() % 8 != 0 {
294            bytes.push(0);
295        }
296        
297        bytes
298    }
299    
300    fn size(&self) -> u64 {
301        // Calculate section header offset
302        let mut section_name_size = 0;
303        
304        for name in &self.section_names {
305            if !name.is_empty() {
306                section_name_size += 1 + name.len();
307            }
308        }
309
310        section_name_size += 1; // null section
311        
312        section_name_size as u64  // Return the calculated size
313    }
314}
315
316#[derive(Debug)]
317pub struct DynamicSection {
318    name: String,
319    name_offset: u32,
320    offset: u64,
321    rel_offset: u64,
322    rel_size: u64,
323    rel_count: u64,
324    dynsym_offset: u64,
325    dynstr_offset: u64,
326    dynstr_size: u64,
327}
328
329impl DynamicSection {
330    pub fn new(name_offset: u32) -> Self {
331        Self {
332            name: String::from(".dynamic"),
333            name_offset,
334            offset: 0,
335            rel_offset: 0,
336            rel_size: 0,
337            rel_count: 0,
338            dynsym_offset: 0,
339            dynstr_offset: 0,
340            dynstr_size: 0,
341        }
342    }
343
344    pub fn set_offset(&mut self, offset: u64) {
345        self.offset = offset;
346    }
347
348    pub fn set_rel_offset(&mut self, offset: u64) {
349        self.rel_offset = offset;
350    }
351
352    pub fn set_rel_size(&mut self, size: u64) {
353        self.rel_size = size;
354    }
355
356    pub fn set_rel_count(&mut self, count: u64) {
357        self.rel_count = count;
358    }
359
360    pub fn set_dynsym_offset(&mut self, offset: u64) {
361        self.dynsym_offset = offset;
362    }
363
364    pub fn set_dynstr_offset(&mut self, offset: u64) {
365        self.dynstr_offset = offset;
366    }
367
368    pub fn set_dynstr_size(&mut self, size: u64) {
369        self.dynstr_size = size;
370    }
371
372    pub fn section_header_bytecode(&self) -> Vec<u8> {
373        SectionHeader::new(
374            self.name_offset,
375            SectionHeader::SHT_DYNAMIC,
376            SectionHeader::SHF_ALLOC | SectionHeader::SHF_WRITE,
377            self.offset,
378            self.offset,
379            self.size(),
380            5,
381            0,
382            8,
383            16
384        ).bytecode()
385    }
386
387}
388
389impl Section for DynamicSection {
390    fn name(&self) -> &str {
391        &self.name
392    }
393
394    fn bytecode(&self) -> Vec<u8> {
395        let mut bytes = Vec::new();
396        
397        // DT_FLAGS (DF_TEXTREL)
398        bytes.extend_from_slice(&0x1e_u64.to_le_bytes());
399        bytes.extend_from_slice(&0x04_u64.to_le_bytes());
400        
401        // DT_REL
402        bytes.extend_from_slice(&0x11_u64.to_le_bytes());
403        bytes.extend_from_slice(&self.rel_offset.to_le_bytes());
404        
405        // DT_RELSZ
406        bytes.extend_from_slice(&0x12_u64.to_le_bytes());
407        bytes.extend_from_slice(&self.rel_size.to_le_bytes());
408        
409        // DT_RELENT
410        bytes.extend_from_slice(&0x13_u64.to_le_bytes());
411        bytes.extend_from_slice(&0x10_u64.to_le_bytes());  // Constant: 16 bytes per entry
412        
413        // DT_RELCOUNT: number of relative relocation entries
414        if self.rel_count > 0 {
415            bytes.extend_from_slice(&0x6fffff_fa_u64.to_le_bytes());
416            bytes.extend_from_slice(&self.rel_count.to_le_bytes());
417        }
418        
419        // DT_SYMTAB
420        bytes.extend_from_slice(&0x06_u64.to_le_bytes());
421        bytes.extend_from_slice(&self.dynsym_offset.to_le_bytes());
422        
423        // DT_SYMENT
424        bytes.extend_from_slice(&0x0b_u64.to_le_bytes());
425        bytes.extend_from_slice(&0x18_u64.to_le_bytes());  // Constant: 24 bytes per symbol
426        
427        // DT_STRTAB
428        bytes.extend_from_slice(&0x05_u64.to_le_bytes());
429        bytes.extend_from_slice(&self.dynstr_offset.to_le_bytes());
430        
431        // DT_STRSZ
432        bytes.extend_from_slice(&0x0a_u64.to_le_bytes());
433        bytes.extend_from_slice(&self.dynstr_size.to_le_bytes());
434        
435        // DT_TEXTREL
436        bytes.extend_from_slice(&0x16_u64.to_le_bytes());
437        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
438        
439        // DT_NULL
440        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
441        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
442        
443        bytes
444    }
445
446    fn size(&self) -> u64 {
447        if self.rel_count > 0 {
448            return 11 * 16;
449        } else {
450            return 10 * 16;
451        }
452    }
453}
454
455#[derive(Debug)]
456pub struct DynStrSection {
457    name: String,
458    name_offset: u32,
459    symbol_names: Vec<String>,
460    offset: u64,
461}
462
463impl DynStrSection {
464    pub fn new(name_offset: u32, symbol_names: Vec<String>) -> Self {
465        Self {
466            name: String::from(".dynstr"),
467            name_offset,
468            symbol_names,
469            offset: 0,
470        }
471    }
472
473    pub fn set_offset(&mut self, offset: u64) {
474        self.offset = offset;
475    }
476
477    pub fn section_header_bytecode(&self) -> Vec<u8> {
478        SectionHeader::new(
479            self.name_offset,
480            SectionHeader::SHT_STRTAB,
481            SectionHeader::SHF_ALLOC,  // Allocatable section
482            self.offset,
483            self.offset,
484            self.size(),
485            0,
486            0,
487            1,
488            0
489        ).bytecode()
490    }
491
492}
493
494impl Section for DynStrSection {
495    fn name(&self) -> &str {
496        &self.name
497    }
498
499    fn bytecode(&self) -> Vec<u8> {
500        let mut bytes = Vec::new();
501        // First byte is null
502        bytes.push(0);
503        
504        // Add each symbol name with null terminator
505        for name in &self.symbol_names {
506            bytes.extend(name.as_bytes());
507            bytes.push(0); // null terminator
508        }
509        // add padding to make size multiple of 8
510        while bytes.len() % 8 != 0 {
511            bytes.push(0);
512        }
513        bytes
514    }
515    
516    fn size(&self) -> u64 {
517        // Calculate total size: initial null byte + sum of (name lengths + null terminators)
518        let mut size = 1 + self.symbol_names.iter()
519            .map(|name| name.len() + 1)
520            .sum::<usize>();
521        // add padding to make size multiple of 8
522        while size % 8 != 0 {
523            size += 1;
524        }
525        size as u64
526    }
527}
528
529#[derive(Debug)]
530pub struct DynSymSection {
531    name: String,
532    name_offset: u32,
533    offset: u64,
534    symbols: Vec<DynamicSymbol>,
535}
536
537impl DynSymSection {
538    pub fn new(name_offset: u32, symbols: Vec<DynamicSymbol>) -> Self {
539        Self {
540            name: String::from(".dynsym"),
541            name_offset,
542            offset: 0,
543            symbols,
544        }
545    }
546
547    pub fn set_offset(&mut self, offset: u64) {
548        self.offset = offset;
549    }
550
551    pub fn section_header_bytecode(&self) -> Vec<u8> {
552        let flags = SectionHeader::SHF_ALLOC;
553        SectionHeader::new(
554            self.name_offset,
555            SectionHeader::SHT_DYNSYM,
556            flags,
557            self.offset,
558            self.offset,
559            self.size(),
560            5,
561            1,
562            8,
563            24
564        ).bytecode()
565    }
566
567}
568
569impl Section for DynSymSection {
570    fn name(&self) -> &str {
571        &self.name
572    }
573
574    fn size(&self) -> u64 {
575        // Each symbol entry is 24 bytes
576        (self.symbols.len() as u64) * 24
577    }
578
579    fn bytecode(&self) -> Vec<u8> {
580        let mut bytes = Vec::new();
581        for symbol in &self.symbols {
582            bytes.extend(symbol.bytecode());
583        }
584        bytes
585    }
586    
587}   
588
589#[derive(Debug)]
590pub struct RelDynSection {
591    name: String,
592    name_offset: u32,
593    offset: u64,
594    entries: Vec<RelDyn>,
595}
596
597impl RelDynSection {
598    pub fn new(name_offset: u32, entries: Vec<RelDyn>) -> Self {
599        Self {
600            name: String::from(".rel.dyn"),
601            name_offset,
602            offset: 0,
603            entries,
604        }
605    }
606
607    pub fn set_offset(&mut self, offset: u64) {
608        self.offset = offset;
609    }
610
611    pub fn size(&self) -> u64 {
612        (self.entries.len() * 16) as u64 // Each RelDyn entry is 16 bytes
613    }
614
615    pub fn section_header_bytecode(&self) -> Vec<u8> {
616        let flags = SectionHeader::SHF_ALLOC;
617        SectionHeader::new(
618            self.name_offset,
619            SectionHeader::SHT_REL,
620            flags,
621            self.offset,
622            self.offset,
623            self.size(),
624            4,
625            0,
626            8,
627            16
628        ).bytecode()
629    }
630
631}
632
633impl Section for RelDynSection {
634    fn name(&self) -> &str {
635        &self.name
636    }
637
638    fn size(&self) -> u64 {
639        self.size()
640    }
641
642    fn bytecode(&self) -> Vec<u8> {
643        let mut bytes = Vec::new();
644        for entry in &self.entries {
645            bytes.extend(entry.bytecode());
646        }
647        bytes
648    }
649
650
651}
652
653#[derive(Debug)]
654pub enum SectionType {
655    Code(CodeSection),
656    Data(DataSection),
657    ShStrTab(ShStrTabSection),
658    Dynamic(DynamicSection),
659    DynStr(DynStrSection),
660    DynSym(DynSymSection),
661    Default(NullSection),
662    RelDyn(RelDynSection),
663}
664
665impl SectionType {
666    pub fn name(&self) -> &str {
667        match self {
668            SectionType::Code(cs) => &cs.name,
669            SectionType::Data(ds) => &ds.name,
670            SectionType::ShStrTab(ss) => &ss.name,
671            SectionType::Dynamic(ds) => &ds.name,
672            SectionType::DynStr(ds) => &ds.name,
673            SectionType::DynSym(ds) => &ds.name,
674            SectionType::Default(ds) => &ds.name,
675            SectionType::RelDyn(ds) => &ds.name,
676        }
677    }
678
679    pub fn bytecode(&self) -> Vec<u8> {
680        match self {
681            SectionType::Code(cs) => cs.bytecode(),
682            SectionType::Data(ds) => ds.bytecode(),
683            SectionType::ShStrTab(ss) => ss.bytecode(),
684            SectionType::Dynamic(ds) => ds.bytecode(),
685            SectionType::DynStr(ds) => ds.bytecode(),
686            SectionType::DynSym(ds) => ds.bytecode(),
687            SectionType::Default(ds) => ds.bytecode(),
688            SectionType::RelDyn(ds) => ds.bytecode(),
689        }
690    }
691
692    pub fn size(&self) -> u64 {
693        match self {
694            SectionType::Code(cs) => cs.size(),
695            SectionType::Data(ds) => ds.size(),
696            SectionType::ShStrTab(ss) => ss.size(),
697            SectionType::Dynamic(ds) => ds.size(),
698            SectionType::DynStr(ds) => ds.size(),
699            SectionType::DynSym(ds) => ds.size(),
700            SectionType::Default(ds) => ds.size(),
701            SectionType::RelDyn(ds) => ds.size(),
702        }
703    }
704
705    pub fn section_header_bytecode(&self) -> Vec<u8> {
706        match self {
707            SectionType::Code(cs) => cs.section_header_bytecode(),
708            SectionType::Data(ds) => ds.section_header_bytecode(),
709            SectionType::ShStrTab(ss) => ss.section_header_bytecode(),
710            SectionType::Dynamic(ds) => ds.section_header_bytecode(),
711            SectionType::DynStr(ds) => ds.section_header_bytecode(),
712            SectionType::DynSym(ds) => ds.section_header_bytecode(),
713            SectionType::Default(ds) => ds.section_header_bytecode(),
714            SectionType::RelDyn(ds) => ds.section_header_bytecode(),
715        }
716    }
717
718    pub fn set_offset(&mut self, offset: u64) {
719        match self {
720            SectionType::Code(cs) => cs.set_offset(offset),
721            SectionType::Data(ds) => ds.set_offset(offset),
722            SectionType::ShStrTab(ss) => ss.set_offset(offset),
723            SectionType::Dynamic(ds) => ds.set_offset(offset),
724            SectionType::DynStr(ds) => ds.set_offset(offset),
725            SectionType::DynSym(ds) => ds.set_offset(offset),
726            SectionType::RelDyn(ds) => ds.set_offset(offset),
727            SectionType::Default(_) => (), // NullSection doesn't need offset
728        }
729    }
730
731    pub fn offset(&self) -> u64 {
732        match self {
733            SectionType::Code(cs) => cs.offset,
734            SectionType::Data(ds) => ds.offset,
735            SectionType::ShStrTab(ss) => ss.offset,
736            SectionType::Dynamic(ds) => ds.offset,
737            SectionType::DynStr(ds) => ds.offset,
738            SectionType::DynSym(ds) => ds.offset,
739            SectionType::Default(ns) => ns.offset,
740            SectionType::RelDyn(rs) => rs.offset,
741        }
742    }
743}
744
745