xref: /DragonOS/kernel/crates/rbpf/src/jit.rs (revision 4dd4856f933be0b4624c7f7ffa9e3d0c8c218873)
1 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
2 // Derived from uBPF <https://github.com/iovisor/ubpf>
3 // Copyright 2015 Big Switch Networks, Inc
4 //      (uBPF: JIT algorithm, originally in C)
5 // Copyright 2016 6WIND S.A. <quentin.monnet@6wind.com>
6 //      (Translation to Rust, MetaBuff addition)
7 
8 use std::{
9     fmt::{Error as FormatterError, Formatter},
10     io::{Error, ErrorKind},
11     mem,
12     ops::{Index, IndexMut},
13 };
14 
15 use crate::{ebpf, HashMap};
16 
17 extern crate libc;
18 
19 type MachineCode = unsafe fn(*mut u8, usize, *mut u8, usize, usize, usize) -> u64;
20 
21 const PAGE_SIZE: usize = 4096;
22 // TODO: check how long the page must be to be sure to support an eBPF program of maximum possible
23 // length
24 const NUM_PAGES: usize = 1;
25 
26 // Special values for target_pc in struct Jump
27 const TARGET_OFFSET: isize = ebpf::PROG_MAX_INSNS as isize;
28 const TARGET_PC_EXIT: isize = TARGET_OFFSET + 1;
29 
30 #[derive(Copy, Clone)]
31 enum OperandSize {
32     S8 = 8,
33     S16 = 16,
34     S32 = 32,
35     S64 = 64,
36 }
37 
38 // Registers
39 const RAX: u8 = 0;
40 const RCX: u8 = 1;
41 const RDX: u8 = 2;
42 const RBX: u8 = 3;
43 const RSP: u8 = 4;
44 const RBP: u8 = 5;
45 const RSI: u8 = 6;
46 const RDI: u8 = 7;
47 const R8: u8 = 8;
48 const R9: u8 = 9;
49 const R10: u8 = 10;
50 const R11: u8 = 11;
51 //const R12: u8 = 12;
52 const R13: u8 = 13;
53 const R14: u8 = 14;
54 const R15: u8 = 15;
55 
56 const REGISTER_MAP_SIZE: usize = 11;
57 const REGISTER_MAP: [u8; REGISTER_MAP_SIZE] = [
58     RAX, // 0  return value
59     RDI, // 1  arg 1
60     RSI, // 2  arg 2
61     RDX, // 3  arg 3
62     R9,  // 4  arg 4
63     R8,  // 5  arg 5
64     RBX, // 6  callee-saved
65     R13, // 7  callee-saved
66     R14, // 8  callee-saved
67     R15, // 9  callee-saved
68     RBP, // 10 stack pointer
69          // R10 and R11 are used to compute store a constant pointer to mem and to compute offset for
70          // LD_ABS_* and LD_IND_* operations, so they are not mapped to any eBPF register.
71 ];
72 
73 // Return the x86 register for the given eBPF register
74 fn map_register(r: u8) -> u8 {
75     assert!(r < REGISTER_MAP_SIZE as u8);
76     REGISTER_MAP[(r % REGISTER_MAP_SIZE as u8) as usize]
77 }
78 
79 macro_rules! emit_bytes {
80     ( $mem:ident, $data:tt, $t:ty ) => {{
81         let size = mem::size_of::<$t>() as usize;
82         assert!($mem.offset + size <= $mem.contents.len());
83         unsafe {
84             let mut ptr = $mem.contents.as_ptr().add($mem.offset) as *mut $t;
85             ptr.write_unaligned($data);
86         }
87         $mem.offset += size;
88     }};
89 }
90 
91 #[derive(Debug)]
92 struct Jump {
93     offset_loc: usize,
94     target_pc: isize,
95 }
96 
97 #[derive(Debug)]
98 struct JitCompiler {
99     pc_locs: Vec<usize>,
100     special_targets: HashMap<isize, usize>,
101     jumps: Vec<Jump>,
102 }
103 
104 impl JitCompiler {
105     fn new() -> JitCompiler {
106         JitCompiler {
107             pc_locs: vec![],
108             jumps: vec![],
109             special_targets: HashMap::new(),
110         }
111     }
112 
113     fn emit1(&self, mem: &mut JitMemory, data: u8) {
114         emit_bytes!(mem, data, u8);
115     }
116 
117     fn emit2(&self, mem: &mut JitMemory, data: u16) {
118         emit_bytes!(mem, data, u16);
119     }
120 
121     fn emit4(&self, mem: &mut JitMemory, data: u32) {
122         emit_bytes!(mem, data, u32);
123     }
124 
125     fn emit8(&self, mem: &mut JitMemory, data: u64) {
126         emit_bytes!(mem, data, u64);
127     }
128 
129     fn emit_modrm(&self, mem: &mut JitMemory, modrm: u8, r: u8, m: u8) {
130         assert_eq!((modrm | 0xc0), 0xc0);
131         self.emit1(mem, (modrm & 0xc0) | ((r & 0b111) << 3) | (m & 0b111));
132     }
133 
134     fn emit_modrm_reg2reg(&self, mem: &mut JitMemory, r: u8, m: u8) {
135         self.emit_modrm(mem, 0xc0, r, m);
136     }
137 
138     fn emit_modrm_and_displacement(&self, mem: &mut JitMemory, r: u8, m: u8, d: i32) {
139         if d == 0 && (m & 0b111) != RBP {
140             self.emit_modrm(mem, 0x00, r, m);
141         } else if (-128..=127).contains(&d) {
142             self.emit_modrm(mem, 0x40, r, m);
143             self.emit1(mem, d as u8);
144         } else {
145             self.emit_modrm(mem, 0x80, r, m);
146             self.emit4(mem, d as u32);
147         }
148     }
149 
150     fn basix_rex_would_set_bits(&self, w: u8, src: u8, dst: u8) -> bool {
151         w != 0 || (src & 0b1000) != 0 || (dst & 0b1000) != 0
152     }
153 
154     fn emit_rex(&self, mem: &mut JitMemory, w: u8, r: u8, x: u8, b: u8) {
155         assert_eq!((w | 1), 1);
156         assert_eq!((r | 1), 1);
157         assert_eq!((x | 1), 1);
158         assert_eq!((b | 1), 1);
159         self.emit1(mem, 0x40 | (w << 3) | (r << 2) | (x << 1) | b);
160     }
161 
162     // Emits a REX prefix with the top bit of src and dst.
163     // Skipped if no bits would be set.
164     fn emit_basic_rex(&self, mem: &mut JitMemory, w: u8, src: u8, dst: u8) {
165         if self.basix_rex_would_set_bits(w, src, dst) {
166             let is_masked = |val, mask| match val & mask {
167                 0 => 0,
168                 _ => 1,
169             };
170             self.emit_rex(mem, w, is_masked(src, 8), 0, is_masked(dst, 8));
171         }
172     }
173 
174     fn emit_push(&self, mem: &mut JitMemory, r: u8) {
175         self.emit_basic_rex(mem, 0, 0, r);
176         self.emit1(mem, 0x50 | (r & 0b111));
177     }
178 
179     fn emit_pop(&self, mem: &mut JitMemory, r: u8) {
180         self.emit_basic_rex(mem, 0, 0, r);
181         self.emit1(mem, 0x58 | (r & 0b111));
182     }
183 
184     // REX prefix and ModRM byte
185     // We use the MR encoding when there is a choice
186     // 'src' is often used as an opcode extension
187     fn emit_alu32(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8) {
188         self.emit_basic_rex(mem, 0, src, dst);
189         self.emit1(mem, op);
190         self.emit_modrm_reg2reg(mem, src, dst);
191     }
192 
193     // REX prefix, ModRM byte, and 32-bit immediate
194     fn emit_alu32_imm32(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8, imm: i32) {
195         self.emit_alu32(mem, op, src, dst);
196         self.emit4(mem, imm as u32);
197     }
198 
199     // REX prefix, ModRM byte, and 8-bit immediate
200     fn emit_alu32_imm8(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8, imm: i8) {
201         self.emit_alu32(mem, op, src, dst);
202         self.emit1(mem, imm as u8);
203     }
204 
205     // REX.W prefix and ModRM byte
206     // We use the MR encoding when there is a choice
207     // 'src' is often used as an opcode extension
208     fn emit_alu64(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8) {
209         self.emit_basic_rex(mem, 1, src, dst);
210         self.emit1(mem, op);
211         self.emit_modrm_reg2reg(mem, src, dst);
212     }
213 
214     // REX.W prefix, ModRM byte, and 32-bit immediate
215     fn emit_alu64_imm32(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8, imm: i32) {
216         self.emit_alu64(mem, op, src, dst);
217         self.emit4(mem, imm as u32);
218     }
219 
220     // REX.W prefix, ModRM byte, and 8-bit immediate
221     fn emit_alu64_imm8(&self, mem: &mut JitMemory, op: u8, src: u8, dst: u8, imm: i8) {
222         self.emit_alu64(mem, op, src, dst);
223         self.emit1(mem, imm as u8);
224     }
225 
226     // Register to register mov
227     fn emit_mov(&self, mem: &mut JitMemory, src: u8, dst: u8) {
228         self.emit_alu64(mem, 0x89, src, dst);
229     }
230 
231     fn emit_cmp_imm32(&self, mem: &mut JitMemory, dst: u8, imm: i32) {
232         self.emit_alu64_imm32(mem, 0x81, 7, dst, imm);
233     }
234 
235     fn emit_cmp(&self, mem: &mut JitMemory, src: u8, dst: u8) {
236         self.emit_alu64(mem, 0x39, src, dst);
237     }
238 
239     fn emit_cmp32_imm32(&self, mem: &mut JitMemory, dst: u8, imm: i32) {
240         self.emit_alu32_imm32(mem, 0x81, 7, dst, imm);
241     }
242 
243     fn emit_cmp32(&self, mem: &mut JitMemory, src: u8, dst: u8) {
244         self.emit_alu32(mem, 0x39, src, dst);
245     }
246 
247     // Load [src + offset] into dst
248     fn emit_load(&self, mem: &mut JitMemory, size: OperandSize, src: u8, dst: u8, offset: i32) {
249         let data = match size {
250             OperandSize::S64 => 1,
251             _ => 0,
252         };
253         self.emit_basic_rex(mem, data, dst, src);
254 
255         match size {
256             OperandSize::S8 => {
257                 // movzx
258                 self.emit1(mem, 0x0f);
259                 self.emit1(mem, 0xb6);
260             }
261             OperandSize::S16 => {
262                 // movzx
263                 self.emit1(mem, 0x0f);
264                 self.emit1(mem, 0xb7);
265             }
266             OperandSize::S32 | OperandSize::S64 => {
267                 // mov
268                 self.emit1(mem, 0x8b);
269             }
270         }
271 
272         self.emit_modrm_and_displacement(mem, dst, src, offset);
273     }
274 
275     // Load sign-extended immediate into register
276     fn emit_load_imm(&self, mem: &mut JitMemory, dst: u8, imm: i64) {
277         if imm >= i32::MIN as i64 && imm <= i32::MAX as i64 {
278             self.emit_alu64_imm32(mem, 0xc7, 0, dst, imm as i32);
279         } else {
280             // movabs $imm,dst
281             self.emit_basic_rex(mem, 1, 0, dst);
282             self.emit1(mem, 0xb8 | (dst & 0b111));
283             self.emit8(mem, imm as u64);
284         }
285     }
286 
287     // Store register src to [dst + offset]
288     fn emit_store(&self, mem: &mut JitMemory, size: OperandSize, src: u8, dst: u8, offset: i32) {
289         match size {
290             OperandSize::S16 => self.emit1(mem, 0x66), // 16-bit override
291             _ => {}
292         };
293         let (is_s8, is_u64, rexw) = match size {
294             OperandSize::S8 => (true, false, 0),
295             OperandSize::S64 => (false, true, 1),
296             _ => (false, false, 0),
297         };
298         if is_u64 || (src & 0b1000) != 0 || (dst & 0b1000) != 0 || is_s8 {
299             let is_masked = |val, mask| match val & mask {
300                 0 => 0,
301                 _ => 1,
302             };
303             self.emit_rex(mem, rexw, is_masked(src, 8), 0, is_masked(dst, 8));
304         }
305         match size {
306             OperandSize::S8 => self.emit1(mem, 0x88),
307             _ => self.emit1(mem, 0x89),
308         };
309         self.emit_modrm_and_displacement(mem, src, dst, offset);
310     }
311 
312     // Store immediate to [dst + offset]
313     fn emit_store_imm32(
314         &self,
315         mem: &mut JitMemory,
316         size: OperandSize,
317         dst: u8,
318         offset: i32,
319         imm: i32,
320     ) {
321         match size {
322             OperandSize::S16 => self.emit1(mem, 0x66), // 16-bit override
323             _ => {}
324         };
325         match size {
326             OperandSize::S64 => self.emit_basic_rex(mem, 1, 0, dst),
327             _ => self.emit_basic_rex(mem, 0, 0, dst),
328         };
329         match size {
330             OperandSize::S8 => self.emit1(mem, 0xc6),
331             _ => self.emit1(mem, 0xc7),
332         };
333         self.emit_modrm_and_displacement(mem, 0, dst, offset);
334         match size {
335             OperandSize::S8 => self.emit1(mem, imm as u8),
336             OperandSize::S16 => self.emit2(mem, imm as u16),
337             _ => self.emit4(mem, imm as u32),
338         };
339     }
340 
341     fn emit_direct_jcc(&self, mem: &mut JitMemory, code: u8, offset: u32) {
342         self.emit1(mem, 0x0f);
343         self.emit1(mem, code);
344         emit_bytes!(mem, offset, u32);
345     }
346 
347     fn emit_call(&self, mem: &mut JitMemory, target: usize) {
348         // TODO use direct call when possible
349         self.emit_load_imm(mem, RAX, target as i64);
350         // callq *%rax
351         self.emit1(mem, 0xff);
352         self.emit1(mem, 0xd0);
353     }
354 
355     fn emit_jump_offset(&mut self, mem: &mut JitMemory, target_pc: isize) {
356         let jump = Jump {
357             offset_loc: mem.offset,
358             target_pc,
359         };
360         self.jumps.push(jump);
361         self.emit4(mem, 0);
362     }
363 
364     fn emit_jcc(&mut self, mem: &mut JitMemory, code: u8, target_pc: isize) {
365         self.emit1(mem, 0x0f);
366         self.emit1(mem, code);
367         self.emit_jump_offset(mem, target_pc);
368     }
369 
370     fn emit_jmp(&mut self, mem: &mut JitMemory, target_pc: isize) {
371         self.emit1(mem, 0xe9);
372         self.emit_jump_offset(mem, target_pc);
373     }
374 
375     fn set_anchor(&mut self, mem: &mut JitMemory, target: isize) {
376         self.special_targets.insert(target, mem.offset);
377     }
378 
379     fn emit_muldivmod(
380         &mut self,
381         mem: &mut JitMemory,
382         pc: u16,
383         opc: u8,
384         src: u8,
385         dst: u8,
386         imm: i32,
387     ) {
388         let mul = (opc & ebpf::BPF_ALU_OP_MASK) == (ebpf::MUL32_IMM & ebpf::BPF_ALU_OP_MASK);
389         let div = (opc & ebpf::BPF_ALU_OP_MASK) == (ebpf::DIV32_IMM & ebpf::BPF_ALU_OP_MASK);
390         let modrm = (opc & ebpf::BPF_ALU_OP_MASK) == (ebpf::MOD32_IMM & ebpf::BPF_ALU_OP_MASK);
391         let is64 = (opc & ebpf::BPF_CLS_MASK) == ebpf::BPF_ALU64;
392         let is_reg = (opc & ebpf::BPF_X) == ebpf::BPF_X;
393 
394         if (div || mul) && !is_reg && imm == 0 {
395             // Division by zero returns 0
396             // Set register to 0: xor with itself
397             self.emit_alu32(mem, 0x31, dst, dst);
398             return;
399         }
400         if modrm && !is_reg && imm == 0 {
401             // Modulo remainder of division by zero keeps destination register unchanged
402             return;
403         }
404         if (div || modrm) && is_reg {
405             self.emit_load_imm(mem, RCX, pc as i64);
406 
407             // test src,src
408             if is64 {
409                 self.emit_alu64(mem, 0x85, src, src);
410             } else {
411                 self.emit_alu32(mem, 0x85, src, src);
412             }
413 
414             if div {
415                 // No division by 0: skip next instructions
416                 // Jump offset: emit_alu32 adds 2 to 3 bytes, emit_jmp adds 5
417                 let offset = match self.basix_rex_would_set_bits(0, dst, dst) {
418                     true => 3 + 5,
419                     false => 2 + 5,
420                 };
421                 self.emit_direct_jcc(mem, 0x85, offset);
422                 // Division by 0: set dst to 0 then go to next instruction
423                 // Set register to 0: xor with itself
424                 self.emit_alu32(mem, 0x31, dst, dst);
425                 self.emit_jmp(mem, (pc + 1) as isize);
426             }
427             if modrm {
428                 // Modulo by zero: keep destination register unchanged
429                 self.emit_jcc(mem, 0x84, (pc + 1) as isize);
430             }
431         }
432 
433         if dst != RAX {
434             self.emit_push(mem, RAX);
435         }
436         if dst != RDX {
437             self.emit_push(mem, RDX);
438         }
439         if imm != 0 {
440             self.emit_load_imm(mem, RCX, imm as i64);
441         } else {
442             self.emit_mov(mem, src, RCX);
443         }
444 
445         self.emit_mov(mem, dst, RAX);
446 
447         if div || modrm {
448             // Set register to 0: xor %edx,%edx
449             self.emit_alu32(mem, 0x31, RDX, RDX);
450         }
451 
452         if is64 {
453             self.emit_rex(mem, 1, 0, 0, 0);
454         }
455 
456         // mul %ecx or div %ecx
457         self.emit_alu32(mem, 0xf7, if mul { 4 } else { 6 }, RCX);
458 
459         if dst != RDX {
460             if modrm {
461                 self.emit_mov(mem, RDX, dst);
462             }
463             self.emit_pop(mem, RDX);
464         }
465         if dst != RAX {
466             if div || mul {
467                 self.emit_mov(mem, RAX, dst);
468             }
469             self.emit_pop(mem, RAX);
470         }
471     }
472 
473     fn jit_compile(
474         &mut self,
475         mem: &mut JitMemory,
476         prog: &[u8],
477         use_mbuff: bool,
478         update_data_ptr: bool,
479         helpers: &HashMap<u32, ebpf::Helper>,
480     ) -> Result<(), Error> {
481         self.emit_push(mem, RBP);
482         self.emit_push(mem, RBX);
483         self.emit_push(mem, R13);
484         self.emit_push(mem, R14);
485         self.emit_push(mem, R15);
486 
487         // RDI: mbuff
488         // RSI: mbuff_len
489         // RDX: mem
490         // RCX: mem_len
491         // R8:  mem_offset
492         // R9:  mem_end_offset
493 
494         // Save mem pointer for use with LD_ABS_* and LD_IND_* instructions
495         self.emit_mov(mem, RDX, R10);
496 
497         match (use_mbuff, update_data_ptr) {
498             (false, _) => {
499                 // We do not use any mbuff. Move mem pointer into register 1.
500                 if map_register(1) != RDX {
501                     self.emit_mov(mem, RDX, map_register(1));
502                 }
503             }
504             (true, false) => {
505                 // We use a mbuff already pointing to mem and mem_end: move it to register 1.
506                 if map_register(1) != RDI {
507                     self.emit_mov(mem, RDI, map_register(1));
508                 }
509             }
510             (true, true) => {
511                 // We have a fixed (simulated) mbuff: update mem and mem_end offset values in it.
512                 // Store mem at mbuff + mem_offset. Trash R8.
513                 self.emit_alu64(mem, 0x01, RDI, R8); // add mbuff to mem_offset in R8
514                 self.emit_store(mem, OperandSize::S64, RDX, R8, 0); // set mem at mbuff + mem_offset
515                                                                     // Store mem_end at mbuff + mem_end_offset. Trash R9.
516                 self.emit_load(mem, OperandSize::S64, RDX, R8, 0); // load mem into R8
517                 self.emit_alu64(mem, 0x01, RCX, R8); // add mem_len to mem (= mem_end)
518                 self.emit_alu64(mem, 0x01, RDI, R9); // add mbuff to mem_end_offset
519                 self.emit_store(mem, OperandSize::S64, R8, R9, 0); // store mem_end
520 
521                 // Move rdi into register 1
522                 if map_register(1) != RDI {
523                     self.emit_mov(mem, RDI, map_register(1));
524                 }
525             }
526         }
527 
528         // Copy stack pointer to R10
529         self.emit_mov(mem, RSP, map_register(10));
530 
531         // Allocate stack space
532         self.emit_alu64_imm32(mem, 0x81, 5, RSP, ebpf::STACK_SIZE as i32);
533 
534         self.pc_locs = vec![0; prog.len() / ebpf::INSN_SIZE + 1];
535 
536         let mut insn_ptr: usize = 0;
537         while insn_ptr * ebpf::INSN_SIZE < prog.len() {
538             let insn = ebpf::get_insn(prog, insn_ptr);
539 
540             self.pc_locs[insn_ptr] = mem.offset;
541 
542             let dst = map_register(insn.dst);
543             let src = map_register(insn.src);
544             let target_pc = insn_ptr as isize + insn.off as isize + 1;
545 
546             match insn.opc {
547                 // BPF_LD class
548                 // R10 is a constant pointer to mem.
549                 ebpf::LD_ABS_B => self.emit_load(mem, OperandSize::S8, R10, RAX, insn.imm),
550                 ebpf::LD_ABS_H => self.emit_load(mem, OperandSize::S16, R10, RAX, insn.imm),
551                 ebpf::LD_ABS_W => self.emit_load(mem, OperandSize::S32, R10, RAX, insn.imm),
552                 ebpf::LD_ABS_DW => self.emit_load(mem, OperandSize::S64, R10, RAX, insn.imm),
553                 ebpf::LD_IND_B => {
554                     self.emit_mov(mem, R10, R11); // load mem into R11
555                     self.emit_alu64(mem, 0x01, src, R11); // add src to R11
556                     self.emit_load(mem, OperandSize::S8, R11, RAX, insn.imm); // ld R0, mem[src+imm]
557                 }
558                 ebpf::LD_IND_H => {
559                     self.emit_mov(mem, R10, R11); // load mem into R11
560                     self.emit_alu64(mem, 0x01, src, R11); // add src to R11
561                     self.emit_load(mem, OperandSize::S16, R11, RAX, insn.imm); // ld R0, mem[src+imm]
562                 }
563                 ebpf::LD_IND_W => {
564                     self.emit_mov(mem, R10, R11); // load mem into R11
565                     self.emit_alu64(mem, 0x01, src, R11); // add src to R11
566                     self.emit_load(mem, OperandSize::S32, R11, RAX, insn.imm); // ld R0, mem[src+imm]
567                 }
568                 ebpf::LD_IND_DW => {
569                     self.emit_mov(mem, R10, R11); // load mem into R11
570                     self.emit_alu64(mem, 0x01, src, R11); // add src to R11
571                     self.emit_load(mem, OperandSize::S64, R11, RAX, insn.imm); // ld R0, mem[src+imm]
572                 }
573 
574                 ebpf::LD_DW_IMM => {
575                     insn_ptr += 1;
576                     let second_part = ebpf::get_insn(prog, insn_ptr).imm as u64;
577                     let imm = (insn.imm as u32) as u64 | second_part.wrapping_shl(32);
578                     self.emit_load_imm(mem, dst, imm as i64);
579                 }
580 
581                 // BPF_LDX class
582                 ebpf::LD_B_REG => self.emit_load(mem, OperandSize::S8, src, dst, insn.off as i32),
583                 ebpf::LD_H_REG => self.emit_load(mem, OperandSize::S16, src, dst, insn.off as i32),
584                 ebpf::LD_W_REG => self.emit_load(mem, OperandSize::S32, src, dst, insn.off as i32),
585                 ebpf::LD_DW_REG => self.emit_load(mem, OperandSize::S64, src, dst, insn.off as i32),
586 
587                 // BPF_ST class
588                 ebpf::ST_B_IMM => {
589                     self.emit_store_imm32(mem, OperandSize::S8, dst, insn.off as i32, insn.imm)
590                 }
591                 ebpf::ST_H_IMM => {
592                     self.emit_store_imm32(mem, OperandSize::S16, dst, insn.off as i32, insn.imm)
593                 }
594                 ebpf::ST_W_IMM => {
595                     self.emit_store_imm32(mem, OperandSize::S32, dst, insn.off as i32, insn.imm)
596                 }
597                 ebpf::ST_DW_IMM => {
598                     self.emit_store_imm32(mem, OperandSize::S64, dst, insn.off as i32, insn.imm)
599                 }
600 
601                 // BPF_STX class
602                 ebpf::ST_B_REG => self.emit_store(mem, OperandSize::S8, src, dst, insn.off as i32),
603                 ebpf::ST_H_REG => self.emit_store(mem, OperandSize::S16, src, dst, insn.off as i32),
604                 ebpf::ST_W_REG => self.emit_store(mem, OperandSize::S32, src, dst, insn.off as i32),
605                 ebpf::ST_DW_REG => {
606                     self.emit_store(mem, OperandSize::S64, src, dst, insn.off as i32)
607                 }
608                 ebpf::ST_W_XADD => unimplemented!(),
609                 ebpf::ST_DW_XADD => unimplemented!(),
610 
611                 // BPF_ALU class
612                 ebpf::ADD32_IMM => self.emit_alu32_imm32(mem, 0x81, 0, dst, insn.imm),
613                 ebpf::ADD32_REG => self.emit_alu32(mem, 0x01, src, dst),
614                 ebpf::SUB32_IMM => self.emit_alu32_imm32(mem, 0x81, 5, dst, insn.imm),
615                 ebpf::SUB32_REG => self.emit_alu32(mem, 0x29, src, dst),
616                 ebpf::MUL32_IMM
617                 | ebpf::MUL32_REG
618                 | ebpf::DIV32_IMM
619                 | ebpf::DIV32_REG
620                 | ebpf::MOD32_IMM
621                 | ebpf::MOD32_REG => {
622                     self.emit_muldivmod(mem, insn_ptr as u16, insn.opc, src, dst, insn.imm)
623                 }
624                 ebpf::OR32_IMM => self.emit_alu32_imm32(mem, 0x81, 1, dst, insn.imm),
625                 ebpf::OR32_REG => self.emit_alu32(mem, 0x09, src, dst),
626                 ebpf::AND32_IMM => self.emit_alu32_imm32(mem, 0x81, 4, dst, insn.imm),
627                 ebpf::AND32_REG => self.emit_alu32(mem, 0x21, src, dst),
628                 ebpf::LSH32_IMM => self.emit_alu32_imm8(mem, 0xc1, 4, dst, insn.imm as i8),
629                 ebpf::LSH32_REG => {
630                     self.emit_mov(mem, src, RCX);
631                     self.emit_alu32(mem, 0xd3, 4, dst);
632                 }
633                 ebpf::RSH32_IMM => self.emit_alu32_imm8(mem, 0xc1, 5, dst, insn.imm as i8),
634                 ebpf::RSH32_REG => {
635                     self.emit_mov(mem, src, RCX);
636                     self.emit_alu32(mem, 0xd3, 5, dst);
637                 }
638                 ebpf::NEG32 => self.emit_alu32(mem, 0xf7, 3, dst),
639                 ebpf::XOR32_IMM => self.emit_alu32_imm32(mem, 0x81, 6, dst, insn.imm),
640                 ebpf::XOR32_REG => self.emit_alu32(mem, 0x31, src, dst),
641                 ebpf::MOV32_IMM => self.emit_alu32_imm32(mem, 0xc7, 0, dst, insn.imm),
642                 ebpf::MOV32_REG => self.emit_mov(mem, src, dst),
643                 ebpf::ARSH32_IMM => self.emit_alu32_imm8(mem, 0xc1, 7, dst, insn.imm as i8),
644                 ebpf::ARSH32_REG => {
645                     self.emit_mov(mem, src, RCX);
646                     self.emit_alu32(mem, 0xd3, 7, dst);
647                 }
648                 ebpf::LE => {} // No-op
649                 ebpf::BE => {
650                     match insn.imm {
651                         16 => {
652                             // rol
653                             self.emit1(mem, 0x66); // 16-bit override
654                             self.emit_alu32_imm8(mem, 0xc1, 0, dst, 8);
655                             // and
656                             self.emit_alu32_imm32(mem, 0x81, 4, dst, 0xffff);
657                         }
658                         32 | 64 => {
659                             // bswap
660                             let bit = match insn.imm {
661                                 64 => 1,
662                                 _ => 0,
663                             };
664                             self.emit_basic_rex(mem, bit, 0, dst);
665                             self.emit1(mem, 0x0f);
666                             self.emit1(mem, 0xc8 | (dst & 0b111));
667                         }
668                         _ => unreachable!(), // Should have been caught by verifier
669                     }
670                 }
671 
672                 // BPF_ALU64 class
673                 ebpf::ADD64_IMM => self.emit_alu64_imm32(mem, 0x81, 0, dst, insn.imm),
674                 ebpf::ADD64_REG => self.emit_alu64(mem, 0x01, src, dst),
675                 ebpf::SUB64_IMM => self.emit_alu64_imm32(mem, 0x81, 5, dst, insn.imm),
676                 ebpf::SUB64_REG => self.emit_alu64(mem, 0x29, src, dst),
677                 ebpf::MUL64_IMM
678                 | ebpf::MUL64_REG
679                 | ebpf::DIV64_IMM
680                 | ebpf::DIV64_REG
681                 | ebpf::MOD64_IMM
682                 | ebpf::MOD64_REG => {
683                     self.emit_muldivmod(mem, insn_ptr as u16, insn.opc, src, dst, insn.imm)
684                 }
685                 ebpf::OR64_IMM => self.emit_alu64_imm32(mem, 0x81, 1, dst, insn.imm),
686                 ebpf::OR64_REG => self.emit_alu64(mem, 0x09, src, dst),
687                 ebpf::AND64_IMM => self.emit_alu64_imm32(mem, 0x81, 4, dst, insn.imm),
688                 ebpf::AND64_REG => self.emit_alu64(mem, 0x21, src, dst),
689                 ebpf::LSH64_IMM => self.emit_alu64_imm8(mem, 0xc1, 4, dst, insn.imm as i8),
690                 ebpf::LSH64_REG => {
691                     self.emit_mov(mem, src, RCX);
692                     self.emit_alu64(mem, 0xd3, 4, dst);
693                 }
694                 ebpf::RSH64_IMM => self.emit_alu64_imm8(mem, 0xc1, 5, dst, insn.imm as i8),
695                 ebpf::RSH64_REG => {
696                     self.emit_mov(mem, src, RCX);
697                     self.emit_alu64(mem, 0xd3, 5, dst);
698                 }
699                 ebpf::NEG64 => self.emit_alu64(mem, 0xf7, 3, dst),
700                 ebpf::XOR64_IMM => self.emit_alu64_imm32(mem, 0x81, 6, dst, insn.imm),
701                 ebpf::XOR64_REG => self.emit_alu64(mem, 0x31, src, dst),
702                 ebpf::MOV64_IMM => self.emit_load_imm(mem, dst, insn.imm as i64),
703                 ebpf::MOV64_REG => self.emit_mov(mem, src, dst),
704                 ebpf::ARSH64_IMM => self.emit_alu64_imm8(mem, 0xc1, 7, dst, insn.imm as i8),
705                 ebpf::ARSH64_REG => {
706                     self.emit_mov(mem, src, RCX);
707                     self.emit_alu64(mem, 0xd3, 7, dst);
708                 }
709 
710                 // BPF_JMP class
711                 ebpf::JA => self.emit_jmp(mem, target_pc),
712                 ebpf::JEQ_IMM => {
713                     self.emit_cmp_imm32(mem, dst, insn.imm);
714                     self.emit_jcc(mem, 0x84, target_pc);
715                 }
716                 ebpf::JEQ_REG => {
717                     self.emit_cmp(mem, src, dst);
718                     self.emit_jcc(mem, 0x84, target_pc);
719                 }
720                 ebpf::JGT_IMM => {
721                     self.emit_cmp_imm32(mem, dst, insn.imm);
722                     self.emit_jcc(mem, 0x87, target_pc);
723                 }
724                 ebpf::JGT_REG => {
725                     self.emit_cmp(mem, src, dst);
726                     self.emit_jcc(mem, 0x87, target_pc);
727                 }
728                 ebpf::JGE_IMM => {
729                     self.emit_cmp_imm32(mem, dst, insn.imm);
730                     self.emit_jcc(mem, 0x83, target_pc);
731                 }
732                 ebpf::JGE_REG => {
733                     self.emit_cmp(mem, src, dst);
734                     self.emit_jcc(mem, 0x83, target_pc);
735                 }
736                 ebpf::JLT_IMM => {
737                     self.emit_cmp_imm32(mem, dst, insn.imm);
738                     self.emit_jcc(mem, 0x82, target_pc);
739                 }
740                 ebpf::JLT_REG => {
741                     self.emit_cmp(mem, src, dst);
742                     self.emit_jcc(mem, 0x82, target_pc);
743                 }
744                 ebpf::JLE_IMM => {
745                     self.emit_cmp_imm32(mem, dst, insn.imm);
746                     self.emit_jcc(mem, 0x86, target_pc);
747                 }
748                 ebpf::JLE_REG => {
749                     self.emit_cmp(mem, src, dst);
750                     self.emit_jcc(mem, 0x86, target_pc);
751                 }
752                 ebpf::JSET_IMM => {
753                     self.emit_alu64_imm32(mem, 0xf7, 0, dst, insn.imm);
754                     self.emit_jcc(mem, 0x85, target_pc);
755                 }
756                 ebpf::JSET_REG => {
757                     self.emit_alu64(mem, 0x85, src, dst);
758                     self.emit_jcc(mem, 0x85, target_pc);
759                 }
760                 ebpf::JNE_IMM => {
761                     self.emit_cmp_imm32(mem, dst, insn.imm);
762                     self.emit_jcc(mem, 0x85, target_pc);
763                 }
764                 ebpf::JNE_REG => {
765                     self.emit_cmp(mem, src, dst);
766                     self.emit_jcc(mem, 0x85, target_pc);
767                 }
768                 ebpf::JSGT_IMM => {
769                     self.emit_cmp_imm32(mem, dst, insn.imm);
770                     self.emit_jcc(mem, 0x8f, target_pc);
771                 }
772                 ebpf::JSGT_REG => {
773                     self.emit_cmp(mem, src, dst);
774                     self.emit_jcc(mem, 0x8f, target_pc);
775                 }
776                 ebpf::JSGE_IMM => {
777                     self.emit_cmp_imm32(mem, dst, insn.imm);
778                     self.emit_jcc(mem, 0x8d, target_pc);
779                 }
780                 ebpf::JSGE_REG => {
781                     self.emit_cmp(mem, src, dst);
782                     self.emit_jcc(mem, 0x8d, target_pc);
783                 }
784                 ebpf::JSLT_IMM => {
785                     self.emit_cmp_imm32(mem, dst, insn.imm);
786                     self.emit_jcc(mem, 0x8c, target_pc);
787                 }
788                 ebpf::JSLT_REG => {
789                     self.emit_cmp(mem, src, dst);
790                     self.emit_jcc(mem, 0x8c, target_pc);
791                 }
792                 ebpf::JSLE_IMM => {
793                     self.emit_cmp_imm32(mem, dst, insn.imm);
794                     self.emit_jcc(mem, 0x8e, target_pc);
795                 }
796                 ebpf::JSLE_REG => {
797                     self.emit_cmp(mem, src, dst);
798                     self.emit_jcc(mem, 0x8e, target_pc);
799                 }
800 
801                 // BPF_JMP32 class
802                 ebpf::JEQ_IMM32 => {
803                     self.emit_cmp32_imm32(mem, dst, insn.imm);
804                     self.emit_jcc(mem, 0x84, target_pc);
805                 }
806                 ebpf::JEQ_REG32 => {
807                     self.emit_cmp32(mem, src, dst);
808                     self.emit_jcc(mem, 0x84, target_pc);
809                 }
810                 ebpf::JGT_IMM32 => {
811                     self.emit_cmp32_imm32(mem, dst, insn.imm);
812                     self.emit_jcc(mem, 0x87, target_pc);
813                 }
814                 ebpf::JGT_REG32 => {
815                     self.emit_cmp32(mem, src, dst);
816                     self.emit_jcc(mem, 0x87, target_pc);
817                 }
818                 ebpf::JGE_IMM32 => {
819                     self.emit_cmp32_imm32(mem, dst, insn.imm);
820                     self.emit_jcc(mem, 0x83, target_pc);
821                 }
822                 ebpf::JGE_REG32 => {
823                     self.emit_cmp32(mem, src, dst);
824                     self.emit_jcc(mem, 0x83, target_pc);
825                 }
826                 ebpf::JLT_IMM32 => {
827                     self.emit_cmp32_imm32(mem, dst, insn.imm);
828                     self.emit_jcc(mem, 0x82, target_pc);
829                 }
830                 ebpf::JLT_REG32 => {
831                     self.emit_cmp32(mem, src, dst);
832                     self.emit_jcc(mem, 0x82, target_pc);
833                 }
834                 ebpf::JLE_IMM32 => {
835                     self.emit_cmp32_imm32(mem, dst, insn.imm);
836                     self.emit_jcc(mem, 0x86, target_pc);
837                 }
838                 ebpf::JLE_REG32 => {
839                     self.emit_cmp32(mem, src, dst);
840                     self.emit_jcc(mem, 0x86, target_pc);
841                 }
842                 ebpf::JSET_IMM32 => {
843                     self.emit_alu32_imm32(mem, 0xf7, 0, dst, insn.imm);
844                     self.emit_jcc(mem, 0x85, target_pc);
845                 }
846                 ebpf::JSET_REG32 => {
847                     self.emit_alu32(mem, 0x85, src, dst);
848                     self.emit_jcc(mem, 0x85, target_pc);
849                 }
850                 ebpf::JNE_IMM32 => {
851                     self.emit_cmp32_imm32(mem, dst, insn.imm);
852                     self.emit_jcc(mem, 0x85, target_pc);
853                 }
854                 ebpf::JNE_REG32 => {
855                     self.emit_cmp32(mem, src, dst);
856                     self.emit_jcc(mem, 0x85, target_pc);
857                 }
858                 ebpf::JSGT_IMM32 => {
859                     self.emit_cmp32_imm32(mem, dst, insn.imm);
860                     self.emit_jcc(mem, 0x8f, target_pc);
861                 }
862                 ebpf::JSGT_REG32 => {
863                     self.emit_cmp32(mem, src, dst);
864                     self.emit_jcc(mem, 0x8f, target_pc);
865                 }
866                 ebpf::JSGE_IMM32 => {
867                     self.emit_cmp32_imm32(mem, dst, insn.imm);
868                     self.emit_jcc(mem, 0x8d, target_pc);
869                 }
870                 ebpf::JSGE_REG32 => {
871                     self.emit_cmp32(mem, src, dst);
872                     self.emit_jcc(mem, 0x8d, target_pc);
873                 }
874                 ebpf::JSLT_IMM32 => {
875                     self.emit_cmp32_imm32(mem, dst, insn.imm);
876                     self.emit_jcc(mem, 0x8c, target_pc);
877                 }
878                 ebpf::JSLT_REG32 => {
879                     self.emit_cmp32(mem, src, dst);
880                     self.emit_jcc(mem, 0x8c, target_pc);
881                 }
882                 ebpf::JSLE_IMM32 => {
883                     self.emit_cmp32_imm32(mem, dst, insn.imm);
884                     self.emit_jcc(mem, 0x8e, target_pc);
885                 }
886                 ebpf::JSLE_REG32 => {
887                     self.emit_cmp32(mem, src, dst);
888                     self.emit_jcc(mem, 0x8e, target_pc);
889                 }
890 
891                 ebpf::CALL => {
892                     // For JIT, helpers in use MUST be registered at compile time. They can be
893                     // updated later, but not created after compiling (we need the address of the
894                     // helper function in the JIT-compiled program).
895                     if let Some(helper) = helpers.get(&(insn.imm as u32)) {
896                         // We reserve RCX for shifts
897                         self.emit_mov(mem, R9, RCX);
898                         self.emit_call(mem, *helper as usize);
899                     } else {
900                         Err(Error::new(
901                             ErrorKind::Other,
902                             format!(
903                                 "[JIT] Error: unknown helper function (id: {:#x})",
904                                 insn.imm as u32
905                             ),
906                         ))?;
907                     };
908                 }
909                 ebpf::TAIL_CALL => {
910                     unimplemented!()
911                 }
912                 ebpf::EXIT => {
913                     if insn_ptr != prog.len() / ebpf::INSN_SIZE - 1 {
914                         self.emit_jmp(mem, TARGET_PC_EXIT);
915                     };
916                 }
917 
918                 _ => {
919                     Err(Error::new(
920                         ErrorKind::Other,
921                         format!(
922                             "[JIT] Error: unknown eBPF opcode {:#2x} (insn #{insn_ptr:?})",
923                             insn.opc
924                         ),
925                     ))?;
926                 }
927             }
928 
929             insn_ptr += 1;
930         }
931 
932         // Epilogue
933         self.set_anchor(mem, TARGET_PC_EXIT);
934 
935         // Move register 0 into rax
936         if map_register(0) != RAX {
937             self.emit_mov(mem, map_register(0), RAX);
938         }
939 
940         // Deallocate stack space
941         self.emit_alu64_imm32(mem, 0x81, 0, RSP, ebpf::STACK_SIZE as i32);
942 
943         self.emit_pop(mem, R15);
944         self.emit_pop(mem, R14);
945         self.emit_pop(mem, R13);
946         self.emit_pop(mem, RBX);
947         self.emit_pop(mem, RBP);
948 
949         self.emit1(mem, 0xc3); // ret
950 
951         Ok(())
952     }
953 
954     fn resolve_jumps(&mut self, mem: &mut JitMemory) -> Result<(), Error> {
955         for jump in &self.jumps {
956             let target_loc = match self.special_targets.get(&jump.target_pc) {
957                 Some(target) => *target,
958                 None => self.pc_locs[jump.target_pc as usize],
959             };
960 
961             // Assumes jump offset is at end of instruction
962             unsafe {
963                 let offset_loc = jump.offset_loc as i32 + std::mem::size_of::<i32>() as i32;
964                 let rel = &(target_loc as i32 - offset_loc) as *const i32;
965 
966                 let offset_ptr = mem.contents.as_ptr().add(jump.offset_loc);
967 
968                 libc::memcpy(
969                     offset_ptr as *mut libc::c_void,
970                     rel as *const libc::c_void,
971                     std::mem::size_of::<i32>(),
972                 );
973             }
974         }
975         Ok(())
976     }
977 } // impl JitCompiler
978 
979 pub struct JitMemory<'a> {
980     contents: &'a mut [u8],
981     offset: usize,
982 }
983 
984 impl<'a> JitMemory<'a> {
985     pub fn new(
986         prog: &[u8],
987         helpers: &HashMap<u32, ebpf::Helper>,
988         use_mbuff: bool,
989         update_data_ptr: bool,
990     ) -> Result<JitMemory<'a>, Error> {
991         let contents: &mut [u8];
992         let mut raw: mem::MaybeUninit<*mut libc::c_void> = mem::MaybeUninit::uninit();
993         unsafe {
994             let size = NUM_PAGES * PAGE_SIZE;
995             libc::posix_memalign(raw.as_mut_ptr(), PAGE_SIZE, size);
996             libc::mprotect(
997                 *raw.as_mut_ptr(),
998                 size,
999                 libc::PROT_EXEC | libc::PROT_READ | libc::PROT_WRITE,
1000             );
1001             std::ptr::write_bytes(*raw.as_mut_ptr(), 0xc3, size); // for now, prepopulate with 'RET' calls
1002             contents =
1003                 std::slice::from_raw_parts_mut(*raw.as_mut_ptr() as *mut u8, NUM_PAGES * PAGE_SIZE);
1004             raw.assume_init();
1005         }
1006 
1007         let mut mem = JitMemory {
1008             contents,
1009             offset: 0,
1010         };
1011 
1012         let mut jit = JitCompiler::new();
1013         jit.jit_compile(&mut mem, prog, use_mbuff, update_data_ptr, helpers)?;
1014         jit.resolve_jumps(&mut mem)?;
1015 
1016         Ok(mem)
1017     }
1018 
1019     pub fn get_prog(&self) -> MachineCode {
1020         unsafe { mem::transmute(self.contents.as_ptr()) }
1021     }
1022 }
1023 
1024 impl<'a> Index<usize> for JitMemory<'a> {
1025     type Output = u8;
1026 
1027     fn index(&self, _index: usize) -> &u8 {
1028         &self.contents[_index]
1029     }
1030 }
1031 
1032 impl<'a> IndexMut<usize> for JitMemory<'a> {
1033     fn index_mut(&mut self, _index: usize) -> &mut u8 {
1034         &mut self.contents[_index]
1035     }
1036 }
1037 
1038 impl<'a> Drop for JitMemory<'a> {
1039     fn drop(&mut self) {
1040         unsafe {
1041             libc::free(self.contents.as_mut_ptr() as *mut libc::c_void);
1042         }
1043     }
1044 }
1045 
1046 impl<'a> std::fmt::Debug for JitMemory<'a> {
1047     fn fmt(&self, fmt: &mut Formatter) -> Result<(), FormatterError> {
1048         fmt.write_str("JIT contents: [")?;
1049         fmt.write_str(" ] | ")?;
1050         fmt.debug_struct("JIT memory")
1051             .field("offset", &self.offset)
1052             .finish()
1053     }
1054 }
1055