wasm-sandbox 0.4.1

// CLI Tool Wrapper Template
// This template is used to wrap a CLI tool in a WebAssembly module

use std::sync::{Arc, Mutex};
use std::process::{Command, Stdio, Child};
use std::io::{Read, Write};
use std::thread;

// Configuration
const TOOL_EXECUTABLE: &str = "{{tool_executable}}";
const INTERACTIVE: bool = {{interactive}};

// Global state
static mut TOOL_INSTANCE: Option<Arc<Mutex<ToolInstance>>> = None;

struct ToolInstance {
    running: bool,
    process: Option<Child>,
    output: Vec<u8>,
}

// Initialize the module
#[no_mangle]
pub extern "C" fn _initialize() {
    unsafe {
        TOOL_INSTANCE = Some(Arc::new(Mutex::new(ToolInstance {
            running: false,
            process: None,
            output: Vec::new(),
        })));
    }
}

// Run the tool with arguments
#[no_mangle]
pub extern "C" fn run(args_ptr: *const u8, args_len: usize) -> bool {
    // Convert args_ptr to a string
    let args_bytes = unsafe {
        if args_ptr.is_null() {
            return false;
        }
        std::slice::from_raw_parts(args_ptr, args_len)
    };
    
    let args_str = match std::str::from_utf8(args_bytes) {
        Ok(s) => s,
        Err(_) => return false,
    };
    
    // Parse arguments from JSON
    let args: Vec<String> = match serde_json::from_str(args_str) {
        Ok(a) => a,
        Err(_) => return false,
    };
    
    // Get tool instance
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let mut instance = instance.lock().unwrap();
    
    if instance.running {
        // Tool is already running
        return false;
    }
    
    // Build command to start the tool
    let mut command = Command::new(TOOL_EXECUTABLE);
    for arg in args {
        command.arg(arg);
    }
    
    // Configure stdio
    command.stdout(Stdio::piped());
    command.stderr(Stdio::piped());
    
    if INTERACTIVE {
        command.stdin(Stdio::piped());
    }
    
    // Start the tool
    match command.spawn() {
        Ok(mut child) => {
            instance.running = true;
            
            if !INTERACTIVE {
                // For non-interactive tools, capture output
                let stdout = child.stdout.take().expect("Failed to capture stdout");
                
                let instance_clone = unsafe { SERVER_INSTANCE.as_ref().unwrap().clone() };
                
                thread::spawn(move || {
                    let mut buffer = Vec::new();
                    let mut stdout_reader = std::io::BufReader::new(stdout);
                    
                    if let Err(_) = stdout_reader.read_to_end(&mut buffer) {
                        return;
                    }
                    
                    let mut instance = instance_clone.lock().unwrap();
                    instance.output = buffer;
                });
            }
            
            instance.process = Some(child);
            true
        },
        Err(_) => false,
    }
}

// Send input to the tool (for interactive tools)
#[no_mangle]
pub extern "C" fn send_input(input_ptr: *const u8, input_len: usize) -> bool {
    if !INTERACTIVE {
        return false;
    }
    
    // Convert input_ptr to a string
    let input_bytes = unsafe {
        if input_ptr.is_null() {
            return false;
        }
        std::slice::from_raw_parts(input_ptr, input_len)
    };
    
    // Get tool instance
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let mut instance = instance.lock().unwrap();
    
    if !instance.running || instance.process.is_none() {
        return false;
    }
    
    // Get stdin of the process
    if let Some(child) = &mut instance.process {
        if let Some(stdin) = child.stdin.as_mut() {
            if let Err(_) = stdin.write_all(input_bytes) {
                return false;
            }
            if let Err(_) = stdin.flush() {
                return false;
            }
            return true;
        }
    }
    
    false
}

// Get output from the tool
#[no_mangle]
pub extern "C" fn get_output() -> *const u8 {
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let instance = instance.lock().unwrap();
    
    if instance.output.is_empty() {
        return std::ptr::null();
    }
    
    let ptr = instance.output.as_ptr();
    std::mem::forget(instance.output.clone());
    ptr
}

// Get output length
#[no_mangle]
pub extern "C" fn get_output_len() -> usize {
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let instance = instance.lock().unwrap();
    instance.output.len()
}

// Stop the tool
#[no_mangle]
pub extern "C" fn stop() {
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let mut instance = instance.lock().unwrap();
    
    if !instance.running {
        return;
    }
    
    // Stop the tool
    if let Some(mut child) = instance.process.take() {
        let _ = child.kill();
        let _ = child.wait();
    }
    
    instance.running = false;
}

// Check if the tool is running
#[no_mangle]
pub extern "C" fn is_running() -> bool {
    let instance = unsafe { TOOL_INSTANCE.as_ref().unwrap().clone() };
    let instance = instance.lock().unwrap();
    instance.running
}

// Main function
fn main() {
    // Initialize
    _initialize();
    
    // The actual logic will be handled by the host through function calls
    println!("CLI tool wrapper initialized");
}