sbpf-runtime 0.2.3

use {
    crate::{
        config::{ExecutionCost, RuntimeConfig, SysvarContext},
        cpi,
        elf::load_elf,
        errors::{RuntimeError, RuntimeResult},
        serialize,
        syscalls::RuntimeSyscallHandler,
    },
    base64::{Engine, engine::general_purpose::STANDARD as BASE64},
    sbpf_common::{execute::Vm, instruction::Instruction},
    sbpf_vm::{
        compute::ComputeMeter,
        memory::Memory,
        vm::{CallFrame, SbpfVm, SbpfVmConfig},
    },
    solana_account::Account,
    solana_address::Address,
    solana_instruction::{AccountMeta, Instruction as SolanaInstruction},
    std::{cell::RefCell, collections::HashMap, rc::Rc},
};

pub type LogCollector = Rc<RefCell<Vec<String>>>;

pub enum ElfSource {
    Path(String),
    Bytes(Vec<u8>),
}

impl From<&str> for ElfSource {
    fn from(path: &str) -> Self {
        ElfSource::Path(path.to_string())
    }
}

impl From<&[u8]> for ElfSource {
    fn from(bytes: &[u8]) -> Self {
        ElfSource::Bytes(bytes.to_vec())
    }
}

impl From<Vec<u8>> for ElfSource {
    fn from(bytes: Vec<u8>) -> Self {
        ElfSource::Bytes(bytes)
    }
}

pub struct ExecutionResult {
    pub exit_code: Option<u64>,
    pub compute_units_consumed: u64,
    pub logs: Vec<String>,
}

pub struct Runtime {
    program_id: Address,
    instructions: Vec<Instruction>,
    rodata: Vec<u8>,
    entrypoint: usize,
    programs: HashMap<Address, Vec<u8>>,
    config: RuntimeConfig,
    sysvars: SysvarContext,
    vm: Option<SbpfVm<RuntimeSyscallHandler>>,
    accounts: HashMap<Address, Account>,
    account_metas: Vec<AccountMeta>,
    pre_lens: Vec<usize>, // original account data lengths at serialization
    log_collector: LogCollector,
}

impl Runtime {
    pub fn new(
        program_id: Address,
        elf: impl Into<ElfSource>,
        config: RuntimeConfig,
    ) -> RuntimeResult<Self> {
        let elf_bytes = match elf.into() {
            ElfSource::Path(path) => std::fs::read(&path)?,
            ElfSource::Bytes(bytes) => bytes,
        };

        let (instructions, rodata, entrypoint) = load_elf(&elf_bytes)?;

        Ok(Self {
            program_id,
            instructions,
            rodata,
            entrypoint,
            programs: HashMap::new(),
            config,
            sysvars: SysvarContext::default(),
            vm: None,
            accounts: HashMap::new(),
            account_metas: Vec::new(),
            pre_lens: Vec::new(),
            log_collector: Rc::new(RefCell::new(Vec::new())),
        })
    }

    pub fn add_program(&mut self, program_id: &Address, elf: impl Into<ElfSource>) {
        let elf_bytes = match elf.into() {
            ElfSource::Path(path) => std::fs::read(&path).expect("Failed to read ELF"),
            ElfSource::Bytes(bytes) => bytes,
        };
        self.programs.insert(*program_id, elf_bytes);
    }

    fn setup_vm(
        &mut self,
        instruction: &SolanaInstruction,
        accounts: &[(Address, Account)],
    ) -> RuntimeResult<()> {
        // Setup accounts (merge with existing account state).
        for (address, account) in accounts.iter() {
            self.accounts
                .entry(*address)
                .or_insert_with(|| account.clone());
        }
        self.account_metas = instruction.accounts.clone();

        let (input, pre_lens, instruction_data_offset) = serialize::serialize_parameters(
            &self.accounts,
            &self.account_metas,
            &instruction.data,
            &self.program_id,
        )?;

        let vm_config = SbpfVmConfig {
            compute_unit_limit: self.config.compute_budget,
            max_call_depth: self.config.max_call_depth,
            heap_size: self.config.heap_size,
        };

        let handler = RuntimeSyscallHandler::new(
            ExecutionCost::default(),
            self.program_id,
            self.sysvars.clone(),
            self.log_collector.clone(),
        );

        let mut vm = SbpfVm::new_with_config(
            self.instructions.clone(),
            input,
            self.rodata.clone(),
            handler,
            vm_config,
        );
        vm.compute_meter = ComputeMeter::new(self.config.compute_budget);
        vm.set_entrypoint(self.entrypoint);
        vm.registers[2] = Memory::INPUT_START + instruction_data_offset as u64;

        self.pre_lens = pre_lens;
        self.vm = Some(vm);
        Ok(())
    }

