rave_engine 0.3.5

pub mod resource_monitor;
use resource_monitor::*;
pub mod error;
use error::*;
pub mod rhai_functions;
use super::types::entries::rave::rave_output::RAVEOutput;
use hdi::prelude::*;
use rhai::{Array, Dynamic, Engine, Map};
use rhai_functions::prelude::*;
use serde_json::{Map as JsonMap, Value};
use std::sync::{Arc, Mutex};

#[cfg(test)]
pub mod tests;

pub struct RhaiEngine {
    config: RhaiEngineConfig,
    resource_monitor: Option<Arc<Mutex<ResourceMonitor>>>,
}

impl Default for RhaiEngine {
    fn default() -> Self {
        Self::new()
    }
}

pub struct PresetVariables {
    pub ea_id: ActionHash,
    pub executor: AgentPubKey,
    pub executed_timestamp: Timestamp,
}

impl RhaiEngine {
    pub fn new() -> Self {
        let resource_monitor = match ResourceMonitor::new(TimeProvider::Holochain, 0, 0) {
            Ok(monitor) => Some(Arc::new(Mutex::new(monitor))),
            Err(_) => {
                // If we can't create the resource monitor, disable it
                // This is a fallback to ensure the engine can still be created
                None
            }
        };

        Self {
            config: RhaiEngineConfig::default(),
            resource_monitor,
        }
    }

    pub fn with_config(config: RhaiEngineConfig) -> ExternResult<Self> {
        let mut engine = Engine::new();
        Self::register_functions(&mut engine);

        if config.enable_sandbox {
            engine.set_max_operations(config.max_operations);
        }

        let time_provider = if config.use_holochain_time {
            TimeProvider::Holochain
        } else {
            TimeProvider::System
        };

        let resource_monitor = if config.enable_sandbox {
            Some(Arc::new(Mutex::new(
                ResourceMonitor::new(
                    time_provider,
                    config.max_operations,
                    config.max_recursion_depth,
                )
                .map_err(|e| wasm_error!(WasmErrorInner::Guest(e.to_string())))?,
            )))
        } else {
            None
        };

        Ok(Self {
            config,
            resource_monitor,
        })
    }

    fn register_functions(engine: &mut Engine) {
        // Parsers
        engine.register_fn("parse_float", parse_float);
        engine.register_fn("add_fuel", add_fuel);
        engine.register_fn("sub_fuel", sub_fuel);
        engine.register_fn("add_units", add_units);
        engine.register_fn("sub_units", sub_units);
        engine.register_fn("is_greater_than_eq", is_greater_than_eq);
        engine.register_fn("extract_indexes", extract_indexes);
        // DNA Helper Functions
        // engine.register_fn("verify_valid_spend_records", verify_valid_spend_records);
        engine.register_fn(
            "parse_record_to_parked_amount_and_source",
            parse_record_to_parked_amount_and_source,
        );
        engine.register_fn("acceding_sort_allocation", acceding_sort_allocation);
        engine.register_fn("get_data_blob", get_data_blob);
        engine.register_fn("get_spend_links_author", get_spend_links_author);
        // HDI Helper Functions
        engine.register_fn("hdi_verify_signature", hdi_verify_signature);
        engine.register_fn("hdi_verify_record_signature", hdi_verify_record_signature);
        engine.register_fn("check_cool_down_period", check_cool_down_period);
    }

