iocaine 1.0.0

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

use anyhow::Result;
use axum::{
    extract::{Path, State},
    http::StatusCode,
    middleware,
    response::{Html, IntoResponse, Response},
    routing::get,
    Router,
};
use std::sync::Arc;

use crate::{
    assembled_statistical_sequences::AssembledStatisticalSequences, config::Config,
    garglebargle::GargleBargle, tenx_programmer, wurstsalat_generator_pro::WurstsalatGeneratorPro,
};

pub type StatefulIocaine = Arc<Iocaine>;

#[derive(Debug)]
pub struct Iocaine {
    pub config: Config,
    pub chain: WurstsalatGeneratorPro,
    pub words: GargleBargle,
}

impl Iocaine {
    pub fn new(config: Config) -> Result<Self> {
        tracing::info!("Loading the markov chain corpus");
        let chain = WurstsalatGeneratorPro::learn_from_files(&config.sources.markov)?;
        tracing::info!("Corpus loaded");

        tracing::info!("Loading word salad");
        let words = GargleBargle::load_words(&config.sources.words)?;
        tracing::info!("Word salad loaded");

        Ok(Self {
            config,
            chain,
            words,
        })
    }

    fn main_app(self) -> Router {
        let state: StatefulIocaine = Arc::new(self);

        let mut app = Router::new()
            .route("/", get(poison))
            .route("/{*path}", get(poison))
            .layer(tower_http::trace::TraceLayer::new_for_http());

        if state.config.metrics.enable {
            app = app.route_layer(middleware::from_fn_with_state(
                state.clone(),
                tenx_programmer::track_metrics,
            ));
        }

        app.with_state(state)
    }

    async fn start_main_server(self) -> std::result::Result<(), std::io::Error> {
        let bind = &self.config.server.bind.clone();

        let app = self.main_app();
        let listener = tokio::net::TcpListener::bind(bind).await?;
        axum::serve(listener, app)
            .with_graceful_shutdown(shutdown_signal())
            .await
    }

    pub async fn run(self) -> Result<()> {
        if self.config.metrics.enable {
            let metrics_bind = &self.config.metrics.bind.clone();
            let metrics_server = tenx_programmer::start_metrics_server(metrics_bind.to_string());
            let (main_server, _metrics_server) =
                tokio::join!(self.start_main_server(), metrics_server);
            Ok(main_server?)
        } else {
            let main_server = self.start_main_server().await;
            Ok(main_server?)
        }
    }
}

pub async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}

async fn poison(
    headers: axum::http::HeaderMap,
    State(iocaine): State<StatefulIocaine>,
    path: Option<Path<String>>,
) -> std::result::Result<Html<String>, AppError> {
    let default_host = axum::http::HeaderValue::from_static("<unknown>");
    let host = headers.get("host").unwrap_or(&default_host).to_str()?;
    let path = path.unwrap_or(Path("/".to_string()));

    let garbage = AssembledStatisticalSequences::generate(&iocaine, host, &path)?;

    if iocaine.config.metrics.enable {
        metrics::counter!("iocaine_garbage_served").increment(garbage.len() as u64);
    }

    Ok(Html(garbage))
}

pub struct AppError(anyhow::Error);

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        tracing::error!("Internal server error: {}", self.0);
        (StatusCode::INTERNAL_SERVER_ERROR, "Something went wrong").into_response()
    }
}

impl From<axum::http::header::ToStrError> for AppError {
    fn from(e: axum::http::header::ToStrError) -> Self {
        Self(e.into())
    }
}

impl From<anyhow::Error> for AppError {
    fn from(e: anyhow::Error) -> Self {
        Self(e)
    }
}

impl From<std::io::Error> for AppError {
    fn from(e: std::io::Error) -> Self {
        Self(e.into())
    }
}

impl From<metrics_exporter_prometheus::BuildError> for AppError {
    fn from(e: metrics_exporter_prometheus::BuildError) -> Self {
        Self(e.into())
    }
}