Skip to main content

solana_sbpf/
program.rs

1//! Common interface for built-in and user supplied programs
2
3#[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/// Defines a set of sbpf_version of an executable
20#[derive(Debug, PartialEq, PartialOrd, Eq, Clone, Copy)]
21pub enum SBPFVersion {
22    /// The legacy format
23    V0,
24    /// SIMD-0166
25    V1,
26    /// SIMD-0174, SIMD-0173
27    V2,
28    /// SIMD-0178, SIMD-0189, SIMD-0377
29    V3,
30    /// SIMD-0177
31    V4,
32    /// Used for future versions
33    Reserved,
34}
35
36impl SBPFVersion {
37    /// Every defined version in ascending order, excluding `Reserved`.
38    pub const ALL: &'static [SBPFVersion] = &[
39        SBPFVersion::V0,
40        SBPFVersion::V1,
41        SBPFVersion::V2,
42        SBPFVersion::V3,
43        SBPFVersion::V4,
44    ];
45
46    /// Enable SIMD-0166: SBPF dynamic stack frames
47    ///
48    /// Allows usage of `add64 r10, imm`.
49    pub const fn manual_stack_frame_bump(self) -> bool {
50        matches!(self, SBPFVersion::V1 | SBPFVersion::V2)
51    }
52    /// ... SIMD-0166
53    pub const fn stack_frame_gaps(self) -> bool {
54        matches!(self, SBPFVersion::V0)
55    }
56
57    /// Enable SIMD-0174: SBPF arithmetics improvements
58    pub const fn enable_pqr(self) -> bool {
59        matches!(self, SBPFVersion::V2)
60    }
61    /// ... SIMD-0174
62    pub const fn explicit_sign_extension_of_results(self) -> bool {
63        matches!(self, SBPFVersion::V2)
64    }
65    /// ... SIMD-0174
66    pub const fn swap_sub_reg_imm_operands(self) -> bool {
67        matches!(self, SBPFVersion::V2)
68    }
69    /// ... SIMD-0174
70    pub const fn disable_neg(self) -> bool {
71        matches!(self, SBPFVersion::V2)
72    }
73
74    /// Enable SIMD-0173: SBPF instruction encoding improvements
75    pub const fn callx_uses_src_reg(self) -> bool {
76        matches!(self, SBPFVersion::V2)
77    }
78    /// ... SIMD-0173
79    pub const fn disable_lddw(self) -> bool {
80        matches!(self, SBPFVersion::V2)
81    }
82    /// ... SIMD-0173
83    pub const fn disable_le(self) -> bool {
84        matches!(self, SBPFVersion::V2)
85    }
86    /// ... SIMD-0173
87    pub const fn move_memory_instruction_classes(self) -> bool {
88        matches!(self, SBPFVersion::V2)
89    }
90
91    /// Enable SIMD-0178: SBPF Static Syscalls
92    pub const fn static_syscalls(self) -> bool {
93        matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
94    }
95    /// Enable SIMD-0189: SBPF stricter ELF headers
96    pub const fn enable_stricter_elf_headers(self) -> bool {
97        matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
98    }
99    /// ... SIMD-0189
100    pub const fn enable_lower_rodata_vaddr(self) -> bool {
101        matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
102    }
103    /// ... SIMD-0377
104    pub const fn enable_jmp32(self) -> bool {
105        matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
106    }
107    /// ... SIMD-0377
108    pub const fn callx_uses_dst_reg(self) -> bool {
109        matches!(self, SBPFVersion::V3 | SBPFVersion::V4)
110    }
111
112    /// Calculate the target program counter for a CALL_IMM instruction depending on
113    /// the SBPF version.
114    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/// Map suitable to hold user-generated sets of function symbols in an [`Executable`].
124///
125/// [`Executable`]: crate::elf::Executable
126#[derive(Debug, PartialEq, Eq)]
127pub struct FunctionRegistry<T> {
128    /// This map does not further hash the keys (which already are hashes), and instead randomizes
129    /// them somewhat to avoid malicious exploitation of big-O factors in hashmaps.
130    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        // Only `write_u32` is used for the keys.
176        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    /// Register a symbol with an explicit key
189    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    /// Used for transitioning from SBPFv0 to SBPFv3
209    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    /// Unregister a symbol again
247    pub fn unregister_function(&mut self, key: u32) -> bool {
248        self.map.remove(&key).is_some()
249    }
250
251    /// Iterate over all keys
252    pub fn keys(&self) -> impl Iterator<Item = u32> + '_ {
253        self.map.keys().copied()
254    }
255
256    /// Iterate over all entries
257    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    /// Get a function by its key
264    #[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    /// Get a function by its name
272    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    /// Calculate memory size
280    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
294/// A program that is in progress of being written out.
295pub struct UnsealedJitProgram {
296    /// Offsets to each BPF instruction's machine code within `text_section`.
297    pc_section: Mapping,
298    /// The machine code.
299    text_section: Mapping,
300}
301
302impl UnsealedJitProgram {
303    /// Allocate a program with `pc` zeroed entries in the pc section and room for `code_capacity`
304    /// bytes of machine code.
305    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        // The pc section relies on zero-initialization to distinguish unfilled forward-jump
321        // targets from filled backward-jump targets in the old JIT. The pool may hand back
322        // recycled memory. Zero just the pc section here.
323        program.pc_section.fill(0);
324        program
325    }
326
327    /// Offset in `text_section` for each BPF instruction.
328    #[inline]
329    pub fn pc_section(&self) -> &[u32] {
330        unsafe {
331            // SAFETY: u32 has the same invariants as u8
332            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    /// The machine code.
340    #[inline]
341    pub fn text_section(&self) -> &[u8] {
342        &self.text_section
343    }
344
345    /// The pc section to fill in.
346    #[inline]
347    pub fn pc_section_mut(&mut self) -> &mut [u32] {
348        // FIXME: if there are more of such casts in the code, consider `bytemuck`.
349        unsafe {
350            // SAFETY: u32 has the same invariants as u8
351            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    /// The machine code to fill in.
359    #[inline]
360    pub fn text_section_mut(&mut self) -> &mut [u8] {
361        &mut self.text_section
362    }
363
364    /// Make the pc section read-only and the text section, shrunk to the `used` length,
365    /// executable.
366    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        // Debugger traps in the unused tail of the last code page.
375        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
385/// A ready-to-execute JIT program.
386///
387/// See [`UnsealedJitProgram::seal`].
388pub struct JitProgram {
389    /// Offset to each BPF instruction's machine code within `text_section`.
390    pc_section: Mapping<Read>,
391    /// The executable machine code.
392    text_section: Mapping<ReadExecute>,
393}
394
395impl JitProgram {
396    /// Offset in `text_section` for each BPF instruction.
397    #[inline]
398    pub fn pc_section(&self) -> &[u32] {
399        unsafe {
400            // SAFETY: u32 has the same invariants as u8
401            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    /// The executable machine code.
409    #[inline]
410    pub fn text_section(&self) -> &[u8] {
411        &self.text_section
412    }
413
414    /// The length of the host machinecode in bytes
415    pub fn machine_code_length(&self) -> usize {
416        self.text_section().len()
417    }
418
419    /// The total pooled allocation size retained by the compiled program.
420    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
437/// Syscall handler function (ContextObject is derived from VM)
438pub type BuiltinFunction<C> = fn(EncryptedHostAddressToEbpfVm<C>, u64, u64, u64, u64, u64);
439/// Re-export of the JIT compiler for the declare_builtin_function! macro
440#[cfg(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64"))]
441pub type JitCompiler<'a, C> = crate::jit::JitCompiler<'a, C>;
442/// Re-export of the JIT compiler for the declare_builtin_function! macro
443#[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    /// Dummy for declare_builtin_function!()
450    #[allow(dead_code)]
451    pub fn emit_external_call(&mut self, _function: BuiltinFunction<C>) {}
452}
453/// Syscall codegen function for JIT compiler
454pub type BuiltinCodegen<C> = fn(&mut JitCompiler<C>);
455
456/// Represents the interface to a fixed functionality program
457#[derive(Eq)]
458pub struct BuiltinProgram<C: ContextObject> {
459    /// Holds the Config if this is a loader program
460    config: Option<Box<Config>>,
461    /// Function pointers by symbol with sparse indexing
462    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    /// Constructs a loader built-in program
473    pub fn new_loader(config: Config) -> Self {
474        Self {
475            config: Some(Box::new(config)),
476            sparse_registry: FunctionRegistry::default(),
477        }
478    }
479
480    /// Constructs a built-in program
481    pub fn new_builtin() -> Self {
482        Self {
483            config: None,
484            sparse_registry: FunctionRegistry::default(),
485        }
486    }
487
488    /// Constructs a mock loader built-in program
489    pub fn new_mock() -> Self {
490        Self {
491            config: Some(Box::default()),
492            sparse_registry: FunctionRegistry::default(),
493        }
494    }
495
496    /// Get the configuration settings assuming this is a loader program
497    pub fn get_config(&self) -> &Config {
498        self.config.as_ref().unwrap()
499    }
500
501    /// Get the function registry depending on the SBPF version
502    pub fn get_function_registry(
503        &self,
504    ) -> &FunctionRegistry<(BuiltinFunction<C>, BuiltinCodegen<C>)> {
505        &self.sparse_registry
506    }
507
508    /// Calculate memory size
509    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    /// Register a function both in the sparse and dense registries
520    ///
521    /// This is a low-level function. Prefer using [`Self::register_definition`].
522    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    /// Register a function both in the sparse and dense registries
534    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    /// Remove a function by name (if it exists)
542    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
548/// Native built-in functions that can be made available to programs to call.
549pub trait BuiltinFunctionDefinition<C>
550where
551    C: crate::vm::ContextObject,
552{
553    /// Error type returned by this built-in function.
554    type Error: Into<Box<dyn core::error::Error>>;
555
556    /// The Rust side of the function logic.
557    ///
558    /// This is the only method you are required to override.
559    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    /// The VM wrapper.
569    #[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            // SAFETY: Sound under the stacked lifetimes model – we've only one `VmAddress`.
573            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    /// Hook for the JIT compiler on how to codegen this built-in function.
592    ///
593    /// You could opt to codegen it in-line, but do note that defining the other methods is still
594    /// required for non-JIT execution modes.
595    fn codegen(jit: &mut JitCompiler<C>) {
596        jit.emit_external_call(Self::vm);
597    }
598
599    /// Register this syscall to the provided program.
600    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/// Generates an adapter for a BuiltinFunction between the Rust and the VM interface
622#[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(&registry, i), Some((format!("f{i}"), i)));
690                }
691                for i in len..len.saturating_add(1000) {
692                    assert_eq!(found(&registry, 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(&registry, 0), None);
711                    for i in 1..len {
712                        assert_eq!(found(&registry, 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(&registry, 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(&registry, 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    /// The pc section of a program is zeroed even if its pages are reused.
754    #[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}