    pub fn execute(
        &self,
        input: &serde_json::Value,
        code: Vec<u8>,
        preset_variables: Option<PresetVariables>,
    ) -> ExternResult<RhaiEngineOutput> {
        unyt_info!("RAVE", "Rhai engine execution started");
        if input.to_string().len() as u64 > self.config.max_script_size {
            return Err(wasm_error!(WasmErrorInner::Guest(
                RhaiError::InputTooLarge.to_string()
            )));
        }
        // trace!("Input size: {:?}", input.to_string().len() as u64);
        if code.len() as u64 > self.config.max_script_size {
            return Err(wasm_error!(WasmErrorInner::Guest(
                RhaiError::ScriptError("Script size limit exceeded".to_string()).to_string()
            )));
        }
        // debug!("Code size: {:?}", code.len() as u64);
        let mut engine = Engine::new();
        Self::register_functions(&mut engine);
        // debug!("Registering functions");
        // Set recursion depth limit
        engine.set_max_expr_depths(
            self.config.max_recursion_depth as usize,
            self.config.max_recursion_depth as usize,
        );
        // debug!("Setting max operations");
        engine.set_max_operations(self.config.max_operations);
        // debug!("Setting max modules");
        engine.set_max_modules(0);
        // debug!("Setting max string size");
        engine.set_max_string_size(1024 * 1024);

        // Register operation tracking
        if let Some(monitor) = &self.resource_monitor {
            let monitor = monitor.clone();
            engine.on_progress(move |_| {
                // Handle potential mutex poisoning gracefully
                if let Ok(mut monitor) = monitor.lock() {
                    monitor.increment_operation_count().ok();
                }
                // If mutex is poisoned, we continue without tracking
                // This prevents the entire engine from failing due to monitoring issues
                None
            });
        }
        trace!("Rhai Engine Input: {:?}", input);
        let mut scope = Self::input_to_scope(input);
        if let Some(preset_variables) = preset_variables {
            scope.set_value("ea_id", preset_variables.ea_id.to_string());
            scope.set_value("executor_pub_key", preset_variables.executor.to_string());
            scope.set_value(
                "executed_timestamp",
                preset_variables.executed_timestamp.to_string(),
            );
        }
        // debug!("Scope: {:?}", scope);
        // Parse and execute the code
        let code: String = rmp_serde::from_slice(&code)
            .map_err(|e| wasm_error!("Unable to parse execution code: {}", e))?;
        // debug!("Executing code: {:?}", code);

        let ast = match engine.compile(&code) {
            Ok(ast) => ast,
            Err(e) => {
                unyt_warn!("RAVE", "Rhai engine compilation error | error={:?}", e);
                let error_str = e.to_string();
                if error_str.contains("execution timeout") {
                    return Err(wasm_error!(WasmErrorInner::Guest(
                        RhaiError::Timeout(Timestamp::from_micros(0)).to_string()
                    )));
                } else if error_str.contains("Expression exceeds maximum complexity") {
                    return Err(wasm_error!(WasmErrorInner::Guest(
                        RhaiError::RecursionDepthExceeded(self.config.max_recursion_depth)
                            .to_string()
                    )));
                } else if error_str.contains("Too many operations") {
                    return Err(wasm_error!(WasmErrorInner::Guest(
                        RhaiError::OperationLimitExceeded(self.config.max_operations).to_string()
                    )));
                }
                return Err(wasm_error!(WasmErrorInner::Guest(
                    RhaiError::SyntaxError(e.to_string()).to_string()
                )));
            }
        };

        // debug!("Ast: {:?}", ast);
        let result = engine.eval_ast_with_scope::<Dynamic>(&mut scope, &ast);
        trace!("Rhai Engine Result: {:?}", result);
        match result {
            Ok(output) => {
                unyt_info!("RAVE", "Rhai engine execution completed successfully");
                let json_output = Self::rhai_dynamic_to_json(output);
                RhaiEngineOutput::try_from(json_output)
                    .map_err(|e| wasm_error!(WasmErrorInner::Guest(e)))
            }
            Err(e) => {
                let error_str = e.to_string();
                if error_str.contains("Too many modules imported") {
                    return Err(wasm_error!(WasmErrorInner::Guest(
                        RhaiError::ModuleImportError.to_string()
                    )));
                }

                Err(wasm_error!(format!("Rhai Engine Execution Error: {:?}", e)))
            }
        }
    }

