Skip to main content

sbpf_assembler/
section.rs

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