1#[cfg(feature = "shuttle-test")]
4use shuttle::rand;
5
6use {
7 crate::{
8 ebpf,
9 elf::ElfError,
10 error::EbpfError,
11 memory_management::{get_system_page_size, round_to_page_size, Mapping, Read, ReadExecute},
12 vm::{Config, ContextObject, EncryptedHostAddressToEbpfVm},
13 },
14 rand::{thread_rng, Rng},
15 std::collections::{hash_map::Entry, HashMap},
16 std::hash::{BuildHasher, Hasher},
17};
18
19#[derive(Debug, PartialEq, PartialOrd, Eq, Clone, Copy)]
21pub enum SBPFVersion {
22 V0,
24 V1,
26 V2,
28 V3,
30 V4,
32 Reserved,
34}
35
36impl SBPFVersion {
37 pub const ALL: &'static [SBPFVersion] = &[
39 SBPFVersion::V0,
40 SBPFVersion::V1,
41 SBPFVersion::V2,
42 SBPFVersion::V3,
43 SBPFVersion::V4,
44 ];
45
46 pub const fn manual_stack_frame_bump(self) -> bool {
50 matches!(self, SBPFVersion::V1 | SBPFVersion::V2)
51 }
52 pub const fn stack_frame_gaps(self) -> bool {
54 matches!(self, SBPFVersion::V0)
55 }
56
57 pub const fn enable_pqr(self) -> bool {
59 matches!(self, SBPFVersion::V2)
60 }
61 pub const fn explicit_sign_extension_of_results(self) -> bool {
63 matches!(self, SBPFVersion::V2)
64 }
65 pub const fn swap_sub_reg_imm_operands(self) -> bool {
67 matches!(self, SBPFVersion::V2)
68 }
69 pub const fn disable_neg(self) -> bool {
71 matches!(self, SBPFVersion::V2)
72 }
73
74 pub const fn callx_uses_src_reg(self) -> bool {
76 matches!(self, SBPFVersion::V2)
77 }
78 pub const fn disable_lddw(self) -> bool {
80 matches!(self, SBPFVersion::V2)
81 }
82 pub const fn disable_le(self) -> bool {
84 matches!(self, SBPFVersion::V2)
85 }
86 pub const fn move_memory_instruction_classes(self) -> bool {
88 matches!(self, SBPFVersion::V2)
89 }
90
91 pub const fn static_syscalls(self) -> bool {
93 matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
94 }
95 pub const fn enable_stricter_elf_headers(self) -> bool {
97 matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
98 }
99 pub const fn enable_lower_rodata_vaddr(self) -> bool {
101 matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
102 }
103 pub const fn enable_jmp32(self) -> bool {
105 matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
106 }
107 pub const fn callx_uses_dst_reg(self) -> bool {
109 matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
110 }
111
112 pub const fn calculate_call_imm_target_pc(self, pc: usize, imm: i64) -> u32 {
115 if self.static_syscalls() {
116 (pc as i64).saturating_add(imm).saturating_add(1) as u32
117 } else {
118 imm as u32
119 }
120 }
121}
122
123#[derive(Debug, PartialEq, Eq)]
127pub struct FunctionRegistry<T> {
128 map: HashMap<u32, (Vec<u8>, T), KeyHasherRandomState>,
131}
132
133impl<T> Default for FunctionRegistry<T> {
134 fn default() -> Self {
135 Self {
136 map: HashMap::default(),
137 }
138 }
139}
140
141#[derive(Clone, Copy)]
142struct KeyHasherRandomState {
143 multiplier: u64,
144}
145
146impl Default for KeyHasherRandomState {
147 fn default() -> Self {
148 Self {
149 multiplier: thread_rng().gen::<u64>() | 1,
150 }
151 }
152}
153
154impl BuildHasher for KeyHasherRandomState {
155 type Hasher = KeyHasher;
156 fn build_hasher(&self) -> KeyHasher {
157 KeyHasher {
158 multiplier: self.multiplier,
159 hash: 0,
160 }
161 }
162}
163
164struct KeyHasher {
165 multiplier: u64,
166 hash: u64,
167}
168
169impl Hasher for KeyHasher {
170 fn finish(&self) -> u64 {
171 self.hash
172 }
173
174 fn write(&mut self, bytes: &[u8]) {
175 for &byte in bytes {
177 self.write_u32(self.hash as u32 ^ u32::from(byte));
178 }
179 }
180
181 fn write_u32(&mut self, key: u32) {
182 let product = u64::from(key).wrapping_mul(self.multiplier);
183 self.hash = product ^ product >> 32;
184 }
185}
186
187impl<T: Copy + PartialEq> FunctionRegistry<T> {
188 pub fn register_function(
190 &mut self,
191 key: u32,
192 name: impl Into<Vec<u8>>,
193 value: T,
194 ) -> Result<(), ElfError> {
195 match self.map.entry(key) {
196 Entry::Vacant(entry) => {
197 entry.insert((name.into(), value));
198 }
199 Entry::Occupied(entry) => {
200 if entry.get().1 != value {
201 return Err(ElfError::SymbolHashCollision(key));
202 }
203 }
204 }
205 Ok(())
206 }
207
208 pub(crate) fn register_function_hashed_legacy<C: ContextObject>(
210 &mut self,
211 loader: &BuiltinProgram<C>,
212 hash_symbol_name: bool,
213 name: impl Into<Vec<u8>>,
214 value: T,
215 ) -> Result<u32, ElfError>
216 where
217 usize: From<T>,
218 {
219 let name = name.into();
220 let config = loader.get_config();
221 let key = if hash_symbol_name {
222 let hash = if name == b"entrypoint" {
223 ebpf::hash_symbol_name(b"entrypoint")
224 } else {
225 ebpf::hash_symbol_name(&usize::from(value).to_le_bytes())
226 };
227 if loader.get_function_registry().lookup_by_key(hash).is_some() {
228 return Err(ElfError::SymbolHashCollision(hash));
229 }
230 hash
231 } else {
232 usize::from(value) as u32
233 };
234 self.register_function(
235 key,
236 if config.enable_symbol_and_section_labels || name == b"entrypoint" {
237 name
238 } else {
239 Vec::default()
240 },
241 value,
242 )?;
243 Ok(key)
244 }
245
246 pub fn unregister_function(&mut self, key: u32) -> bool {
248 self.map.remove(&key).is_some()
249 }
250
251 pub fn keys(&self) -> impl Iterator<Item = u32> + '_ {
253 self.map.keys().copied()
254 }
255
256 pub fn iter(&self) -> impl Iterator<Item = (u32, (&[u8], T))> + '_ {
258 self.map
259 .iter()
260 .map(|(key, (name, value))| (*key, (name.as_slice(), *value)))
261 }
262
263 #[inline]
265 pub fn lookup_by_key(&self, key: u32) -> Option<(&[u8], T)> {
266 self.map
267 .get(&key)
268 .map(|(function_name, value)| (function_name.as_slice(), *value))
269 }
270
271 pub fn lookup_by_name(&self, name: &[u8]) -> Option<(&[u8], T)> {
273 self.map
274 .values()
275 .find(|(function_name, _value)| function_name == name)
276 .map(|(function_name, value)| (function_name.as_slice(), *value))
277 }
278
279 pub fn mem_size(&self) -> usize {
281 std::mem::size_of::<Self>().saturating_add(self.map.values().fold(
282 0,
283 |state: usize, (name, value)| {
284 state.saturating_add(
285 std::mem::size_of_val(value).saturating_add(
286 std::mem::size_of_val(name).saturating_add(name.capacity()),
287 ),
288 )
289 },
290 ))
291 }
292}
293
294pub struct UnsealedJitProgram {
296 pc_section: Mapping,
298 text_section: Mapping,
300}
301
302impl UnsealedJitProgram {
303 pub fn new(pc: usize, code_capacity: usize) -> Self {
306 let page_size = get_system_page_size();
307 let pc_len = pc.saturating_mul(std::mem::size_of::<u32>());
308 let pc_size = round_to_page_size(pc_len, page_size);
309 let text_capacity = round_to_page_size(code_capacity, page_size);
310 let mapping = Mapping::pooled(pc_size.saturating_add(text_capacity));
311 let Ok([pc_section, text_section]) =
312 mapping.split([0..pc_len, pc_size..pc_size.saturating_add(text_capacity)])
313 else {
314 unreachable!("mapping split failed for some reason");
315 };
316 let mut program = Self {
317 pc_section,
318 text_section,
319 };
320 program.pc_section.fill(0);
324 program
325 }
326
327 #[inline]
329 pub fn pc_section(&self) -> &[u32] {
330 unsafe {
331 let (&[], out, &[]) = self.pc_section.align_to() else {
333 unreachable!("pc_section can't be unaligned to u32");
334 };
335 out
336 }
337 }
338
339 #[inline]
341 pub fn text_section(&self) -> &[u8] {
342 &self.text_section
343 }
344
345 #[inline]
347 pub fn pc_section_mut(&mut self) -> &mut [u32] {
348 unsafe {
350 let (&mut [], out, &mut []) = self.pc_section.align_to_mut() else {
352 unreachable!("pc_section can't be unaligned to u32");
353 };
354 out
355 }
356 }
357
358 #[inline]
360 pub fn text_section_mut(&mut self) -> &mut [u8] {
361 &mut self.text_section
362 }
363
364 pub fn seal(self, used: usize) -> Result<JitProgram, EbpfError> {
367 let Self {
368 pc_section,
369 mut text_section,
370 } = self;
371 if used > text_section.len() {
372 return Err(EbpfError::ExhaustedTextSegment(used));
373 }
374 let code_size = round_to_page_size(used, get_system_page_size());
376 text_section[used..code_size].fill(0xcc);
377 text_section.truncate(used);
378 Ok(JitProgram {
379 pc_section: pc_section.protect()?,
380 text_section: text_section.protect()?,
381 })
382 }
383}
384
385pub struct JitProgram {
389 pc_section: Mapping<Read>,
391 text_section: Mapping<ReadExecute>,
393}
394
395impl JitProgram {
396 #[inline]
398 pub fn pc_section(&self) -> &[u32] {
399 unsafe {
400 let (&[], out, &[]) = self.pc_section.align_to() else {
402 unreachable!("pc_section can't be unaligned to u32");
403 };
404 out
405 }
406 }
407
408 #[inline]
410 pub fn text_section(&self) -> &[u8] {
411 &self.text_section
412 }
413
414 pub fn machine_code_length(&self) -> usize {
416 self.text_section().len()
417 }
418
419 pub fn mem_size(&self) -> usize {
421 self.text_section.mapped_len()
422 }
423}
424
425impl std::fmt::Debug for JitProgram {
426 fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
427 fmt.write_fmt(format_args!("JitProgram {:?}", self as *const _))
428 }
429}
430
431impl PartialEq for JitProgram {
432 fn eq(&self, other: &Self) -> bool {
433 std::ptr::eq(self, other)
434 }
435}
436
437pub type BuiltinFunction<C> = fn(EncryptedHostAddressToEbpfVm<C>, u64, u64, u64, u64, u64);
439#[cfg(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64"))]
441pub type JitCompiler<'a, C> = crate::jit::JitCompiler<'a, C>;
442#[cfg(not(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64")))]
444pub struct JitCompiler<'a, C> {
445 _phantom: std::marker::PhantomData<&'a C>,
446}
447#[cfg(not(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64")))]
448impl<'a, C: ContextObject> JitCompiler<'a, C> {
449 #[allow(dead_code)]
451 pub fn emit_external_call(&mut self, _function: BuiltinFunction<C>) {}
452}
453pub type BuiltinCodegen<C> = fn(&mut JitCompiler<C>);
455
456#[derive(Eq)]
458pub struct BuiltinProgram<C: ContextObject> {
459 config: Option<Box<Config>>,
461 sparse_registry: FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)>,
463}
464
465impl<C: ContextObject> PartialEq for BuiltinProgram<C> {
466 fn eq(&self, other: &Self) -> bool {
467 self.config.eq(&other.config) && self.sparse_registry.eq(&other.sparse_registry)
468 }
469}
470
471impl<C: ContextObject> BuiltinProgram<C> {
472 pub fn new_loader(config: Config) -> Self {
474 Self {
475 config: Some(Box::new(config)),
476 sparse_registry: FunctionRegistry::default(),
477 }
478 }
479
480 pub fn new_builtin() -> Self {
482 Self {
483 config: None,
484 sparse_registry: FunctionRegistry::default(),
485 }
486 }
487
488 pub fn new_mock() -> Self {
490 Self {
491 config: Some(Box::default()),
492 sparse_registry: FunctionRegistry::default(),
493 }
494 }
495
496 pub fn get_config(&self) -> &Config {
498 self.config.as_ref().unwrap()
499 }
500
501 pub fn get_function_registry(
503 &self,
504 ) -> &FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)> {
505 &self.sparse_registry
506 }
507
508 pub fn mem_size(&self) -> usize {
510 std::mem::size_of::<Self>()
511 .saturating_add(if self.config.is_some() {
512 std::mem::size_of::<Config>()
513 } else {
514 0
515 })
516 .saturating_add(self.sparse_registry.mem_size())
517 }
518
519 pub fn register_function(
523 &mut self,
524 name: &str,
525 entry: (BuiltinFunction<C>, BuiltinCodegen<C>),
526 ) -> Result<(), ElfError> {
527 let key = ebpf::hash_symbol_name(name.as_bytes());
528 self.sparse_registry
529 .register_function(key, name, entry)
530 .map(|_| ())
531 }
532
533 pub fn register_definition<BFD: BuiltinFunctionDefinition<C>>(
535 &mut self,
536 name: &str,
537 ) -> Result<(), ElfError> {
538 self.register_function(name, (BFD::vm, BFD::codegen))
539 }
540
541 pub fn unregister_function(&mut self, name: &str) -> bool {
543 let key = ebpf::hash_symbol_name(name.as_bytes());
544 self.sparse_registry.unregister_function(key)
545 }
546}
547
548pub trait BuiltinFunctionDefinition<C>
550where
551 C: crate::vm::ContextObject,
552{
553 type Error: Into<Box<dyn core::error::Error>>;
555
556 fn rust(
560 vm: &mut C,
561 arg_a: u64,
562 arg_b: u64,
563 arg_c: u64,
564 arg_d: u64,
565 arg_e: u64,
566 ) -> Result<u64, Self::Error>;
567
568 #[expect(clippy::arithmetic_side_effects)]
570 fn vm(mut vm: EncryptedHostAddressToEbpfVm<C>, a: u64, b: u64, c: u64, d: u64, e: u64) {
571 unsafe {
572 vm.with_vm(|vm| {
574 let enable_insn_meter = vm.loader.get_config().enable_instruction_meter;
575 if enable_insn_meter {
576 let used_cus = vm.previous_instruction_meter - vm.due_insn_count;
577 vm.context().consume(used_cus);
578 }
579 let converted_result: crate::error::ProgramResult =
580 Self::rust(vm.context(), a, b, c, d, e)
581 .map_err(|err| crate::error::EbpfError::SyscallError(err.into()))
582 .into();
583 vm.program_result = converted_result;
584 if enable_insn_meter {
585 vm.previous_instruction_meter = vm.context().get_remaining();
586 }
587 })
588 }
589 }
590
591 fn codegen(jit: &mut JitCompiler<C>) {
596 jit.emit_external_call(Self::vm);
597 }
598
599 fn register(program: &mut BuiltinProgram<C>, name: &str) -> Result<(), ElfError>
601 where
602 Self: Sized,
603 {
604 program.register_definition::<Self>(name)
605 }
606}
607
608impl<C: ContextObject> std::fmt::Debug for BuiltinProgram<C> {
609 fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
610 f.debug_struct("BuiltinProgram")
611 .field("registry", unsafe {
612 std::mem::transmute::<
613 &FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)>,
614 &FunctionRegistry<(usize, usize)>,
615 >(&self.sparse_registry)
616 })
617 .finish()
618 }
619}
620
621#[macro_export]
623macro_rules! declare_builtin_function {
624 (
625 $(#[$attr:meta])*
626 $name:ident $(<$($generic_ident:tt : $generic_type:tt),+>)?,
627 fn rust(
628 $vm:ident : &mut $ContextObject:ty,
629 $arg_a:ident : u64,
630 $arg_b:ident : u64,
631 $arg_c:ident : u64,
632 $arg_d:ident : u64,
633 $arg_e:ident : u64,
634 ) -> Result<$Ok:ty, $Err:ty> {
635 $($rust:tt)*
636 }
637 $(fn codegen($jit:ident : &mut JitCompiler<$ContextObject2:ty>) {
638 $($codegen:tt)*
639 })?
640 ) => {
641 $(#[$attr])*
642 pub struct $name $(<$($generic_ident),+>)? (
643 $(std::marker::PhantomData<($($generic_ident,)+)>)?
644 );
645 impl $(<$($generic_ident : $generic_type),+>)?
646 $crate::program::BuiltinFunctionDefinition<$ContextObject> for
647 $name $(<$($generic_ident),+>)?
648 {
649 type Error = $Err;
650 fn rust(
651 $vm: &mut $ContextObject,
652 $arg_a: u64,
653 $arg_b: u64,
654 $arg_c: u64,
655 $arg_d: u64,
656 $arg_e: u64,
657 ) -> core::result::Result<$Ok, $Err> {
658 $($rust)*
659 }
660 $(fn codegen(
661 $jit: &mut $crate::program::JitCompiler<$ContextObject2>,
662 ) {
663 $($codegen)*
664 })?
665 }
666 };
667}
668
669#[cfg(test)]
670mod tests {
671 use super::*;
672
673 macro_rules! check_lookups {
674 ($registry:ident, $lens:expr) => {
675 let key = |i: u32| ebpf::hash_symbol_name(&i.to_le_bytes());
676 for len in $lens {
677 let mut registry = $registry::default();
678 for i in 0..len {
679 registry
680 .register_function(key(i), format!("f{i}"), i)
681 .unwrap();
682 }
683 let found = |registry: &$registry<u32>, i: u32| {
684 registry
685 .lookup_by_key(key(i))
686 .map(|(name, value)| (String::from_utf8(name.to_vec()).unwrap(), value))
687 };
688 for i in 0..len {
689 assert_eq!(found(®istry, i), Some((format!("f{i}"), i)));
690 }
691 for i in len..len.saturating_add(1000) {
692 assert_eq!(found(®istry, i), None);
693 }
694 let mut keys: Vec<_> = (0..len).map(key).collect();
695 keys.sort();
696 let mut registered: Vec<_> = registry.keys().collect();
697 registered.sort();
698 assert_eq!(registered, keys);
699 let mut iterated: Vec<_> = registry.iter().map(|(key, _)| key).collect();
700 iterated.sort();
701 assert_eq!(iterated, keys);
702 if len > 0 {
703 assert!(registry.register_function(key(0), "f0", 0).is_ok());
704 assert!(matches!(
705 registry.register_function(key(0), "f0", 1),
706 Err(ElfError::SymbolHashCollision(_))
707 ));
708 assert!(registry.unregister_function(key(0)));
709 assert!(!registry.unregister_function(key(0)));
710 assert_eq!(found(®istry, 0), None);
711 for i in 1..len {
712 assert_eq!(found(®istry, i), Some((format!("f{i}"), i)));
713 }
714 assert_eq!(registry.keys().count(), (len as usize).saturating_sub(1));
715 registry.register_function(key(0), "f0", 0).unwrap();
716 assert_eq!(found(®istry, 0), Some(("f0".to_string(), 0)));
717 for i in 0..len {
718 assert!(registry.unregister_function(key(i)));
719 }
720 assert_eq!(registry.keys().count(), 0);
721 for i in 0..len.saturating_add(1000) {
722 assert_eq!(found(®istry, i), None);
723 }
724 }
725 }
726 };
727 }
728
729 #[test]
730 fn function_registry_lookups() {
731 check_lookups!(FunctionRegistry, [0u32, 1, 3, 64, 1000]);
732 }
733
734 #[test]
735 fn jit_program_sections() {
736 let page_size = get_system_page_size();
737 let mut program = UnsealedJitProgram::new(3, page_size + 1);
738 assert_eq!(program.pc_section(), [0; 3]);
739 assert_eq!(program.text_section().len(), 2 * page_size);
740 program.pc_section_mut().copy_from_slice(&[1, 2, 3]);
741 program.text_section_mut()[..2].copy_from_slice(&[0x90, 0xc3]);
742 assert!(matches!(
743 UnsealedJitProgram::new(3, page_size + 1).seal(2 * page_size + 1),
744 Err(EbpfError::ExhaustedTextSegment(_))
745 ));
746 let program = program.seal(2).unwrap();
747 assert_eq!(program.pc_section(), [1, 2, 3]);
748 assert_eq!(program.text_section(), [0x90, 0xc3]);
749 assert_eq!(program.machine_code_length(), 2);
750 assert!(program.mem_size() >= 3 * page_size);
751 }
752
753 #[test]
755 fn jit_program_pc_section_zeroed() {
756 for _ in 0..4 {
757 let mut program = UnsealedJitProgram::new(16, 16);
758 assert!(program.pc_section().iter().all(|&pc| pc == 0));
759 program.pc_section_mut().fill(u32::MAX);
760 program.seal(0).unwrap();
761 }
762 }
763}