solana_optimizer/
optimizer.rs

1use solana_rbpf::ebpf;
2use solana_rbpf::elf::Executable;
3use solana_rbpf::program::{BuiltinProgram, SBPFVersion, FunctionRegistry};
4use solana_rbpf::vm::Config;
5use std::sync::Arc;
6use elf::ElfBytes;
7use elf::endian::{AnyEndian, EndianParse}; // 导入 EndianParse
8use elf::file::Class;
9use std::fs;
10
11#[derive(Debug, serde::Serialize)]
12pub struct Issue {
13    kind: String,
14    offset: usize,
15    desc: String,
16}
17
18pub struct Optimizer {
19    insns: Vec<solana_rbpf::ebpf::Insn>,
20    issues: Vec<Issue>,
21    elf_bytes: Vec<u8>,
22    text_section_idx: usize,
23}
24
25impl Optimizer {
26    pub fn new(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
27        let elf_bytes = fs::read(path)?;
28        let elf = ElfBytes::<AnyEndian>::minimal_parse(&elf_bytes)?;
29        let (shdrs_opt, strtab_opt) = elf.section_headers_with_strtab()?;
30        let shdrs = shdrs_opt.ok_or("No section headers")?;
31        let strtab = strtab_opt.ok_or("No string table")?;
32        let text_section_idx = shdrs
33            .iter()
34            .position(|sh| strtab.get(sh.sh_name as usize).ok() == Some(".text"))
35            .ok_or("No .text section")?;
36        // 修改第 38、39 行：将 Option 转为 Result
37        let text_section = shdrs.get(text_section_idx).map_err(|_| "Invalid text section index")?;
38        let text_bytes = elf.section_data(&text_section)?.0;
39        let insns = Self::disassemble_text_bytes(text_bytes)?;
40        Ok(Self { insns, issues: Vec::new(), elf_bytes, text_section_idx })
41    }
42
43    fn disassemble_text_bytes(bytes: &[u8]) -> Result<Vec<solana_rbpf::ebpf::Insn>, Box<dyn std::error::Error>> {
44        let mut insns = Vec::new();
45        let mut offset = 0;
46        while offset + 8 <= bytes.len() {
47            let chunk = &bytes[offset..offset + 8];
48            insns.push(ebpf::Insn {
49                ptr: 0,
50                opc: chunk[0],
51                dst: chunk[1] & 0x0F,
52                src: (chunk[1] >> 4) & 0x0F,
53                off: i16::from_le_bytes([chunk[2], chunk[3]]),
54                imm: i64::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7], 0, 0, 0, 0]),
55            });
56            offset += 8;
57        }
58        Ok(insns)
59    }
60
61    pub fn remove_logs(&mut self) {
62        let original_len = self.insns.len();
63        self.insns.retain(|insn| {
64            if insn.opc == 0x91 { // 假设 sol_log opcode
65                self.issues.push(Issue {
66                    kind: "LogRemoved".to_string(),
67                    offset: insn.off as usize,
68                    desc: "Removed redundant sol_log call".to_string(),
69                });
70                false
71            } else {
72                true
73            }
74        });
75        println!("Removed {} log instructions", original_len - self.insns.len());
76    }
77
78    pub fn merge_loads(&mut self) {
79        let mut i = 0;
80        while i < self.insns.len() - 1 {
81            if self.insns[i].opc == ebpf::LD_DW_IMM && self.insns[i + 1].opc == ebpf::LD_DW_IMM {
82                if self.insns[i].src == self.insns[i + 1].src {
83                    self.insns.remove(i + 1);
84                    self.issues.push(Issue {
85                        kind: "LoadMerged".to_string(),
86                        offset: i,
87                        desc: "Merged duplicate load instruction".to_string(),
88                    });
89                    continue;
90                }
91            }
92            i += 1;
93        }
94        println!("Merged duplicate load instructions");
95    }
96
97    pub fn check_size(&mut self) {
98        let size = self.insns.len() * 8;
99        if size > 128 * 1024 {
100            self.issues.push(Issue {
101                kind: "SizeExceeded".to_string(),
102                offset: 0,
103                desc: format!("Program size {} bytes exceeds 128KB", size),
104            });
105        }
106    }
107
108    pub fn generate(&self) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
109        use solana_rbpf::vm::TestContextObject;
110
111        let loader = Arc::new(BuiltinProgram::<TestContextObject>::new_loader(
112            Config::default(),
113            FunctionRegistry::default(),
114        ));
115
116        let executable = Executable::from_text_bytes(
117            &self.insns.iter().flat_map(|insn| {
118                let imm_bytes = insn.imm.to_le_bytes();
119                [
120                    insn.opc,
121                    (insn.dst & 0x0F) | ((insn.src & 0x0F) << 4),
122                    insn.off.to_le_bytes()[0],
123                    insn.off.to_le_bytes()[1],
124                    imm_bytes[0],
125                    imm_bytes[1],
126                    imm_bytes[2],
127                    imm_bytes[3],
128                ]
129            }).collect::<Vec<u8>>(),
130            loader,
131            SBPFVersion::V2,
132            FunctionRegistry::default(),
133        )?;
134        let optimized_text = executable.get_text_bytes().1.to_vec();
135        let elf = ElfBytes::<AnyEndian>::minimal_parse(&self.elf_bytes)?;
136        let ehdr = elf.ehdr;
137        let (shdrs_opt, _) = elf.section_headers_with_strtab()?;
138        let shdrs = shdrs_opt.ok_or("No section headers")?;
139        let mut new_shdrs: Vec<_> = shdrs.iter().collect();
140        let mut elf_bytes = Vec::new();
141
142        // 写入 ELF 头部
143        let class_val = match ehdr.class {
144            Class::ELF32 => elf::abi::ELFCLASS32,
145            Class::ELF64 => elf::abi::ELFCLASS64,
146            _ => elf::abi::ELFCLASS32,
147        };
148        // 修改第 148 行：使用 EndianParse 的 is_little 方法
149        let endian_val = if ehdr.endianness.is_little() {
150            elf::abi::ELFDATA2LSB
151        } else {
152            elf::abi::ELFDATA2MSB
153        };
154        let mut ehdr_bytes = Vec::new();
155        ehdr_bytes.extend_from_slice(&[0x7f, b'E', b'L', b'F']); // ei_magic，4 bytes
156        ehdr_bytes.extend_from_slice(&[class_val, endian_val, ehdr.version.try_into()?, ehdr.osabi]); // 4 bytes
157        ehdr_bytes.extend_from_slice(&[ehdr.abiversion, 0, 0, 0, 0, 0, 0, 0]); // EI_PAD，8 bytes
158        ehdr_bytes.extend_from_slice(&ehdr.e_type.to_le_bytes()); // 2 bytes
159        ehdr_bytes.extend_from_slice(&ehdr.e_machine.to_le_bytes()); // 2 bytes
160        ehdr_bytes.extend_from_slice(&ehdr.version.to_le_bytes()); // 4 bytes
161        ehdr_bytes.extend_from_slice(&(ehdr.e_entry as u32).to_le_bytes()); // 4 bytes
162        ehdr_bytes.extend_from_slice(&(ehdr.e_phoff as u32).to_le_bytes()); // 4 bytes
163        ehdr_bytes.extend_from_slice(&(ehdr.e_shoff as u32).to_le_bytes()); // 4 bytes
164        ehdr_bytes.extend_from_slice(&ehdr.e_flags.to_le_bytes()); // 4 bytes
165        ehdr_bytes.extend_from_slice(&ehdr.e_ehsize.to_le_bytes()); // 2 bytes
166        ehdr_bytes.extend_from_slice(&ehdr.e_phentsize.to_le_bytes()); // 2 bytes
167        ehdr_bytes.extend_from_slice(&ehdr.e_phnum.to_le_bytes()); // 2 bytes
168        ehdr_bytes.extend_from_slice(&ehdr.e_shentsize.to_le_bytes()); // 2 bytes
169        ehdr_bytes.extend_from_slice(&ehdr.e_shnum.to_le_bytes()); // 2 bytes
170        ehdr_bytes.extend_from_slice(&ehdr.e_shstrndx.to_le_bytes()); // 2 bytes
171        elf_bytes.extend_from_slice(&ehdr_bytes);
172
173        // 计算段数据起始偏移
174        let mut offset = ehdr.e_ehsize as usize;
175        let mut section_data = Vec::new();
176
177        // 修改第 185、186 行：使用 iter_mut 避免借用冲突
178        for (i, sh) in new_shdrs.iter_mut().enumerate() {
179            let data = if i == self.text_section_idx {
180                optimized_text.clone()
181            } else {
182                elf.section_data(sh)?.0.to_vec()
183            };
184            section_data.push((sh.sh_offset, data.clone()));
185            sh.sh_offset = offset as u64; // 直接修改可变引用
186            sh.sh_size = data.len() as u64; // 直接修改可变引用
187            offset += data.len();
188            offset = (offset + 7) & !7; // 8 字节对齐
189        }
190
191        // 更新段表偏移
192        ehdr_bytes[32..36].copy_from_slice(&(offset as u32).to_le_bytes()); // e_shoff，4 bytes
193        elf_bytes[0..52].copy_from_slice(&ehdr_bytes); // 固定为 52 字节
194
195        // 写入段数据
196        for (_, data) in section_data.iter() {
197            elf_bytes.extend_from_slice(data);
198            let padding = (8 - (data.len() % 8)) % 8;
199            elf_bytes.extend_from_slice(&vec![0; padding]);
200        }
201
202        // 写入段表
203        for sh in new_shdrs.iter() {
204            let sh_bytes = [
205                sh.sh_name.to_le_bytes(), // 4 bytes
206                sh.sh_type.to_le_bytes(), // 4 bytes
207                (sh.sh_flags as u32).to_le_bytes(), // 4 bytes
208                (sh.sh_addr as u32).to_le_bytes(), // 4 bytes
209                (sh.sh_offset as u32).to_le_bytes(), // 4 bytes
210                (sh.sh_size as u32).to_le_bytes(), // 4 bytes
211                sh.sh_link.to_le_bytes(), // 4 bytes
212                sh.sh_info.to_le_bytes(), // 4 bytes
213                (sh.sh_addralign as u32).to_le_bytes(), // 4 bytes
214                (sh.sh_entsize as u32).to_le_bytes(), // 4 bytes
215            ].concat();
216            elf_bytes.extend_from_slice(&sh_bytes);
217        }
218
219        Ok(elf_bytes)
220    }
221
222    pub fn report(&self) -> String {
223        serde_json::to_string_pretty(&self.issues).unwrap_or("[]".to_string())
224    }
225}
226
227pub fn optimize_sbf(input_path: &str, output_path: &str) -> Result<String, Box<dyn std::error::Error>> {
228    let mut optimizer = Optimizer::new(input_path)?;
229    optimizer.remove_logs();
230    optimizer.merge_loads();
231    optimizer.check_size();
232    let optimized_bytes = optimizer.generate()?;
233    fs::write(output_path, optimized_bytes)?;
234    Ok(optimizer.report())
235}
solana_optimizer/optimizer.rs

solana_optimizer/
optimizer.rs