    fn sync_accounts(&mut self) -> RuntimeResult<()> {
        if let Some(ref vm) = self.vm {
            serialize::deserialize_parameters(
                &mut self.accounts,
                &self.account_metas,
                &vm.memory.input,
                &self.pre_lens,
                &self.program_id,
            )?;
        }
        Ok(())
    }

    pub fn run(
        &mut self,
        instruction: &SolanaInstruction,
        accounts: &[(Address, Account)],
    ) -> RuntimeResult<ExecutionResult> {
        self.log_collector.borrow_mut().clear();
        self.setup_vm(instruction, accounts)?;

        // Get pre-execution lamports from the account state.
        let pre_lamports: HashMap<Address, u64> = self
            .account_metas
            .iter()
            .filter_map(|meta| {
                self.accounts
                    .get(&meta.pubkey)
                    .map(|a| (meta.pubkey, a.lamports))
            })
            .collect();

        self.log_collector
            .borrow_mut()
            .push(format!("Program {} invoke [1]", self.program_id));

        loop {
            let vm = self.vm.as_mut().unwrap();
            if let Err(e) = vm.step() {
                self.log_collector
                    .borrow_mut()
                    .push(format!("Program failed: {}", e));
                return Err(e.into());
            }

            if let Some(request) = vm.syscall_handler.pending_cpi.take() {
                if let Err(e) = self.handle_cpi(request) {
                    self.log_collector
                        .borrow_mut()
                        .push(format!("Program failed: {}", e));
                    return Err(e);
                }
                continue;
            }

            if vm.halted {
                break;
            }
        }

        self.sync_accounts()?;

        // Verify total lamport balance is conserved across all instruction accounts.
        let pre_total: u64 = pre_lamports.values().sum();
        let post_total: u64 = pre_lamports
            .keys()
            .filter_map(|pk| self.accounts.get(pk))
            .map(|a| a.lamports)
            .sum();
        if pre_total != post_total {
            return Err(RuntimeError::UnbalancedInstruction(pre_total, post_total));
        }

        let vm = self.vm.as_ref().unwrap();
        let consumed = vm.compute_meter.get_consumed();
        let exit_code = vm.exit_code;

        if let Some(ref return_data) = vm.syscall_handler.return_data
            && !return_data.1.is_empty()
        {
            self.log_collector.borrow_mut().push(format!(
                "Program return: {} {}",
                return_data.0,
                BASE64.encode(&return_data.1)
            ));
        }

        self.log_collector.borrow_mut().push(format!(
            "Program {} consumed {} of {} compute units",
            self.program_id, consumed, self.config.compute_budget
        ));

        if exit_code.unwrap_or(0) == 0 {
            self.log_collector
                .borrow_mut()
                .push(format!("Program {} success", self.program_id));
        } else {
            self.log_collector.borrow_mut().push(format!(
                "Program {} failed: exit code {}",
                self.program_id,
                exit_code.unwrap_or(0)
            ));
        }

        let logs = self.log_collector.borrow().clone();

        Ok(ExecutionResult {
            exit_code,
            compute_units_consumed: consumed,
            logs,
        })
    }

    pub fn prepare(
        &mut self,
        instruction: &SolanaInstruction,
        accounts: &[(Address, Account)],
    ) -> RuntimeResult<()> {
        self.log_collector.borrow_mut().clear();
        self.setup_vm(instruction, accounts)?;
        self.log_collector
            .borrow_mut()
            .push(format!("Program {} invoke [1]", self.program_id));
        Ok(())
    }

