casper_wasm_utils/

stack_height.rs

//! The pass that tries to make stack overflows deterministic, by introducing
//! an upper bound of the stack size.
//!
//! This pass introduces a global mutable variable to track stack height,
//! and instruments all calls with preamble and postamble.
//!
//! Stack height is increased prior the call. Otherwise, the check would
//! be made after the stack frame is allocated.
//!
//! The preamble is inserted before the call. It increments
//! the global stack height variable with statically determined "stack cost"
//! of the callee. If after the increment the stack height exceeds
//! the limit (specified by the `rules`) then execution traps.
//! Otherwise, the call is executed.
//!
//! The postamble is inserted after the call. The purpose of the postamble is to decrease
//! the stack height by the "stack cost" of the callee function.
//!
//! Note, that we can't instrument all possible ways to return from the function. The simplest
//! example would be a trap issued by the host function.
//! That means stack height global won't be equal to zero upon the next execution after such trap.
//!
//! # Thunks
//!
//! Because stack height is increased prior the call few problems arises:
//!
//! - Stack height isn't increased upon an entry to the first function, i.e. exported function.
//! - Start function is executed externally (similar to exported functions).
//! - It is statically unknown what function will be invoked in an indirect call.
//!
//! The solution for this problems is to generate a intermediate functions, called 'thunks', which
//! will increase before and decrease the stack height after the call to original function, and
//! then make exported function and table entries, start section to point to a corresponding thunks.
//!
//! # Stack cost
//!
//! Stack cost of the function is calculated as a sum of it's locals
//! and the maximal height of the value stack.
//!
//! All values are treated equally, as they have the same size.
//!
//! The rationale is that this makes it possible to use the following very naive wasm executor:
//!
//! - values are implemented by a union, so each value takes a size equal to
//!   the size of the largest possible value type this union can hold. (In MVP it is 8 bytes)
//! - each value from the value stack is placed on the native stack.
//! - each local variable and function argument is placed on the native stack.
//! - arguments pushed by the caller are copied into callee stack rather than shared
//!   between the frames.
//! - upon entry into the function entire stack frame is allocated.

use crate::std::{mem, string::String, vec::Vec};

use casper_wasm::{
    builder,
    elements::{self, Instruction, Instructions, Type},
};

/// Const function to generate preamble and postamble.
const fn instrument_call(
    callee_idx: u32,
    callee_stack_cost: u32,
    stack_height_global_idx: u32,
    stack_limit: u32,
) -> [Instruction; INSTRUMENT_CALL_LENGTH] {
    use casper_wasm::elements::Instruction::*;
    [
        // stack_height += stack_cost(F)
        GetGlobal(stack_height_global_idx),
        I32Const(callee_stack_cost as i32),
        I32Add,
        SetGlobal(stack_height_global_idx),
        // if stack_counter > LIMIT: unreachable
        GetGlobal(stack_height_global_idx),
        I32Const(stack_limit as i32),
        I32GtU,
        If(elements::BlockType::NoResult),
        Unreachable,
        End,
        // Original call
        Call(callee_idx),
        // stack_height -= stack_cost(F)
        GetGlobal(stack_height_global_idx),
        I32Const(callee_stack_cost as i32),
        I32Sub,
        SetGlobal(stack_height_global_idx),
    ]
}

const INSTRUMENT_CALL_LENGTH: usize = 15;

mod max_height;
mod thunk;

/// Error that occured during processing the module.
///
/// This means that the module is invalid.
#[derive(Debug)]
#[allow(unused)]
pub struct Error(String);

pub(crate) struct Context {
    stack_height_global_idx: u32,
    func_stack_costs: Vec<u32>,
    stack_limit: u32,
}

impl Context {
    /// Returns index in a global index space of a stack_height global variable.
    fn stack_height_global_idx(&self) -> u32 {
        self.stack_height_global_idx
    }

