rbpf/examples/rbpf_plugin.rs

// Copyright Microsoft Corporation
// SPDX-License-Identifier: (Apache-2.0 OR MIT)

// Path: examples/rbpf_plugin.rs
use std::io::Read;

// Helper function used by https://github.com/Alan-Jowett/bpf_conformance/blob/main/tests/call_unwind_fail.data
fn _unwind(a: u64, _b: u64, _c: u64, _d: u64, _e: u64) -> u64
{
    return a;
}

// This is a plugin for the bpf_conformance test suite (https://github.com/Alan-Jowett/bpf_conformance)
// It accepts a single argument, the memory contents to pass to the VM.
// It reads the program from stdin.
fn main() {
    let mut args: Vec<String> = std::env::args().collect();
    let mut jit : bool = false;
    let mut program_text = String::new();
    let mut memory_text = String::new();

    args.remove(0);
    
    // Memory is always the first argument.
    if args.len() > 0 {
        memory_text = args[0].clone();
        // Strip whitespace
        memory_text.retain(|c| !c.is_whitespace());
        args.remove(0);
    }

    // Process the rest of the arguments.
    while args.len() > 0 {
        match args[0].as_str() {
            "--help" => { 
                println!("Usage: rbpf_plugin [memory] < program");
                return;
            },
            "--jit" => { 
                #[cfg(windows)] {
                    println!("JIT not supported on Windows");
                    return;
                }
                #[cfg(not(windows))] {
                    jit = true; 
                }
            },
            "--program" => { 
                if args.len() < 2 {
                    println!("Missing argument to --program");
                    return;
                }
                args.remove(0);
                if args.len() > 0 {
                    program_text = args[0].clone();
                    args.remove(0);
                }
            },
            _ => panic!("Unknown argument {}", args[0]),
        }
        args.remove(0);
    }

    if program_text.len() == 0 {
        // Read program text from stdin
        std::io::stdin().read_to_string(&mut program_text).unwrap();
    }

    // Strip whitespace
    program_text.retain(|c| !c.is_whitespace());

    // Convert program from hex to bytecode
    let bytecode = hex::decode(program_text).unwrap();

    // Convert memory from hex to bytes
    let mut memory: Vec<u8> = hex::decode(memory_text).unwrap();

    // Create rbpf vm
    let mut vm = rbpf::EbpfVmRaw::new(Some(&bytecode)).unwrap();
    
    // Register the helper function used by call_unwind_fail.data test.
    vm.register_helper(5, _unwind).unwrap();
    
    let result : u64;
    if jit {
        #[cfg(windows)] {
            println!("JIT not supported on Windows");
            return;
        }
        #[cfg(not(windows))] {
            unsafe { 
                vm.jit_compile().unwrap();
                result = vm.execute_program_jit(&mut memory).unwrap();
            }
        }
    }
    else {
        result = vm.execute_program(&mut memory).unwrap();
    }
    println!("{result:x}");
}