    pub fn step(&mut self) -> RuntimeResult<()> {
        let vm = self.vm.as_mut().ok_or(RuntimeError::VmNotPrepared)?;
        if vm.halted {
            return Ok(());
        }
        if let Err(e) = vm.step() {
            self.log_collector
                .borrow_mut()
                .push(format!("Program failed: {}", e));
            return Err(e.into());
        }

        if let Some(request) = vm.syscall_handler.pending_cpi.take()
            && let Err(e) = self.handle_cpi(request)
        {
            self.log_collector
                .borrow_mut()
                .push(format!("Program failed: {}", e));
            return Err(e);
        }

        let vm_ref = self.vm.as_ref().unwrap();
        if vm_ref.halted {
            self.sync_accounts()?;

            let vm = self.vm.as_ref().unwrap();
            let consumed = vm.compute_meter.get_consumed();
            let exit_code = vm.exit_code;

            if let Some(ref return_data) = vm.syscall_handler.return_data
                && !return_data.1.is_empty()
            {
                self.log_collector.borrow_mut().push(format!(
                    "Program return: {} {}",
                    return_data.0,
                    BASE64.encode(&return_data.1)
                ));
            }

            self.log_collector.borrow_mut().push(format!(
                "Program {} consumed {} of {} compute units",
                self.program_id, consumed, self.config.compute_budget
            ));

            if exit_code.unwrap_or(0) == 0 {
                self.log_collector
                    .borrow_mut()
                    .push(format!("Program {} success", self.program_id));
            } else {
                self.log_collector.borrow_mut().push(format!(
                    "Program {} failed: exit code {}",
                    self.program_id,
                    exit_code.unwrap_or(0)
                ));
            }
        }

        Ok(())
    }

    fn handle_cpi(&mut self, request: cpi::request::CpiRequest) -> RuntimeResult<()> {
        let vm = self.vm.as_ref().unwrap();
        let compute_remaining = self.config.compute_budget - vm.compute_meter.get_consumed();

        // Sync latest account state from caller VM memory into account store.
        cpi::sync::sync_from_caller(&vm.memory, &request.caller_accounts, &mut self.accounts)?;

        let caller_accounts = request.caller_accounts;
        let cpi_request = cpi::request::CpiRequest {
            program_id: request.program_id,
            accounts: request.accounts,
            data: request.data,
            caller_accounts: Vec::new(),
            signers: request.signers,
        };

        let mut ctx = cpi::CpiContext {
            request: cpi_request,
            programs: &self.programs,
            accounts: &mut self.accounts,
            config: &self.config,
            sysvars: &self.sysvars,
            compute_remaining,
            cpi_depth: 1,
            caller_account_metas: &self.account_metas,
            log_collector: &self.log_collector,
        };

        let output = cpi::execute_cpi(&mut ctx)?;

        let vm = self.vm.as_mut().unwrap();
        vm.compute_meter.consume(output.compute_consumed)?;
        vm.syscall_handler.return_data = output.return_data;

        if output.exit_code != 0 {
            return Err(RuntimeError::VmError(
                sbpf_vm::errors::SbpfVmError::SyscallError(format!(
                    "CPI callee returned error: {}",
                    output.exit_code
                )),
            ));
        }

        // Sync updated accounts back to caller VM memory.
        let vm = self.vm.as_mut().unwrap();
        cpi::sync::sync_to_caller(&mut vm.memory, &caller_accounts, &self.accounts)?;

        Ok(())
    }

    pub fn get_pc(&self) -> usize {
        self.vm.as_ref().map(|vm| vm.pc).unwrap_or(0)
    }

    pub fn get_registers(&self) -> Option<&[u64; 11]> {
        self.vm.as_ref().map(|vm| &vm.registers)
    }

    pub fn current_program_id(&self) -> &Address {
        &self.program_id
    }

    pub fn is_halted(&self) -> bool {
        self.vm.as_ref().map(|vm| vm.halted).unwrap_or(false)
    }

    pub fn exit_code(&self) -> Option<u64> {
        self.vm.as_ref().and_then(|vm| vm.exit_code)
    }

    pub fn compute_units_consumed(&self) -> u64 {
        self.vm
            .as_ref()
            .map(|vm| vm.compute_meter.get_consumed())
            .unwrap_or(0)
    }

    pub fn get_account(&self, pubkey: &Address) -> Option<Account> {
        self.accounts.get(pubkey).cloned()
    }

    pub fn get_accounts(&self) -> &HashMap<Address, Account> {
        &self.accounts
    }

