#[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(())
}