iocaine 2.5.0

The deadliest poison known to AI
Documentation
// SPDX-FileCopyrightText: 2025 Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

#[cfg(all(not(target_env = "musl"), feature = "jemalloc"))]
#[global_allocator]
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

use anyhow::Result;
use axum::http::header::{self, HeaderMap, HeaderName};
use clap::Parser;
use std::collections::BTreeMap;
use std::error::Error;

use iocaine::sex_dungeon::{self, Language, LanguageOptions, Request};

#[derive(Debug, Parser)]
#[command(version, about)]
pub struct Args {
    #[arg(short = 's', long)]
    pub script_path: String,
    #[arg(short = 'l', long, value_parser = Self::parse_language, default_value_t = Language::Roto)]
    pub script_language: Language,
    #[arg(short = 'o', long = "language-option", value_parser = Self::parse_key_val::<String, String>)]
    pub language_options: Vec<(String, String)>,

    #[arg(short = 'P', long)]
    pub path: Option<String>,
    #[arg(short = 'A', long, default_value_t = Self::default_user_agent())]
    pub user_agent: String,
    #[arg(short = 'H', long = "header", value_parser = Self::parse_key_val::<String, String>)]
    pub headers: Vec<(String, String)>,
    #[arg(short = 'Q', long = "query", value_parser = Self::parse_key_val::<String, String>)]
    pub query_params: Vec<(String, String)>,
    #[arg(short = 'M', long)]
    pub method: Option<String>,
}

impl Args {
    fn default_user_agent() -> String {
        "Mozilla/5.0 (X11; Linux x86_64; rv:137.0) Gecko/20100101 Firefox/137.0".to_owned()
    }

    fn parse_key_val<T, U>(
        s: &str,
    ) -> std::result::Result<(T, U), Box<dyn Error + Send + Sync + 'static>>
    where
        T: std::str::FromStr,
        T::Err: Error + Send + Sync + 'static,
        U: std::str::FromStr,
        U::Err: Error + Send + Sync + 'static,
    {
        let pos = s
            .find('=')
            .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
        Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
    }

    fn parse_language(
        s: &str,
    ) -> std::result::Result<Language, Box<dyn Error + Send + Sync + 'static>> {
        match s {
            "roto" => Ok(Language::Roto),
            "lua" => Ok(Language::Lua),
            "fennel" => Ok(Language::Fennel),
            l => Err(format!("Unsupported script language: {l}").into()),
        }
    }
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
        )
        .with_writer(std::io::stderr)
        .init();

    let args = Args::parse();

    tracing::info!("Preparing the runtime");

    let mut lang_opts = LanguageOptions::new();
    for (key, value) in args.language_options {
        lang_opts.insert(key, value);
    }

    let handler = sex_dungeon::create(args.script_language, &lang_opts, &args.script_path, None)?;

    let mut headers = HeaderMap::new();
    headers.insert(header::HOST, "example.com".parse()?);
    headers.insert(header::USER_AGENT, args.user_agent.parse()?);

    for (key, value) in args.headers {
        headers.insert(HeaderName::from_bytes(key.as_bytes())?, value.parse()?);
    }

    let mut params = BTreeMap::new();
    for (key, value) in args.query_params {
        params.insert(key, value);
    }

    let request = Request {
        method: args.method.unwrap_or_else(|| "GET".to_owned()),
        path: args.path.unwrap_or_default(),
        headers,
        params,
    };

    tracing::info!("Executing `decide()`");

    let verdict = handler.decide(request);
    println!("{verdict:?}");

    Ok(())
}