    pub fn get_register(&self, idx: usize) -> Option<u64> {
        self.vm
            .as_ref()
            .and_then(|vm| vm.registers.get(idx).copied())
    }

    pub fn set_register(&mut self, idx: usize, value: u64) -> RuntimeResult<()> {
        let vm = self.vm.as_mut().ok_or(RuntimeError::VmNotPrepared)?;
        if idx >= vm.registers.len() {
            return Err(RuntimeError::RegisterOutOfRange(idx));
        }
        vm.set_register(idx, value);
        Ok(())
    }

    pub fn read_memory(&self, addr: u64, size: usize) -> Option<Vec<u8>> {
        self.vm
            .as_ref()
            .and_then(|vm| vm.memory.read_bytes(addr, size).ok().map(|s| s.to_vec()))
    }

    pub fn get_instruction(&self) -> Option<&Instruction> {
        let vm = self.vm.as_ref()?;
        vm.program.get(vm.pc)
    }

    pub fn get_program(&self) -> &[Instruction] {
        &self.instructions
    }

    pub fn get_call_stack(&self) -> Option<&[CallFrame]> {
        self.vm.as_ref().map(|vm| vm.call_stack.as_slice())
    }

    pub fn config(&self) -> &RuntimeConfig {
        &self.config
    }

    pub fn sysvars(&self) -> &SysvarContext {
        &self.sysvars
    }

    pub fn sysvars_mut(&mut self) -> &mut SysvarContext {
        &mut self.sysvars
    }

    pub fn log_collector(&self) -> &LogCollector {
        &self.log_collector
    }

    pub fn drain_logs(&self) -> Vec<String> {
        self.log_collector.borrow_mut().drain(..).collect()
    }
}

#[cfg(test)]
mod tests {
    use {super::*, std::path::PathBuf};

    const PROGRAM_ID: Address =
        Address::from_str_const("22222222222222222222222222222222222222222222");

    fn escrow_elf_path() -> String {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("tests/fixtures/libupstream_pinocchio_escrow.so")
            .to_str()
            .unwrap()
            .to_string()
    }

    fn new_runtime() -> Runtime {
        Runtime::new(
            PROGRAM_ID,
            escrow_elf_path().as_str(),
            RuntimeConfig::default(),
        )
        .unwrap()
    }

    fn empty_instruction() -> SolanaInstruction {
        SolanaInstruction {
            program_id: PROGRAM_ID,
            accounts: Vec::new(),
            data: Vec::new(),
        }
    }

    #[test]
    fn elf_source_from_str() {
        let path = escrow_elf_path();
        assert!(matches!(ElfSource::from(path.as_str()), ElfSource::Path(_)));
    }

    #[test]
    fn elf_source_from_bytes_slice() {
        let bytes: &[u8] = &[1, 2, 3];
        assert!(matches!(ElfSource::from(bytes), ElfSource::Bytes(_)));
    }

    #[test]
    fn elf_source_from_vec() {
        assert!(matches!(
            ElfSource::from(vec![1u8, 2, 3]),
            ElfSource::Bytes(_)
        ));
    }

    #[test]
    fn new_from_path() {
        let rt = new_runtime();
        assert_eq!(*rt.current_program_id(), PROGRAM_ID);
        assert!(!rt.get_program().is_empty());
    }

    #[test]
    fn new_from_bytes() {
        let bytes = std::fs::read(escrow_elf_path()).unwrap();
        let rt = Runtime::new(PROGRAM_ID, bytes, RuntimeConfig::default()).unwrap();
        assert!(!rt.get_program().is_empty());
    }

    #[test]
    fn new_bad_path_errors() {
        match Runtime::new(
            PROGRAM_ID,
            "/nonexistent/path/to/program.so",
            RuntimeConfig::default(),
        ) {
            Err(RuntimeError::ElfReadError(_)) => {}
            other => panic!("expected ElfReadError, got {:?}", other.map(|_| ())),
        }
    }

    #[test]
    fn new_invalid_elf_errors() {
        match Runtime::new(PROGRAM_ID, vec![0u8; 8], RuntimeConfig::default()) {
            Err(RuntimeError::ElfReadError(_)) => panic!("expected parse error, got io error"),
            Err(_) => {}
            Ok(_) => panic!("expected error for invalid ELF"),
        }
    }