    fn input_to_scope(input: &serde_json::Value) -> rhai::Scope<'_> {
        let mut scope = rhai::Scope::new();
        if let Some(object) = input.as_object() {
            for (key, value) in object.clone() {
                match value {
                    Value::Null => {
                        scope.set_value(key, Dynamic::UNIT);
                    }
                    Value::Bool(b) => {
                        scope.set_value(key, b);
                    }
                    Value::Number(n) => {
                        if let Some(i) = n.as_i64() {
                            scope.set_value(key, i);
                        } else if let Some(f) = n.as_f64() {
                            scope.set_value(key, f);
                        }
                    }
                    Value::String(s) => {
                        scope.set_value(key, s);
                    }
                    Value::Array(arr) => {
                        let rhai_arr = Self::json_to_rhai_array(&Value::Array(arr.clone()));
                        scope.set_value(key, rhai_arr);
                    }
                    Value::Object(_) => {
                        let nested_scope = Self::json_to_rhai_map(&value);
                        scope.set_value(key, nested_scope);
                    }
                }
            }
        }
        scope
    }

    fn json_to_rhai_array(json: &Value) -> Array {
        let mut rhai_array = Array::new();

        if let Value::Array(arr) = json {
            for item in arr {
                let dynamic_item = Self::json_to_rhai_dynamic(item);
                rhai_array.push(dynamic_item);
            }
        }

        rhai_array
    }

    fn json_to_rhai_dynamic(value: &Value) -> Dynamic {
        match value {
            Value::Number(n) => {
                if let Some(i) = n.as_i64() {
                    Dynamic::from(i)
                } else if let Some(f) = n.as_f64() {
                    Dynamic::from(f)
                } else {
                    Dynamic::UNIT
                }
            }
            Value::String(s) => Dynamic::from(s.clone()),
            Value::Bool(b) => Dynamic::from(*b),
            Value::Array(arr) => {
                Dynamic::from(Self::json_to_rhai_array(&Value::Array(arr.clone())))
            }
            Value::Object(_) => Dynamic::from(Self::json_to_rhai_map(value)),
            _ => Dynamic::UNIT,
        }
    }

    fn json_to_rhai_map(json: &Value) -> rhai::Map {
        let mut rhai_map = rhai::Map::new();
        if let Value::Object(obj) = json {
            for (key, value) in obj {
                let dynamic_value = Self::json_to_rhai_dynamic(value);
                rhai_map.insert(key.clone().into(), dynamic_value);
            }
        }
        rhai_map
    }

    fn rhai_map_to_json(map: Map) -> Value {
        let json_map: JsonMap<String, Value> = map
            .into_iter()
            .map(|(key, value)| (key.into(), Self::rhai_dynamic_to_json(value)))
            .collect();
        Value::Object(json_map)
    }

    fn rhai_dynamic_to_json(value: Dynamic) -> Value {
        match value.type_name() {
            "bool" => Value::Bool(value.as_bool().unwrap_or(false)),
            "int" | "i64" | "f64" => {
                Value::Number(serde_json::Number::from(value.as_int().unwrap_or(0)))
            }
            "float" => {
                let float_value = value.as_float().unwrap_or(0.0);
                serde_json::Number::from_f64(float_value)
                    .map(Value::Number)
                    .unwrap_or(Value::Null)
            }
            "string" => Value::String(value.cast::<String>()),
            "map" => Self::rhai_map_to_json(value.cast::<Map>()),
            "array" => {
                let array = value.cast::<Vec<Dynamic>>();
                Value::Array(array.into_iter().map(Self::rhai_dynamic_to_json).collect())
            }
            _ => serde_json::json!({"unsupported_type": value.type_name()}),
        }
    }
}

pub struct RhaiEngineConfig {
    pub max_operations: u64,      // Enforced through Rhai's set_max_operations
    pub max_recursion_depth: u32, // Enforced through Rhai's set_max_expr_depths
    pub max_script_size: u64,     // Enforced through input validation
    pub enable_sandbox: bool,     // Enforced through Rhai's sandboxing
    pub use_holochain_time: bool, // Enforced through time provider
}

