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