    #[test]
    fn add_program_accepts_path_and_bytes() {
        let mut rt = new_runtime();
        rt.add_program(&Address::new_unique(), escrow_elf_path().as_str());
        let bytes = std::fs::read(escrow_elf_path()).unwrap();
        rt.add_program(&Address::new_unique(), bytes);
    }

    #[test]
    fn getters_before_prepare_are_defaults() {
        let rt = new_runtime();
        assert_eq!(rt.get_pc(), 0);
        assert!(rt.get_registers().is_none());
        assert!(rt.get_register(2).is_none());
        assert!(rt.read_memory(Memory::INPUT_START, 8).is_none());
        assert!(rt.get_instruction().is_none());
        assert!(rt.get_call_stack().is_none());
        assert!(!rt.is_halted());
        assert!(rt.exit_code().is_none());
        assert_eq!(rt.compute_units_consumed(), 0);
        assert!(rt.get_account(&PROGRAM_ID).is_none());
        assert!(rt.get_accounts().is_empty());
        assert!(rt.drain_logs().is_empty());
    }

    #[test]
    fn set_register_before_prepare_errors() {
        let mut rt = new_runtime();
        let err = rt.set_register(0, 42).unwrap_err();
        assert!(matches!(err, RuntimeError::VmNotPrepared));
    }

    #[test]
    fn step_before_prepare_errors() {
        let mut rt = new_runtime();
        let err = rt.step().unwrap_err();
        assert!(matches!(err, RuntimeError::VmNotPrepared));
    }

    #[test]
    fn config_and_sysvar_accessors() {
        let mut rt = new_runtime();
        assert_eq!(
            rt.config().compute_budget,
            RuntimeConfig::default().compute_budget
        );
        let slot = rt.sysvars().clock.slot;
        rt.sysvars_mut().clock.slot = slot + 5;
        assert_eq!(rt.sysvars().clock.slot, slot + 5);
        assert!(rt.log_collector().borrow().is_empty());
    }

    #[test]
    fn prepare_initializes_vm_state() {
        let mut rt = new_runtime();
        rt.prepare(&empty_instruction(), &[]).unwrap();

        // After prepare the VM exists: getters now return Some / live values.
        assert!(rt.get_registers().is_some());
        assert!(rt.get_call_stack().is_some());
        assert!(rt.get_instruction().is_some());
        let r2 = rt.get_register(2).unwrap();
        assert!(r2 >= Memory::INPUT_START);
        assert!(rt.read_memory(Memory::INPUT_START, 8).is_some());
        let logs = rt.log_collector().borrow().clone();
        assert!(logs.iter().any(|l| l.contains("invoke [1]")));
    }

    #[test]
    fn set_register_after_prepare_and_out_of_range() {
        let mut rt = new_runtime();
        rt.prepare(&empty_instruction(), &[]).unwrap();

        rt.set_register(3, 0xdead_beef).unwrap();
        assert_eq!(rt.get_register(3), Some(0xdead_beef));

        let err = rt.set_register(99, 1).unwrap_err();
        assert!(matches!(err, RuntimeError::RegisterOutOfRange(99)));
    }

    #[test]
    fn step_drives_program_to_clean_halt() {
        let mut rt = new_runtime();
        rt.prepare(&empty_instruction(), &[]).unwrap();

        let mut steps = 0;
        while !rt.is_halted() {
            rt.step().expect("step should not trap on this program");
            steps += 1;
            assert!(steps <= 100_000, "program did not terminate");
        }
        assert!(steps > 0);
        assert!(rt.is_halted());
        assert!(rt.exit_code().is_some());
    }

    #[test]
    fn run_returns_ok_with_nonzero_exit_and_logs() {
        let mut rt = new_runtime();
        let exec = rt.run(&empty_instruction(), &[]).unwrap();

        assert!(matches!(exec.exit_code, Some(code) if code != 0));
        assert!(exec.compute_units_consumed > 0);

        assert!(exec.logs.iter().any(|l| l.contains("invoke [1]")));
        assert!(exec.logs.iter().any(|l| l.contains("consumed")));
        assert!(exec.logs.iter().any(|l| l.contains("failed: exit code")));
    }
}