    /// Returns `stack_cost` for `func_idx`.
    fn stack_cost(&self, func_idx: u32) -> Option<u32> {
        self.func_stack_costs.get(func_idx as usize).cloned()
    }

    /// Returns stack limit specified by the rules.
    fn stack_limit(&self) -> u32 {
        self.stack_limit
    }
}

/// Instrument a module with stack height limiter.
///
/// See module-level documentation for more details.
///
/// # Errors
///
/// Returns `Err` if module is invalid and can't be
pub fn inject_limiter(
    mut module: elements::Module,
    stack_limit: u32,
) -> Result<elements::Module, Error> {
    let mut ctx = Context {
        stack_height_global_idx: generate_stack_height_global(&mut module),
        func_stack_costs: compute_stack_costs(&module)?,
        stack_limit,
    };

    instrument_functions(&mut ctx, &mut module)?;
    let module = thunk::generate_thunks(&mut ctx, module)?;

    Ok(module)
}

/// Generate a new global that will be used for tracking current stack height.
fn generate_stack_height_global(module: &mut elements::Module) -> u32 {
    let global_entry = builder::global()
        .value_type()
        .i32()
        .mutable()
        .init_expr(Instruction::I32Const(0))
        .build();

    // Try to find an existing global section.
    for section in module.sections_mut() {
        if let elements::Section::Global(gs) = section {
            gs.entries_mut().push(global_entry);
            return (gs.entries().len() as u32) - 1;
        }
    }

    // Existing section not found, create one!
    module.sections_mut().push(elements::Section::Global(
        elements::GlobalSection::with_entries(vec![global_entry]),
    ));
    0
}

/// Calculate stack costs for all functions.
///
/// Returns a vector with a stack cost for each function, including imports.
fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, Error> {
    let func_imports = module.import_count(elements::ImportCountType::Function);

    // TODO: optimize!
    (0..module.functions_space())
        .map(|func_idx| {
            if func_idx < func_imports {
                // We can't calculate stack_cost of the import functions.
                Ok(0)
            } else {
                compute_stack_cost(func_idx as u32, module)
            }
        })
        .collect()
}

/// Stack cost of the given *defined* function is the sum of it's locals count (that is,
/// number of arguments plus number of local variables) and the maximal stack
/// height.
fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, Error> {
    // To calculate the cost of a function we need to convert index from
    // function index space to defined function spaces.
    let func_imports = module.import_count(elements::ImportCountType::Function) as u32;
    let defined_func_idx = func_idx
        .checked_sub(func_imports)
        .ok_or_else(|| Error("This should be a index of a defined function".into()))?;

    let code_section = module
        .code_section()
        .ok_or_else(|| Error("Due to validation code section should exists".into()))?;
    let body = &code_section
        .bodies()
        .get(defined_func_idx as usize)
        .ok_or_else(|| Error("Function body is out of bounds".into()))?;

    let mut locals_count: u32 = 0;
    for local_group in body.locals() {
        locals_count = locals_count
            .checked_add(local_group.count())
            .ok_or_else(|| Error("Overflow in local count".into()))?;
    }

    let max_stack_height = max_height::compute(defined_func_idx, module)?;

    locals_count
        .checked_add(max_stack_height)
        .ok_or_else(|| Error("Overflow in adding locals_count and max_stack_height".into()))
}

fn instrument_functions(ctx: &mut Context, module: &mut elements::Module) -> Result<(), Error> {
    for section in module.sections_mut() {
        if let elements::Section::Code(code_section) = section {
            for func_body in code_section.bodies_mut() {
                let opcodes = func_body.code_mut();
                instrument_function(ctx, opcodes)?;
            }
        }
    }
    Ok(())
}

