Code Sandboxing

const DEFAULT_FUEL: u64 = 1_000_000;      // Instruction budget
const DEFAULT_TIMEOUT_SECS: u64 = 30;     // Wall-clock timeout
const DEFAULT_MEMORY_PAGES: u64 = 256;    // 16MB (64KB per page)

const WASM_ALLOWED_FUNCTIONS: &[&str] = &[
    "memory::recall",
    "memory::store",
    "tool::web_fetch",
    "tool::file_read",
    "tool::file_list",
    "security::scan_injection",
    "embedding::generate",
];

fn execute_with_dual_metering(
    store: &mut Store<()>,
    func: &TypedFunc<(i32, i32), i64>,
    args: (i32, i32),
    fuel_limit: u64,
    timeout: Duration,
) -> Result<i64, ExecutionError> {
    // Set fuel limit
    store.set_fuel(fuel_limit)?;
    
    // Enable epoch-based interruption
    store.epoch_deadline_trap();
    store.set_epoch_deadline(1);

    let engine = store.engine().clone();
    let killed = Arc::new(AtomicBool::new(false));
    let killed_clone = killed.clone();

    // Watchdog thread increments epoch every 1ms
    let watchdog = std::thread::spawn(move || {
        let start = Instant::now();
        loop {
            std::thread::sleep(Duration::from_millis(1));
            engine.increment_epoch();
            if start.elapsed() > timeout || killed_clone.load(Ordering::SeqCst) {
                break;
            }
        }
    });

    // Execute with both limits enforced
    let result = func.call(store, args);

    killed.store(true, Ordering::SeqCst);
    let _ = watchdog.join();

    match result {
        Ok(val) => Ok(val),
        Err(e) => {
            let msg = e.to_string();
            if msg.contains("epoch") || msg.contains("interrupt") {
                Err(ExecutionError::Timeout)
            } else if msg.contains("fuel") {
                Err(ExecutionError::FuelExhausted)
            } else {
                Err(ExecutionError::Trapped(msg))
            }
        }
    }
}

// Verify module doesn't exceed memory limit
let memory = instance.get_memory(&mut store, "memory")
    .ok_or_else(|| IIIError::Handler("Module does not export 'memory'".into()))?;

let current_pages = memory.size(&store);
if current_pages > memory_pages {
    return Err(IIIError::Handler(format!(
        "Module memory {} pages exceeds limit of {}",
        current_pages, memory_pages
    )));
}

pub fn register(iii: &III) {
    iii.register_function_with_description(
        "sandbox::docker_exec",
        "Execute code in isolated Docker container",
        move |input: Value| async move {
            let code = input["code"].as_str().unwrap_or("");
            let language = input["language"].as_str().unwrap_or("python");
            let timeout = input["timeout"].as_u64().unwrap_or(30);
            
            // Create temporary directory
            let temp_dir = tempfile::tempdir()?;
            let code_path = temp_dir.path().join("code");
            std::fs::write(&code_path, code)?;
            
            // Run in container
            let output = Command::new("docker")
                .args(&[
                    "run",
                    "--rm",
                    "--network=none",
                    "--memory=512m",
                    "--cpus=1",
                    "--pids-limit=100",
                    "-v", &format!("{}:/workspace:ro", temp_dir.path().display()),
                    &format!("agentos-{}", language),
                    "/workspace/code",
                ])
                .timeout(Duration::from_secs(timeout))
                .output()
                .await?;
            
            Ok(json!({
                "stdout": String::from_utf8_lossy(&output.stdout),
                "stderr": String::from_utf8_lossy(&output.stderr),
                "exitCode": output.status.code(),
            }))
        },
    );
}

static INJECTION_PATTERNS: &[&str] = &[
    r"(?i)ignore\s+(all\s+)?(previous|above|prior)\s+(instructions|prompts)",
    r"(?i)you\s+are\s+now\s+",
    r"(?i)system\s*:\s*",
    r"(?i)\bDAN\b.*\bmode\b",
    r"(?i)pretend\s+you\s+are",
    r"(?i)act\s+as\s+if\s+you",
    r"(?i)disregard\s+(your|all)",
    r"(?i)override\s+(your|system)",
    r"(?i)jailbreak",
];

fn scan_injection(text: &str) -> Result<Value, IIIError> {
    let mut matches = Vec::new();
    let compiled = compiled_injection_patterns();

    for re in compiled {
        if re.is_match(text) {
            matches.push(re.as_str().to_string());
        }
    }

    let risk_score = (matches.len() as f64 * 0.25).min(1.0);

    Ok(json!({
        "safe": matches.is_empty(),
        "matches": matches,
        "riskScore": risk_score,
    }))
}

interface RateLimitConfig {
  requestsPerMinute: number;
  burstSize?: number;
}

async function checkRateLimit(
  agentId: string,
  config: RateLimitConfig
): Promise<boolean> {
  const now = Date.now();
  const interval = 60_000 / config.requestsPerMinute;  // ms per request
  const burst = config.burstSize || config.requestsPerMinute;
  
  const state = await getState(agentId) || {
    tat: now,  // Theoretical arrival time
    allowance: burst,
  };
  
  const newTat = Math.max(state.tat, now) + interval;
  const allowance = burst - (newTat - now) / interval;
  
  if (allowance < 1) {
    return false;  // Rate limit exceeded
  }
  
  await setState(agentId, { tat: newTat, allowance });
  return true;
}

interface LoopState {
  iterations: number;
  lastReset: number;
  tripped: boolean;
}

const MAX_ITERATIONS = 1000;
const RESET_WINDOW_MS = 60_000;

async function checkLoopGuard(agentId: string, loopId: string): Promise<void> {
  const key = `${agentId}:${loopId}`;
  const state = await getLoopState(key) || {
    iterations: 0,
    lastReset: Date.now(),
    tripped: false,
  };
  
  // Reset after window
  if (Date.now() - state.lastReset > RESET_WINDOW_MS) {
    state.iterations = 0;
    state.lastReset = Date.now();
    state.tripped = false;
  }
  
  // Check if tripped
  if (state.tripped) {
    throw new Error(`Loop guard tripped for ${loopId}`);
  }
  
  // Increment and check limit
  state.iterations++;
  if (state.iterations > MAX_ITERATIONS) {
    state.tripped = true;
    await setLoopState(key, state);
    throw new Error(`Loop guard: exceeded ${MAX_ITERATIONS} iterations`);
  }
  
  await setLoopState(key, state);
}