website/src/handlers/mod.rs

use axum::{
    body::Body,
    extract::{Request, State},
    http::{header, HeaderMap, StatusCode},
    middleware::Next,
    response::{Html, IntoResponse, Response},
    routing::get,
    Router,
};
use cached::once_cell::sync::Lazy;
use chrono::{DateTime, FixedOffset};
use prometheus::{opts, Encoder, IntCounterVec, TextEncoder};
use std::sync::Arc;
use tower_http::services::ServeDir;
use tracing::{
    instrument,
    log::{debug, error, info},
};

use crate::{AppState, WebsiteError};

pub mod posts;
pub mod tags;

pub static HIT_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
    prometheus::register_int_counter_vec!(
        opts!(
            "http_requests_total",
            "Total amount of http requests received"
        ),
        &["route", "method", "status"]
    )
    .unwrap()
});

pub fn routes(state: &Arc<AppState>) -> Router<Arc<AppState>> {
    Router::new()
        .route("/", get(index))
        .merge(posts::router())
        .merge(tags::router())
        .merge(posts::alias_router(state.posts.values()))
        .route("/healthcheck", get(healthcheck))
        .route_service("/posts/:slug/*path", ServeDir::new("./"))
        .route_service("/static/*path", ServeDir::new("./"))
        .layer(axum::middleware::from_fn(metrics_middleware))
        .route("/metrics", get(metrics))
}

#[instrument(skip(state))]
pub async fn index(
    State(state): State<Arc<AppState>>,
    headers: HeaderMap,
) -> std::result::Result<Response, WebsiteError> {
    if should_return_304(&headers, Some(state.startup_time.into())) {
        return Ok(StatusCode::NOT_MODIFIED.into_response());
    }
    let mut ctx = tera::Context::new();
    ctx.insert("base_url", &state.base_url.to_string());
    let res = state.tera.render("index.html", &ctx).map_err(|e| {
        error!("Failed rendering index: {}", e);
        WebsiteError::InternalError(e.into())
    })?;
    Ok((
        StatusCode::OK,
        [(header::LAST_MODIFIED, state.startup_time.to_rfc2822())],
        Html(res),
    )
        .into_response())
}

async fn healthcheck() -> &'static str {
    "OK"
}

#[instrument]
async fn metrics() -> impl IntoResponse {
    let encoder = TextEncoder::new();
    let metric_families = prometheus::gather();
    let mut buffer = vec![];
    encoder.encode(&metric_families, &mut buffer).unwrap();

    Response::builder()
        .status(200)
        .header(header::CONTENT_TYPE, encoder.format_type())
        .body(Body::from(buffer))
        .unwrap()
}

pub async fn not_found() -> impl IntoResponse {
    (StatusCode::NOT_FOUND, ())
}

#[instrument(skip(request, next))]
pub async fn metrics_middleware(request: Request, next: Next) -> Response {
    let path = request.uri().path().to_string();
    let method = request.method().clone();

    let response = next.run(request).await;

    if !response.status().is_client_error() {
        HIT_COUNTER
            .with_label_values(&[&path, method.as_str(), response.status().as_str()])
            .inc();
    } else if response.status() == StatusCode::NOT_FOUND {
        HIT_COUNTER
            .with_label_values(&["not found", method.as_str(), response.status().as_str()])
            .inc();
    }

    response
}

fn should_return_304(headers: &HeaderMap, last_changed: Option<DateTime<FixedOffset>>) -> bool {
    let Some(date) = last_changed else {
        debug!("no last modified date");
        return false;
    };
    let Some(since) = headers.get(header::IF_MODIFIED_SINCE) else {
        debug!("no IF_MODIFIED_SINCE header");
        return false;
    };

    let Ok(parsed) = DateTime::<FixedOffset>::parse_from_rfc2822(since.to_str().unwrap()) else {
        debug!("failed to parse IF_MODIFIED_SINCE header");
        return false;
    };

    date > parsed
}

impl IntoResponse for WebsiteError {
    fn into_response(self) -> Response {
        match self {
            WebsiteError::NotFound => {
                info!("not found");
                (StatusCode::NOT_FOUND, ()).into_response()
            }
            WebsiteError::InternalError(e) => {
                if let Some(s) = e.source() {
                    error!("internal error: {}: {}", e, s);
                } else {
                    error!("internal error: {}", e);
                }
                (StatusCode::INTERNAL_SERVER_ERROR, ()).into_response()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use crate::AppState;

    #[test]
    fn render_index() {
        let tera = tera::Tera::new("templates/**/*").unwrap();
        let mut ctx = tera::Context::new();
        ctx.insert("base_url", "http://localhost/");
        let _res = tera.render("index.html", &ctx).unwrap();
    }

    #[tokio::test]
    async fn setup_routes() {
        let mut state = AppState {
            startup_time: chrono::offset::Utc::now(),
            base_url: "http://localhost:8180".parse().unwrap(),
            tera: tera::Tera::new("templates/**/*").unwrap(),
            ..Default::default()
        };
        // Load the actual posts, just to make this test fail if
        // aliases overlap with themselves or other routes
        let posts = crate::post::load_all(&state).await.unwrap();
        state.tags = crate::tag::get_tags(posts.values());
        state.posts = posts;
        let state = Arc::new(state);

        super::routes(&state).with_state(state).into_make_service();
    }
}