xref: /DragonOS/kernel/crates/rbpf/examples/rbpf_plugin.rs (revision 7b0ef10895108a0de5ff5ef3d2f93f40cf2e33a5)
1 // Copyright Microsoft Corporation
2 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
3 
4 // Path: examples/rbpf_plugin.rs
5 use std::io::Read;
6 
7 // Helper function used by https://github.com/Alan-Jowett/bpf_conformance/blob/main/tests/call_unwind_fail.data
8 fn _unwind(a: u64, _b: u64, _c: u64, _d: u64, _e: u64) -> u64 {
9     a
10 }
11 
12 // This is a plugin for the bpf_conformance test suite (https://github.com/Alan-Jowett/bpf_conformance)
13 // It accepts a single argument, the memory contents to pass to the VM.
14 // It reads the program from stdin.
15 fn main() {
16     let mut args: Vec<String> = std::env::args().collect();
17     #[allow(unused_mut)] // In no_std the jit variable isn't mutated.
18     let mut jit: bool = false;
19     let mut cranelift: bool = false;
20     let mut program_text = String::new();
21     let mut memory_text = String::new();
22 
23     args.remove(0);
24 
25     // Memory is always the first argument.
26     if !args.is_empty() {
27         memory_text.clone_from(&args[0]);
28         // Strip whitespace
29         memory_text.retain(|c| !c.is_whitespace());
30         args.remove(0);
31     }
32 
33     // Process the rest of the arguments.
34     while !args.is_empty() {
35         match args[0].as_str() {
36             "--help" => {
37                 println!("Usage: rbpf_plugin [memory] < program");
38                 return;
39             }
40             "--jit" => {
41                 #[cfg(any(windows, not(feature = "std")))]
42                 {
43                     println!("JIT not supported");
44                     return;
45                 }
46                 #[cfg(all(not(windows), feature = "std"))]
47                 {
48                     jit = true;
49                 }
50             }
51             "--cranelift" => {
52                 cranelift = true;
53 
54                 #[cfg(not(feature = "cranelift"))]
55                 {
56                     let _ = cranelift;
57                     println!("Cranelift is not enabled");
58                     return;
59                 }
60             }
61             "--program" => {
62                 if args.len() < 2 {
63                     println!("Missing argument to --program");
64                     return;
65                 }
66                 args.remove(0);
67                 if !args.is_empty() {
68                     program_text.clone_from(&args[0]);
69                     args.remove(0);
70                 }
71             }
72             _ => panic!("Unknown argument {}", args[0]),
73         }
74         args.remove(0);
75     }
76 
77     if program_text.is_empty() {
78         // Read program text from stdin
79         std::io::stdin().read_to_string(&mut program_text).unwrap();
80     }
81 
82     // Strip whitespace
83     program_text.retain(|c| !c.is_whitespace());
84 
85     // Convert program from hex to bytecode
86     let bytecode = hex::decode(program_text).unwrap();
87 
88     // Convert memory from hex to bytes
89     let mut memory: Vec<u8> = hex::decode(memory_text).unwrap();
90 
91     // Create rbpf vm
92     let mut vm = rbpf::EbpfVmRaw::new(Some(&bytecode)).unwrap();
93 
94     // Register the helper function used by call_unwind_fail.data test.
95     vm.register_helper(5, _unwind).unwrap();
96 
97     let result: u64;
98     if jit {
99         #[cfg(any(windows, not(feature = "std")))]
100         {
101             println!("JIT not supported");
102             return;
103         }
104         #[cfg(all(not(windows), feature = "std"))]
105         {
106             unsafe {
107                 vm.jit_compile().unwrap();
108                 result = vm.execute_program_jit(&mut memory).unwrap();
109             }
110         }
111     } else if cranelift {
112         #[cfg(not(feature = "cranelift"))]
113         {
114             println!("Cranelift is not enabled");
115             return;
116         }
117         #[cfg(feature = "cranelift")]
118         {
119             vm.cranelift_compile().unwrap();
120             result = vm.execute_program_cranelift(&mut memory).unwrap();
121         }
122     } else {
123         result = vm.execute_program(&mut memory).unwrap();
124     }
125     println!("{result:x}");
126 }
127