impl Default for RhaiEngineConfig {
    fn default() -> Self {
        Self {
            max_operations: 100000,
            max_recursion_depth: 32,
            max_script_size: 1024 * 1024, // 1MB
            enable_sandbox: true,
            use_holochain_time: true,
        }
    }
}

impl RhaiEngineConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_max_operations(mut self, operations: usize) -> Self {
        self.max_operations = operations as u64;
        self
    }

    pub fn with_max_call_depth(mut self, depth: usize) -> Self {
        self.max_recursion_depth = depth as u32;
        self
    }

    pub fn with_max_script_size(mut self, size_bytes: usize) -> Self {
        self.max_script_size = size_bytes as u64;
        self
    }

    pub fn with_enable_sandbox(mut self, enable: bool) -> Self {
        self.enable_sandbox = enable;
        self
    }

    pub fn with_use_holochain_time(mut self, use_holochain_time: bool) -> Self {
        self.use_holochain_time = use_holochain_time;
        self
    }
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, SerializedBytes)]
pub struct RhaiEngineOutput {
    pub output: RAVEOutput,
    // these links will only be removed from the initial inputs
    pub rejected_links: Vec<ProcessedLinks>,
    // these links will be deleted and removed from the initial inputs
    pub redacted_links: Vec<ProcessedLinks>,
}

impl TryFrom<Value> for RhaiEngineOutput {
    type Error = String;

    fn try_from(value: Value) -> Result<Self, Self::Error> {
        match value {
            Value::Object(map) => {
                // output is required
                let output = if let Some(output) = map.get("output") {
                    RAVEOutput::try_from(output.clone()).unwrap_or_default()
                } else {
                    RAVEOutput::default()
                };

                let rejected_links = if let Some(links) = map.get("rejected_links") {
                    links
                        .as_array()
                        .ok_or("rejected_links must be an array")?
                        .iter()
                        .map(ProcessedLinks::try_from)
                        .collect::<Result<Vec<ProcessedLinks>, String>>()
                        .map_err(|e| format!("Failed to parse rejected_links: {e}"))?
                } else {
                    vec![]
                };

                let redacted_links = if let Some(links) = map.get("redacted_links") {
                    links
                        .as_array()
                        .ok_or("redacted_links must be an array")?
                        .iter()
                        .map(ProcessedLinks::try_from)
                        .collect::<Result<Vec<ProcessedLinks>, String>>()
                        .map_err(|e| format!("Failed to parse redacted_links: {e}"))?
                } else {
                    vec![]
                };

                Ok(RhaiEngineOutput {
                    output,
                    rejected_links,
                    redacted_links,
                })
            }
            _ => Err("Expected JSON object for RhaiEngineOutput".to_string()),
        }
    }
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, SerializedBytes)]
pub struct ProcessedLinks {
    pub hash: ActionHashB64,
    reason: String,
}
impl TryFrom<&Value> for ProcessedLinks {
    type Error = String;

    fn try_from(value: &Value) -> Result<Self, Self::Error> {
        let hash_str = value
            .get("hash")
            .ok_or("ProcessedLinks missing required 'hash' field")?
            .as_str()
            .ok_or("ProcessedLinks 'hash' field must be a string")?;

        let action_hash = ActionHash::try_from(hash_str.to_string())
            .map_err(|e| format!("Failed to parse hash '{hash_str}' as ActionHash: {e}"))?;

        let reason = value
            .get("reason")
            .ok_or("ProcessedLinks missing required 'reason' field")?
            .as_str()
            .ok_or("ProcessedLinks 'reason' field must be a string")?
            .to_string();

        Ok(ProcessedLinks {
            hash: ActionHashB64::from(action_hash),
            reason,
        })
    }
}