miasma 0.1.22

Trap malicious AI web scrapers in an endless pit of poisoned training data.
Documentation
mod poison;

pub use poison::LinkSettings;
pub use poison::serve_poison;

use crate::{MiasmaConfig, routes};

use std::sync::Arc;

use axum::{
    Router,
    body::Body,
    http::Request,
    response::{IntoResponse, Response},
    routing::get,
};
use reqwest::{StatusCode, header};
use serde::Deserialize;
use tokio::sync::{Semaphore, TryAcquireError};

#[derive(Deserialize)]
pub struct QueryParams {
    /// We use 'page' instead of depth to look more convincing to scrapers
    page: Option<u32>,
}
impl QueryParams {
    pub const CURRENT_DEPTH_QUERY_PARAM: &str = "page";
}

/// Build a new `axum::Router` for Miasma's routes.
pub fn new_miasma_router(config: &'static MiasmaConfig) -> Router {
    let in_flight_sem = Arc::new(Semaphore::new(config.max_in_flight as usize));

    Router::new().fallback(get(move |req: Request<Body>| async move {
        let in_flight_permit = match in_flight_sem.try_acquire_owned() {
            Ok(p) => p,
            Err(e) => match e {
                TryAcquireError::NoPermits => {
                    return Response::builder()
                        .status(StatusCode::TOO_MANY_REQUESTS)
                        .header(header::RETRY_AFTER, 5)
                        .body(Body::empty())
                        .unwrap();
                }
                TryAcquireError::Closed => {
                    return StatusCode::INTERNAL_SERVER_ERROR.into_response();
                }
            },
        };

        let gzip_response = config.force_gzip
            || req
                .headers()
                .get(header::ACCEPT_ENCODING)
                .map(|acc| {
                    acc.to_str()
                        .unwrap_or("")
                        .split(',')
                        // Don't you dare allocate anything !
                        .any(|tok| tok.trim().eq_ignore_ascii_case("gzip"))
                })
                .unwrap_or(false);

        let current_depth = axum::extract::Query::<QueryParams>::try_from_uri(req.uri())
            .ok()
            .and_then(|q| q.page)
            .unwrap_or(1);

        let link_settings = LinkSettings::next(config, current_depth);

        routes::serve_poison(config, in_flight_permit, gzip_response, link_settings)
            .await
            .into_response()
    }))
}