/// This function searches `call` instructions and wrap each call
/// with preamble and postamble.
///
/// Before:
///
/// ```text
/// get_local 0
/// get_local 1
/// call 228
/// drop
/// ```
///
/// After:
///
/// ```text
/// get_local 0
/// get_local 1
///
/// < ... preamble ... >
///
/// call 228
///
/// < .. postamble ... >
///
/// drop
/// ```
fn instrument_function(ctx: &mut Context, func: &mut Instructions) -> Result<(), Error> {
    use Instruction::*;

    struct InstrumentCall {
        offset: usize,
        callee: u32,
        cost: u32,
    }

    let calls: Vec<_> = func
        .elements()
        .iter()
        .enumerate()
        .filter_map(|(offset, instruction)| {
            if let Call(callee) = *instruction {
                ctx.stack_cost(callee)
                    .filter(|&cost| cost > 0)
                    .map(|cost| InstrumentCall {
                        callee,
                        offset,
                        cost,
                    })
            } else {
                None
            }
        })
        .collect();

    // The `instrumented_call!` contains the call itself. This is why we need to subtract one.
    let len = func.elements().len() + calls.len() * (INSTRUMENT_CALL_LENGTH - 1);
    let original_instrs = mem::replace(func.elements_mut(), Vec::with_capacity(len));
    let new_instrs = func.elements_mut();

    let mut calls = calls.into_iter().peekable();
    for (original_pos, instr) in original_instrs.into_iter().enumerate() {
        // whether there is some call instruction at this position that needs to be instrumented
        let did_instrument = if let Some(call) = calls.peek() {
            if call.offset == original_pos {
                let new_seq = instrument_call(
                    call.callee,
                    call.cost,
                    ctx.stack_height_global_idx(),
                    ctx.stack_limit(),
                );
                new_instrs.extend(new_seq);
                true
            } else {
                false
            }
        } else {
            false
        };

        if did_instrument {
            calls.next();
        } else {
            new_instrs.push(instr);
        }
    }

    if calls.next().is_some() {
        return Err(Error("Not all calls were used".into()));
    }

    Ok(())
}

fn resolve_func_type(
    func_idx: u32,
    module: &elements::Module,
) -> Result<&elements::FunctionType, Error> {
    let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
    let functions = module
        .function_section()
        .map(|fs| fs.entries())
        .unwrap_or(&[]);

    let func_imports = module.import_count(elements::ImportCountType::Function);
    let sig_idx = if func_idx < func_imports as u32 {
        module
            .import_section()
            .expect("function import count is not zero; import section must exists; qed")
            .entries()
            .iter()
            .filter_map(|entry| match entry.external() {
                elements::External::Function(idx) => Some(*idx),
                _ => None,
            })
            .nth(func_idx as usize)
            .expect(
                "func_idx is less than function imports count;
				nth function import must be `Some`;
				qed",
            )
    } else {
        functions
            .get(func_idx as usize - func_imports)
            .ok_or_else(|| Error(format!("Function at index {} is not defined", func_idx)))?
            .type_ref()
    };
    let Type::Function(ty) = types.get(sig_idx as usize).ok_or_else(|| {
        Error(format!(
            "Signature {} (specified by func {}) isn't defined",
            sig_idx, func_idx
        ))
    })?;
    Ok(ty)
}

#[cfg(test)]
mod tests {
    use super::*;
    use casper_wasm::elements;

    fn parse_wat(source: &str) -> elements::Module {
        elements::deserialize_buffer(&wabt::wat2wasm(source).expect("Failed to wat2wasm"))
            .expect("Failed to deserialize the module")
    }

    fn validate_module(module: elements::Module) {
        let binary = elements::serialize(module).expect("Failed to serialize");
        wabt::Module::read_binary(binary, &Default::default())
            .expect("Wabt failed to read final binary")
            .validate()
            .expect("Invalid module");
    }

    #[test]
    fn test_with_params_and_result() {
        let module = parse_wat(
            r#"
(module
	(func (export "i32.add") (param i32 i32) (result i32)
		get_local 0
	get_local 1
	i32.add
	)
)
"#,
        );

        let module = inject_limiter(module, 1024).expect("Failed to inject stack counter");
        validate_module(module);
    }
}