Skip to main content

sbpf_assembler/
section.rs

1use crate::{
2    astnode::{ASTNode, ROData},
3    dynsym::{DynamicSymbol, RelDyn},
4    header::SectionHeader,
5    parser::Token,
6};
7
8// Base Section trait
9pub trait Section {
10    fn name(&self) -> &str {
11        ".unknown" // Default section name
12    }
13
14    fn bytecode(&self) -> Vec<u8> {
15        Vec::new() // Default empty bytecode
16    }
17
18    // fn get_size(&self) -> u64
19    fn size(&self) -> u64 {
20        self.bytecode().len() as u64
21    }
22
23    // fn get_aligned_size(&self) -> u64
24
25    // fn section_header_bytecode(&self) -> Vec<u8>
26}
27
28// Code Section implementation
29#[derive(Debug)]
30pub struct CodeSection {
31    name: String,
32    nodes: Vec<ASTNode>,
33    size: u64,
34    offset: u64,
35    vaddr: u64,
36}
37
38impl CodeSection {
39    pub fn new(nodes: Vec<ASTNode>, size: u64) -> Self {
40        Self {
41            name: String::from(".text"),
42            nodes,
43            size,
44            offset: 0,
45            vaddr: 0,
46        }
47    }
48
49    pub fn get_nodes(&self) -> &Vec<ASTNode> {
50        &self.nodes
51    }
52
53    pub fn get_size(&self) -> u64 {
54        self.size
55    }
56
57    pub fn set_offset(&mut self, offset: u64) {
58        self.offset = offset;
59        self.vaddr = offset;
60    }
61
62    pub fn set_vaddr(&mut self, vaddr: u64) {
63        self.vaddr = vaddr;
64    }
65
66    pub fn section_header_bytecode(&self) -> Vec<u8> {
67        let flags = SectionHeader::SHF_ALLOC | SectionHeader::SHF_EXECINSTR;
68        SectionHeader::new(
69            1,
70            SectionHeader::SHT_PROGBITS,
71            flags,
72            self.vaddr,
73            self.offset,
74            self.size,
75            0,
76            0,
77            4,
78            0,
79        )
80        .bytecode()
81    }
82}
83
84impl Section for CodeSection {
85    fn name(&self) -> &str {
86        &self.name
87    }
88
89    fn bytecode(&self) -> Vec<u8> {
90        let mut bytecode = Vec::new();
91        for node in &self.nodes {
92            if let Some(node_bytes) = node.bytecode() {
93                bytecode.extend(node_bytes);
94            }
95        }
96        bytecode
97    }
98
99    fn size(&self) -> u64 {
100        self.size
101    }
102}
103
104// Data Section implementation
105#[derive(Debug)]
106pub struct DataSection {
107    name: String,
108    nodes: Vec<ASTNode>,
109    size: u64,
110    offset: u64,
111    vaddr: u64,
112}
113
114impl DataSection {
115    pub fn new(nodes: Vec<ASTNode>, size: u64) -> Self {
116        Self {
117            name: String::from(".rodata"),
118            nodes,
119            size,
120            offset: 0,
121            vaddr: 0,
122        }
123    }
124
125    pub fn get_nodes(&self) -> &Vec<ASTNode> {
126        &self.nodes
127    }
128
129    pub fn get_size(&self) -> u64 {
130        self.size
131    }
132
133    pub fn set_offset(&mut self, offset: u64) {
134        self.offset = offset;
135        self.vaddr = offset;
136    }
137
138    pub fn set_vaddr(&mut self, vaddr: u64) {
139        self.vaddr = vaddr;
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 {
146                rodata: ROData { name, args, .. },
147                offset,
148            } = node
149                && let Some(Token::StringLiteral(str_literal, _)) = args.get(1)
150            {
151                ro_data_labels.push((name.clone(), *offset as usize, str_literal.clone()));
152            }
153        }
154        ro_data_labels
155    }
156
157    pub fn section_header_bytecode(&self) -> Vec<u8> {
158        let flags = SectionHeader::SHF_ALLOC; // Read-only data
159        SectionHeader::new(
160            7,
161            SectionHeader::SHT_PROGBITS,
162            flags,
163            self.vaddr,
164            self.offset,
165            self.size,
166            0,
167            0,
168            1,
169            0,
170        )
171        .bytecode()
172    }
173}
174
175impl Section for DataSection {
176    fn name(&self) -> &str {
177        &self.name
178    }
179
180    fn size(&self) -> u64 {
181        // Return 8-byte aligned size
182        (self.size + 7) & !7
183    }
184
185    fn bytecode(&self) -> Vec<u8> {
186        let mut bytecode = Vec::new();
187        for node in &self.nodes {
188            if let Some(node_bytes) = node.bytecode() {
189                bytecode.extend(node_bytes);
190            }
191        }
192        // Add padding to make size multiple of 8
193        while bytecode.len() % 8 != 0 {
194            bytecode.push(0);
195        }
196
197        bytecode
198    }
199}
200
201#[derive(Debug, Default)]
202pub struct NullSection {
203    name: String,
204    offset: u64,
205}
206
207impl NullSection {
208    pub fn new() -> Self {
209        Self::default()
210    }
211
212    pub fn section_header_bytecode(&self) -> Vec<u8> {
213        SectionHeader::new(0, SectionHeader::SHT_NULL, 0, 0, 0, 0, 0, 0, 0, 0).bytecode()
214    }
215}
216
217impl Section for NullSection {
218    // We can use all default implementations from the Section trait
219}
220
221#[derive(Debug)]
222pub struct ShStrTabSection {
223    name: String,
224    name_offset: u32,
225    section_names: Vec<String>,
226    offset: u64,
227}
228
229impl ShStrTabSection {
230    pub fn new(name_offset: u32, section_names: Vec<String>) -> Self {
231        Self {
232            name: String::from(".s"),
233            name_offset,
234            section_names: {
235                let mut names = section_names;
236                names.push(".s".to_string());
237                names
238            },
239            offset: 0,
240        }
241    }
242
243    pub fn set_offset(&mut self, offset: u64) {
244        self.offset = offset;
245    }
246
247    pub fn section_header_bytecode(&self) -> Vec<u8> {
248        SectionHeader::new(
249            self.name_offset,
250            SectionHeader::SHT_STRTAB,
251            0,
252            0,
253            self.offset,
254            self.size(),
255            0,
256            0,
257            1,
258            0,
259        )
260        .bytecode()
261    }
262}
263
264impl Section for ShStrTabSection {
265    fn name(&self) -> &str {
266        &self.name
267    }
268
269    fn bytecode(&self) -> Vec<u8> {
270        let mut bytes = Vec::new();
271        // First byte is null
272        bytes.push(0);
273
274        // Add each non-empty section name with null terminator
275        for name in &self.section_names {
276            if !name.is_empty() {
277                bytes.extend(name.as_bytes());
278                bytes.push(0); // null terminator
279            }
280        }
281
282        // Add padding to make size multiple of 8
283        while bytes.len() % 8 != 0 {
284            bytes.push(0);
285        }
286
287        bytes
288    }
289
290    fn size(&self) -> u64 {
291        // Calculate section header offset
292        let mut section_name_size = 0;
293
294        for name in &self.section_names {
295            if !name.is_empty() {
296                section_name_size += 1 + name.len();
297            }
298        }
299
300        section_name_size += 1; // null section
301
302        section_name_size as u64 // Return the calculated size
303    }
304}
305
306#[derive(Debug)]
307pub struct DynamicSection {
308    name: String,
309    name_offset: u32,
310    offset: u64,
311    link: u32,
312    rel_offset: u64,
313    rel_size: u64,
314    rel_count: u64,
315    dynsym_offset: u64,
316    dynstr_offset: u64,
317    dynstr_size: u64,
318}
319
320impl DynamicSection {
321    pub fn new(name_offset: u32) -> Self {
322        Self {
323            name: String::from(".dynamic"),
324            name_offset,
325            offset: 0,
326            link: 0,
327            rel_offset: 0,
328            rel_size: 0,
329            rel_count: 0,
330            dynsym_offset: 0,
331            dynstr_offset: 0,
332            dynstr_size: 0,
333        }
334    }
335
336    pub fn set_offset(&mut self, offset: u64) {
337        self.offset = offset;
338    }
339
340    pub fn set_link(&mut self, link: u32) {
341        self.link = link;
342    }
343
344    pub fn set_rel_offset(&mut self, offset: u64) {
345        self.rel_offset = offset;
346    }
347
348    pub fn set_rel_size(&mut self, size: u64) {
349        self.rel_size = size;
350    }
351
352    pub fn set_rel_count(&mut self, count: u64) {
353        self.rel_count = count;
354    }
355
356    pub fn set_dynsym_offset(&mut self, offset: u64) {
357        self.dynsym_offset = offset;
358    }
359
360    pub fn set_dynstr_offset(&mut self, offset: u64) {
361        self.dynstr_offset = offset;
362    }
363
364    pub fn set_dynstr_size(&mut self, size: u64) {
365        self.dynstr_size = size;
366    }
367
368    pub fn section_header_bytecode(&self) -> Vec<u8> {
369        SectionHeader::new(
370            self.name_offset,
371            SectionHeader::SHT_DYNAMIC,
372            SectionHeader::SHF_ALLOC | SectionHeader::SHF_WRITE,
373            self.offset,
374            self.offset,
375            self.size(),
376            self.link,
377            0,
378            8,
379            16,
380        )
381        .bytecode()
382    }
383}
384
385impl Section for DynamicSection {
386    fn name(&self) -> &str {
387        &self.name
388    }
389
390    fn bytecode(&self) -> Vec<u8> {
391        let mut bytes = Vec::new();
392
393        // DT_FLAGS (DF_TEXTREL)
394        bytes.extend_from_slice(&0x1e_u64.to_le_bytes());
395        bytes.extend_from_slice(&0x04_u64.to_le_bytes());
396
397        // DT_REL
398        bytes.extend_from_slice(&0x11_u64.to_le_bytes());
399        bytes.extend_from_slice(&self.rel_offset.to_le_bytes());
400
401        // DT_RELSZ
402        bytes.extend_from_slice(&0x12_u64.to_le_bytes());
403        bytes.extend_from_slice(&self.rel_size.to_le_bytes());
404
405        // DT_RELENT
406        bytes.extend_from_slice(&0x13_u64.to_le_bytes());
407        bytes.extend_from_slice(&0x10_u64.to_le_bytes()); // Constant: 16 bytes per entry
408
409        // DT_RELCOUNT: number of relative relocation entries
410        if self.rel_count > 0 {
411            bytes.extend_from_slice(&0x6fff_fffa_u64.to_le_bytes());
412            bytes.extend_from_slice(&self.rel_count.to_le_bytes());
413        }
414
415        // DT_SYMTAB
416        bytes.extend_from_slice(&0x06_u64.to_le_bytes());
417        bytes.extend_from_slice(&self.dynsym_offset.to_le_bytes());
418
419        // DT_SYMENT
420        bytes.extend_from_slice(&0x0b_u64.to_le_bytes());
421        bytes.extend_from_slice(&0x18_u64.to_le_bytes()); // Constant: 24 bytes per symbol
422
423        // DT_STRTAB
424        bytes.extend_from_slice(&0x05_u64.to_le_bytes());
425        bytes.extend_from_slice(&self.dynstr_offset.to_le_bytes());
426
427        // DT_STRSZ
428        bytes.extend_from_slice(&0x0a_u64.to_le_bytes());
429        bytes.extend_from_slice(&self.dynstr_size.to_le_bytes());
430
431        // DT_TEXTREL
432        bytes.extend_from_slice(&0x16_u64.to_le_bytes());
433        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
434
435        // DT_NULL
436        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
437        bytes.extend_from_slice(&0x00_u64.to_le_bytes());
438
439        bytes
440    }
441
442    fn size(&self) -> u64 {
443        if self.rel_count > 0 { 11 * 16 } else { 10 * 16 }
444    }
445}
446
447#[derive(Debug)]
448pub struct DynStrSection {
449    name: String,
450    name_offset: u32,
451    symbol_names: Vec<String>,
452    offset: u64,
453}
454
455impl DynStrSection {
456    pub fn new(name_offset: u32, symbol_names: Vec<String>) -> Self {
457        Self {
458            name: String::from(".dynstr"),
459            name_offset,
460            symbol_names,
461            offset: 0,
462        }
463    }
464
465    pub fn set_offset(&mut self, offset: u64) {
466        self.offset = offset;
467    }
468
469    pub fn section_header_bytecode(&self) -> Vec<u8> {
470        SectionHeader::new(
471            self.name_offset,
472            SectionHeader::SHT_STRTAB,
473            SectionHeader::SHF_ALLOC, // Allocatable section
474            self.offset,
475            self.offset,
476            self.size(),
477            0,
478            0,
479            1,
480            0,
481        )
482        .bytecode()
483    }
484}
485
486impl Section for DynStrSection {
487    fn name(&self) -> &str {
488        &self.name
489    }
490
491    fn bytecode(&self) -> Vec<u8> {
492        let mut bytes = Vec::new();
493        // First byte is null
494        bytes.push(0);
495
496        // Add each symbol name with null terminator
497        for name in &self.symbol_names {
498            bytes.extend(name.as_bytes());
499            bytes.push(0); // null terminator
500        }
501        // add padding to make size multiple of 8
502        while bytes.len() % 8 != 0 {
503            bytes.push(0);
504        }
505        bytes
506    }
507
508    fn size(&self) -> u64 {
509        // Calculate total size: initial null byte + sum of (name lengths + null terminators)
510        let mut size = 1 + self
511            .symbol_names
512            .iter()
513            .map(|name| name.len() + 1)
514            .sum::<usize>();
515        // add padding to make size multiple of 8
516        while size % 8 != 0 {
517            size += 1;
518        }
519        size as u64
520    }
521}
522
523#[derive(Debug)]
524pub struct DynSymSection {
525    name: String,
526    name_offset: u32,
527    offset: u64,
528    link: u32,
529    symbols: Vec<DynamicSymbol>,
530}
531
532impl DynSymSection {
533    pub fn new(name_offset: u32, symbols: Vec<DynamicSymbol>) -> Self {
534        Self {
535            name: String::from(".dynsym"),
536            name_offset,
537            offset: 0,
538            link: 0,
539            symbols,
540        }
541    }
542
543    pub fn set_offset(&mut self, offset: u64) {
544        self.offset = offset;
545    }
546
547    pub fn set_link(&mut self, link: u32) {
548        self.link = link;
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            self.link,
561            1,
562            8,
563            24,
564        )
565        .bytecode()
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#[derive(Debug)]
589pub struct RelDynSection {
590    name: String,
591    name_offset: u32,
592    offset: u64,
593    link: u32,
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            link: 0,
604            entries,
605        }
606    }
607
608    pub fn set_offset(&mut self, offset: u64) {
609        self.offset = offset;
610    }
611
612    pub fn set_link(&mut self, link: u32) {
613        self.link = link;
614    }
615
616    pub fn size(&self) -> u64 {
617        (self.entries.len() * 16) as u64 // Each RelDyn entry is 16 bytes
618    }
619
620    pub fn section_header_bytecode(&self) -> Vec<u8> {
621        let flags = SectionHeader::SHF_ALLOC;
622        SectionHeader::new(
623            self.name_offset,
624            SectionHeader::SHT_REL,
625            flags,
626            self.offset,
627            self.offset,
628            self.size(),
629            self.link,
630            0,
631            8,
632            16,
633        )
634        .bytecode()
635    }
636}
637
638impl Section for RelDynSection {
639    fn name(&self) -> &str {
640        &self.name
641    }
642
643    fn size(&self) -> u64 {
644        self.size()
645    }
646
647    fn bytecode(&self) -> Vec<u8> {
648        let mut bytes = Vec::new();
649        for entry in &self.entries {
650            bytes.extend(entry.bytecode());
651        }
652        bytes
653    }
654}
655#[derive(Debug, Clone)]
656pub struct DebugSection {
657    name: String,
658    name_offset: u32,
659    data: Vec<u8>,
660    offset: u64,
661}
662
663impl DebugSection {
664    pub fn new(name: &str, name_offset: u32, data: Vec<u8>) -> Self {
665        Self {
666            name: name.to_string(),
667            name_offset,
668            data,
669            offset: 0,
670        }
671    }
672
673    pub fn name(&self) -> &str {
674        &self.name
675    }
676
677    pub fn size(&self) -> u64 {
678        let raw = self.data.len();
679        let padding = (8 - (raw % 8)) % 8;
680        (raw + padding) as u64
681    }
682
683    pub fn bytecode(&self) -> Vec<u8> {
684        let mut bytes = self.data.clone();
685        while !bytes.len().is_multiple_of(8) {
686            bytes.push(0);
687        }
688        bytes
689    }
690
691    pub fn set_offset(&mut self, offset: u64) {
692        self.offset = offset;
693    }
694
695    pub fn offset(&self) -> u64 {
696        self.offset
697    }
698
699    pub fn set_name_offset(&mut self, name_offset: u32) {
700        self.name_offset = name_offset;
701    }
702
703    pub fn name_offset(&self) -> u32 {
704        self.name_offset
705    }
706
707    pub fn section_header_bytecode(&self) -> Vec<u8> {
708        SectionHeader::new(
709            self.name_offset,
710            SectionHeader::SHT_PROGBITS,
711            0,
712            0,
713            self.offset,
714            self.data.len() as u64, // size without padding
715            0,
716            0,
717            1,
718            0,
719        )
720        .bytecode()
721    }
722}
723
724#[derive(Debug)]
725pub enum SectionType {
726    Code(CodeSection),
727    Data(DataSection),
728    ShStrTab(ShStrTabSection),
729    Dynamic(DynamicSection),
730    DynStr(DynStrSection),
731    DynSym(DynSymSection),
732    Default(NullSection),
733    RelDyn(RelDynSection),
734    DebugAbbrev(DebugSection),
735    DebugInfo(DebugSection),
736    DebugLine(DebugSection),
737    DebugLineStr(DebugSection),
738    DebugStr(DebugSection),
739    DebugFrame(DebugSection),
740    DebugLoc(DebugSection),
741    DebugRanges(DebugSection),
742}
743
744impl SectionType {
745    pub fn name(&self) -> &str {
746        match self {
747            SectionType::Code(cs) => &cs.name,
748            SectionType::Data(ds) => &ds.name,
749            SectionType::ShStrTab(ss) => &ss.name,
750            SectionType::Dynamic(ds) => &ds.name,
751            SectionType::DynStr(ds) => &ds.name,
752            SectionType::DynSym(ds) => &ds.name,
753            SectionType::Default(ds) => &ds.name,
754            SectionType::RelDyn(ds) => &ds.name,
755            SectionType::DebugAbbrev(ds) => ds.name(),
756            SectionType::DebugInfo(ds) => ds.name(),
757            SectionType::DebugLine(ds) => ds.name(),
758            SectionType::DebugLineStr(ds) => ds.name(),
759            SectionType::DebugStr(ds) => ds.name(),
760            SectionType::DebugFrame(ds) => ds.name(),
761            SectionType::DebugLoc(ds) => ds.name(),
762            SectionType::DebugRanges(ds) => ds.name(),
763        }
764    }
765
766    pub fn bytecode(&self) -> Vec<u8> {
767        match self {
768            SectionType::Code(cs) => cs.bytecode(),
769            SectionType::Data(ds) => ds.bytecode(),
770            SectionType::ShStrTab(ss) => ss.bytecode(),
771            SectionType::Dynamic(ds) => ds.bytecode(),
772            SectionType::DynStr(ds) => ds.bytecode(),
773            SectionType::DynSym(ds) => ds.bytecode(),
774            SectionType::Default(ds) => ds.bytecode(),
775            SectionType::RelDyn(ds) => ds.bytecode(),
776            SectionType::DebugAbbrev(ds) => ds.bytecode(),
777            SectionType::DebugInfo(ds) => ds.bytecode(),
778            SectionType::DebugLine(ds) => ds.bytecode(),
779            SectionType::DebugLineStr(ds) => ds.bytecode(),
780            SectionType::DebugStr(ds) => ds.bytecode(),
781            SectionType::DebugFrame(ds) => ds.bytecode(),
782            SectionType::DebugLoc(ds) => ds.bytecode(),
783            SectionType::DebugRanges(ds) => ds.bytecode(),
784        }
785    }
786
787    pub fn size(&self) -> u64 {
788        match self {
789            SectionType::Code(cs) => cs.size(),
790            SectionType::Data(ds) => ds.size(),
791            SectionType::ShStrTab(ss) => ss.size(),
792            SectionType::Dynamic(ds) => ds.size(),
793            SectionType::DynStr(ds) => ds.size(),
794            SectionType::DynSym(ds) => ds.size(),
795            SectionType::Default(ds) => ds.size(),
796            SectionType::RelDyn(ds) => ds.size(),
797            SectionType::DebugAbbrev(ds) => ds.size(),
798            SectionType::DebugInfo(ds) => ds.size(),
799            SectionType::DebugLine(ds) => ds.size(),
800            SectionType::DebugLineStr(ds) => ds.size(),
801            SectionType::DebugStr(ds) => ds.size(),
802            SectionType::DebugFrame(ds) => ds.size(),
803            SectionType::DebugLoc(ds) => ds.size(),
804            SectionType::DebugRanges(ds) => ds.size(),
805        }
806    }
807
808    pub fn section_header_bytecode(&self) -> Vec<u8> {
809        match self {
810            SectionType::Code(cs) => cs.section_header_bytecode(),
811            SectionType::Data(ds) => ds.section_header_bytecode(),
812            SectionType::ShStrTab(ss) => ss.section_header_bytecode(),
813            SectionType::Dynamic(ds) => ds.section_header_bytecode(),
814            SectionType::DynStr(ds) => ds.section_header_bytecode(),
815            SectionType::DynSym(ds) => ds.section_header_bytecode(),
816            SectionType::Default(ds) => ds.section_header_bytecode(),
817            SectionType::RelDyn(ds) => ds.section_header_bytecode(),
818            SectionType::DebugAbbrev(ds) => ds.section_header_bytecode(),
819            SectionType::DebugInfo(ds) => ds.section_header_bytecode(),
820            SectionType::DebugLine(ds) => ds.section_header_bytecode(),
821            SectionType::DebugLineStr(ds) => ds.section_header_bytecode(),
822            SectionType::DebugStr(ds) => ds.section_header_bytecode(),
823            SectionType::DebugFrame(ds) => ds.section_header_bytecode(),
824            SectionType::DebugLoc(ds) => ds.section_header_bytecode(),
825            SectionType::DebugRanges(ds) => ds.section_header_bytecode(),
826        }
827    }
828
829    pub fn set_offset(&mut self, offset: u64) {
830        match self {
831            SectionType::Code(cs) => cs.set_offset(offset),
832            SectionType::Data(ds) => ds.set_offset(offset),
833            SectionType::ShStrTab(ss) => ss.set_offset(offset),
834            SectionType::Dynamic(ds) => ds.set_offset(offset),
835            SectionType::DynStr(ds) => ds.set_offset(offset),
836            SectionType::DynSym(ds) => ds.set_offset(offset),
837            SectionType::RelDyn(ds) => ds.set_offset(offset),
838            SectionType::Default(_) => (), // NullSection doesn't need offset
839            SectionType::DebugAbbrev(ds) => ds.set_offset(offset),
840            SectionType::DebugInfo(ds) => ds.set_offset(offset),
841            SectionType::DebugLine(ds) => ds.set_offset(offset),
842            SectionType::DebugLineStr(ds) => ds.set_offset(offset),
843            SectionType::DebugStr(ds) => ds.set_offset(offset),
844            SectionType::DebugFrame(ds) => ds.set_offset(offset),
845            SectionType::DebugLoc(ds) => ds.set_offset(offset),
846            SectionType::DebugRanges(ds) => ds.set_offset(offset),
847        }
848    }
849
850    pub fn set_vaddr(&mut self, vaddr: u64) {
851        match self {
852            SectionType::Code(cs) => cs.set_vaddr(vaddr),
853            SectionType::Data(ds) => ds.set_vaddr(vaddr),
854            _ => (),
855        }
856    }
857
858    pub fn offset(&self) -> u64 {
859        match self {
860            SectionType::Code(cs) => cs.offset,
861            SectionType::Data(ds) => ds.offset,
862            SectionType::ShStrTab(ss) => ss.offset,
863            SectionType::Dynamic(ds) => ds.offset,
864            SectionType::DynStr(ds) => ds.offset,
865            SectionType::DynSym(ds) => ds.offset,
866            SectionType::Default(ns) => ns.offset,
867            SectionType::RelDyn(rs) => rs.offset,
868            SectionType::DebugAbbrev(ds) => ds.offset(),
869            SectionType::DebugInfo(ds) => ds.offset(),
870            SectionType::DebugLine(ds) => ds.offset(),
871            SectionType::DebugLineStr(ds) => ds.offset(),
872            SectionType::DebugStr(ds) => ds.offset(),
873            SectionType::DebugFrame(ds) => ds.offset(),
874            SectionType::DebugLoc(ds) => ds.offset(),
875            SectionType::DebugRanges(ds) => ds.offset(),
876        }
877    }
878}
879
880#[cfg(test)]
881mod tests {
882    use {
883        super::*,
884        sbpf_common::{instruction::Instruction, opcode::Opcode},
885    };
886
887    #[test]
888    fn test_code_section_new() {
889        let inst = Instruction {
890            opcode: Opcode::Exit,
891            dst: None,
892            src: None,
893            off: None,
894            imm: None,
895            span: 0..4,
896        };
897        let nodes = vec![ASTNode::Instruction {
898            instruction: inst,
899            offset: 0,
900        }];
901
902        let section = CodeSection::new(nodes, 8);
903        assert_eq!(section.name(), ".text");
904        assert_eq!(section.get_size(), 8);
905    }
906
907    #[test]
908    fn test_code_section_bytecode() {
909        let inst = Instruction {
910            opcode: Opcode::Exit,
911            dst: None,
912            src: None,
913            off: None,
914            imm: None,
915            span: 0..4,
916        };
917        let nodes = vec![ASTNode::Instruction {
918            instruction: inst,
919            offset: 0,
920        }];
921
922        let section = CodeSection::new(nodes, 8);
923        let bytes = section.bytecode();
924        assert_eq!(bytes.len(), 8);
925    }
926
927    #[test]
928    fn test_data_section_new() {
929        let rodata = ROData {
930            name: "msg".to_string(),
931            args: vec![
932                Token::Directive("ascii".to_string(), 0..5),
933                Token::StringLiteral("Hi".to_string(), 6..10),
934            ],
935            span: 0..10,
936        };
937        let nodes = vec![ASTNode::ROData { rodata, offset: 0 }];
938
939        let section = DataSection::new(nodes, 2);
940        assert_eq!(section.name(), ".rodata");
941        assert_eq!(section.get_size(), 2);
942    }
943
944    #[test]
945    fn test_data_section_rodata() {
946        let rodata = ROData {
947            name: "my_str".to_string(),
948            args: vec![
949                Token::Directive("ascii".to_string(), 0..5),
950                Token::StringLiteral("test".to_string(), 6..12),
951            ],
952            span: 0..12,
953        };
954        let nodes = vec![ASTNode::ROData { rodata, offset: 0 }];
955
956        let section = DataSection::new(nodes, 4);
957        let rodata = section.rodata();
958        assert_eq!(rodata.len(), 1);
959        assert_eq!(rodata[0].0, "my_str");
960        assert_eq!(rodata[0].2, "test");
961    }
962
963    #[test]
964    fn test_null_section() {
965        let section = NullSection::new();
966        assert_eq!(section.name(), ".unknown");
967        assert_eq!(section.bytecode().len(), 0);
968        assert_eq!(section.size(), 0);
969    }
970
971    #[test]
972    fn test_shstrtab_section() {
973        let names = vec![".text".to_string(), ".data".to_string()];
974        let mut section = ShStrTabSection::new(10, names);
975        section.set_offset(100);
976
977        assert_eq!(section.name(), ".s");
978        assert!(section.size() > 0);
979
980        let bytes = section.bytecode();
981        assert!(!bytes.is_empty());
982        assert_eq!(bytes[0], 0); // First byte is null
983    }
984
985    #[test]
986    fn test_dynamic_section_setters() {
987        let mut section = DynamicSection::new(5);
988        section.set_offset(100);
989        section.set_link(3);
990        section.set_rel_offset(200);
991        section.set_rel_size(48);
992        section.set_rel_count(2);
993        section.set_dynsym_offset(300);
994        section.set_dynstr_offset(400);
995        section.set_dynstr_size(50);
996
997        assert_eq!(section.name(), ".dynamic");
998        assert!(section.size() > 0);
999    }
1000
1001    #[test]
1002    fn test_dynamic_section_bytecode_with_rel_count() {
1003        let mut section = DynamicSection::new(0);
1004        section.set_rel_count(3);
1005        assert_eq!(section.size(), 11 * 16); // With rel_count
1006    }
1007
1008    #[test]
1009    fn test_dynamic_section_bytecode_without_rel_count() {
1010        let section = DynamicSection::new(0);
1011        assert_eq!(section.size(), 10 * 16); // Without rel_count
1012    }
1013
1014    #[test]
1015    fn test_dynstr_section() {
1016        let names = vec!["entrypoint".to_string(), "function".to_string()];
1017        let mut section = DynStrSection::new(8, names);
1018        section.set_offset(200);
1019
1020        assert_eq!(section.name(), ".dynstr");
1021
1022        let bytes = section.bytecode();
1023        assert_eq!(bytes[0], 0);
1024        assert_eq!(bytes.len() % 8, 0);
1025    }
1026
1027    #[test]
1028    fn test_dynsym_section() {
1029        let symbols = vec![
1030            DynamicSymbol::new(0, 0, 0, 0, 0, 0),
1031            DynamicSymbol::new(1, 0x12, 0, 1, 0x120, 48),
1032        ];
1033        let mut section = DynSymSection::new(15, symbols);
1034        section.set_offset(300);
1035        section.set_link(4);
1036
1037        assert_eq!(section.name(), ".dynsym");
1038        assert_eq!(section.size(), 2 * 24); // 2 symbols * 24 bytes each
1039    }
1040
1041    #[test]
1042    fn test_rel_dyn_section() {
1043        let entries = vec![RelDyn::new(0x120, 0x08, 0), RelDyn::new(0x200, 0x0a, 1)];
1044        let mut section = RelDynSection::new(18, entries);
1045        section.set_offset(400);
1046        section.set_link(3);
1047
1048        assert_eq!(section.name(), ".rel.dyn");
1049        assert_eq!(section.size(), 2 * 16); // 2 entries * 16 bytes each
1050    }
1051
1052    #[test]
1053    fn test_section_type_enum_code() {
1054        let inst = Instruction {
1055            opcode: Opcode::Exit,
1056            dst: None,
1057            src: None,
1058            off: None,
1059            imm: None,
1060            span: 0..4,
1061        };
1062        let code_section = CodeSection::new(
1063            vec![ASTNode::Instruction {
1064                instruction: inst,
1065                offset: 0,
1066            }],
1067            8,
1068        );
1069
1070        let section_type = SectionType::Code(code_section);
1071        assert_eq!(section_type.name(), ".text");
1072        assert_eq!(section_type.size(), 8);
1073    }
1074
1075    #[test]
1076    fn test_section_type_set_offset() {
1077        let mut section = SectionType::Default(NullSection::new());
1078        section.set_offset(100);
1079
1080        let mut dyn_section = SectionType::Dynamic(DynamicSection::new(0));
1081        dyn_section.set_offset(200);
1082        assert_eq!(dyn_section.offset(), 200);
1083    }
1084
1085    #[test]
1086    fn test_debug_section_types() {
1087        let sections = [
1088            SectionType::DebugAbbrev(DebugSection::new(".debug_abbrev", 0, vec![1, 2, 3])),
1089            SectionType::DebugInfo(DebugSection::new(".debug_info", 0, vec![4, 5, 6])),
1090            SectionType::DebugLine(DebugSection::new(".debug_line", 0, vec![7, 8, 9])),
1091            SectionType::DebugLineStr(DebugSection::new(".debug_line_str", 0, vec![10, 11, 12])),
1092        ];
1093
1094        let expected_names = [
1095            ".debug_abbrev",
1096            ".debug_info",
1097            ".debug_line",
1098            ".debug_line_str",
1099        ];
1100        for (i, section) in sections.iter().enumerate() {
1101            assert_eq!(section.name(), expected_names[i]);
1102            assert_eq!(section.size(), 8);
1103            assert_eq!(section.bytecode().len(), 8);
1104            assert!(!section.section_header_bytecode().is_empty());
1105        }
1106    }
1107}