elf_assembler/assembler/x64/
code_builder.rs

1//! X64 代码构建器
2//!
3//! 提供生成 x64 汇编指令的高级接口
4
5use super::context::{RelocationType, X64Context};
6
7/// X64 代码构建器
8///
9/// 现在使用 context 来管理代码生成状态
10#[derive(Debug)]
11pub struct X64CodeBuilder {
12    context: X64Context,
13}
14
15impl X64CodeBuilder {
16    /// 创建新的代码构建器
17    pub fn new() -> Self {
18        Self { context: X64Context::new() }
19    }
20
21    /// 获取上下文的可变引用
22    pub fn context_mut(&mut self) -> &mut X64Context {
23        &mut self.context
24    }
25
26    /// 获取上下文的不可变引用
27    pub fn context(&self) -> &X64Context {
28        &self.context
29    }
30
31    /// 获取生成的代码
32    pub fn code(&self) -> &[u8] {
33        &self.context.code
34    }
35
36    /// 完成代码生成并返回上下文
37    pub fn finish(self) -> X64Context {
38        self.context
39    }
40
41    /// 生成函数序言
42    pub fn function_prologue(&mut self) -> &mut Self {
43        let bytes = vec![
44            0x55, // push rbp
45            0x48, 0x89, 0xE5, // mov rbp, rsp
46            0x48, 0x83, 0xEC, 0x20, // sub rsp, 32 (shadow space)
47        ];
48        self.context.emit_bytes(&bytes);
49        self
50    }
51
52    /// 生成函数尾声
53    pub fn function_epilogue(&mut self) -> &mut Self {
54        let bytes = vec![
55            0x48, 0x89, 0xEC, // mov rsp, rbp
56            0x5D, // pop rbp
57            0xC3, // ret
58        ];
59        self.context.emit_bytes(&bytes);
60        self
61    }
62
63    /// 生成退出程序代码
64    pub fn exit_program(&mut self, exit_code: i32) -> &mut Self {
65        // mov ecx, exit_code (第一个参数)
66        self.context.emit_bytes(&[0xB9]);
67        self.context.emit_bytes(&exit_code.to_le_bytes());
68
69        // call ExitProcess - 添加函数调用信息
70        self.context.add_function_call("ExitProcess", true);
71        self.context.add_relocation(RelocationType::RipRel32, "ExitProcess");
72        self.context.emit_bytes(&[0xFF, 0x15, 0x00, 0x00, 0x00, 0x00]); // call [ExitProcess]
73
74        self
75    }
76
77    /// 实现 context 管理功能
78    pub fn load_immediate(&mut self, value: i64) {
79        // mov rax, value
80        if value >= i32::MIN as i64 && value <= i32::MAX as i64 {
81            self.context.emit_bytes(&[0xB8]);
82            self.context.emit_bytes(&(value as i32).to_le_bytes());
83        }
84        else {
85            self.context.emit_bytes(&[0x48, 0xB8]);
86            self.context.emit_bytes(&value.to_le_bytes());
87        }
88    }
89
90    pub fn load_string_address(&mut self, string_id: &str) {
91        // 添加字符串常量
92        self.context.add_string_constant(string_id);
93
94        // lea rax, [rip + string_offset]
95        self.context.emit_bytes(&[0x48, 0x8D, 0x05]);
96        let _offset = self.context.reference_label(&format!("str_{}", string_id));
97        self.context.emit_bytes(&[0x00, 0x00, 0x00, 0x00]); // 占位符
98    }
99
100    pub fn store_local(&mut self, _offset: i32) {
101        let stack_offset = self.context.allocate_stack(8);
102        // mov [rbp + stack_offset], rax
103        if stack_offset >= -128 && stack_offset <= 127 {
104            self.context.emit_bytes(&[0x48, 0x89, 0x45]);
105            self.context.emit_bytes(&[stack_offset as u8]);
106        }
107        else {
108            self.context.emit_bytes(&[0x48, 0x89, 0x85]);
109            self.context.emit_bytes(&stack_offset.to_le_bytes());
110        }
111    }
112
113    pub fn load_local(&mut self, offset: i32) {
114        // mov rax, [rbp + offset]
115        if offset >= -128 && offset <= 127 {
116            self.context.emit_bytes(&[0x48, 0x8B, 0x45]);
117            self.context.emit_bytes(&[offset as u8]);
118        }
119        else {
120            self.context.emit_bytes(&[0x48, 0x8B, 0x85]);
121            self.context.emit_bytes(&offset.to_le_bytes());
122        }
123    }
124
125    pub fn add_operation(&mut self) {
126        // pop rbx; pop rax; add rax, rbx; push rax
127        self.context.emit_bytes(&[
128            0x5B, // pop rbx
129            0x58, // pop rax
130            0x48, 0x01, 0xD8, // add rax, rbx
131            0x50, // push rax
132        ]);
133    }
134
135    pub fn sub_operation(&mut self) {
136        // pop rbx; pop rax; sub rax, rbx; push rax
137        self.context.emit_bytes(&[
138            0x5B, // pop rbx
139            0x58, // pop rax
140            0x48, 0x29, 0xD8, // sub rax, rbx
141            0x50, // push rax
142        ]);
143    }
144
145    pub fn mul_operation(&mut self) {
146        // pop rbx; pop rax; imul rax, rbx; push rax
147        self.context.emit_bytes(&[
148            0x5B, // pop rbx
149            0x58, // pop rax
150            0x48, 0x0F, 0xAF, 0xC3, // imul rax, rbx
151            0x50, // push rax
152        ]);
153    }
154
155    pub fn call_printf(&mut self) {
156        // 调用 printf 函数
157        self.context.add_function_call("printf", true);
158
159        // call [printf] - RIP相对寻址
160        self.context.emit_bytes(&[0xFF, 0x15, 0x00, 0x00, 0x00, 0x00]);
161    }
162
163    pub fn pop_stack(&mut self) {
164        // pop rax
165        self.context.emit_bytes(&[0x58]);
166    }
167
168    pub fn conditional_jump_false(&mut self, label: &str) {
169        // test rax, rax; jz label
170        self.context.emit_bytes(&[0x48, 0x85, 0xC0]); // test rax, rax
171        self.context.emit_bytes(&[0x0F, 0x84]); // jz
172        let _offset = self.context.reference_label(label);
173        self.context.emit_bytes(&[0x00, 0x00, 0x00, 0x00]); // 占位符
174    }
175
176    pub fn unconditional_jump(&mut self, label: &str) {
177        // jmp label
178        self.context.emit_bytes(&[0xE9]); // jmp
179        let _offset = self.context.reference_label(label);
180        self.context.emit_bytes(&[0x00, 0x00, 0x00, 0x00]); // 占位符
181    }
182
183    /// 生成 Hello World 程序
184    pub fn hello_world_program() -> Vec<u8> {
185        let mut builder = X64CodeBuilder::new();
186
187        // 函数序言
188        builder.function_prologue();
189
190        // 加载字符串地址到 RCX (第一个参数)
191        builder.load_string_address("Hello, World!\n");
192
193        // 调用 printf
194        builder.call_printf();
195
196        // 退出程序
197        builder.exit_program(0);
198
199        builder.finish().code
200    }
201
202    /// 生成消息框程序
203    pub fn message_box_program() -> Vec<u8> {
204        let mut builder = X64CodeBuilder::new();
205
206        // 函数序言
207        builder.function_prologue();
208
209        // MessageBoxA 参数设置
210        // mov r9d, 0 (uType = MB_OK)
211        builder.context_mut().emit_bytes(&[0x41, 0xB9, 0x00, 0x00, 0x00, 0x00]);
212
213        // mov r8, title_addr (lpCaption)
214        builder.context_mut().emit_bytes(&[0x49, 0xB8]);
215        builder.context_mut().emit_bytes(&0x2100u64.to_le_bytes()); // 假设标题在 0x2100
216
217        // mov rdx, message_addr (lpText)
218        builder.context_mut().emit_bytes(&[0x48, 0xBA]);
219        builder.context_mut().emit_bytes(&0x2000u64.to_le_bytes()); // 假设消息在 0x2000
220
221        // mov rcx, 0 (hWnd = NULL)
222        builder.context_mut().emit_bytes(&[0x48, 0xB9, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
223
224        // call MessageBoxA
225        builder.context_mut().emit_bytes(&[0xFF, 0x15, 0x00, 0x00, 0x00, 0x00]);
226
227        // 退出程序
228        builder.exit_program(0);
229
230        builder.finish().code
231    }
232}