treehouse/crates/treehouse/src/cli/serve.rs

#[cfg(debug_assertions)]
mod live_reload;

use std::fmt::Write;
use std::{net::Ipv4Addr, sync::Arc};

use axum::{
    extract::{Path, Query, RawQuery, State},
    http::{
        header::{CACHE_CONTROL, CONTENT_TYPE},
        HeaderValue, StatusCode,
    },
    response::{Html, IntoResponse, Response},
    routing::get,
    Router,
};
use serde::Deserialize;
use tokio::net::TcpListener;
use tracing::{info, instrument};

use crate::sources::Sources;
use crate::vfs::asynch::AsyncDir;
use crate::vfs::VPath;
use crate::{html::EscapeHtml, state::Source};

mod system {
    use crate::vfs::VPath;

    pub const INDEX: &VPath = VPath::new_const("index");
    pub const FOUR_OH_FOUR: &VPath = VPath::new_const("_treehouse/404");
    pub const B_DOCS: &VPath = VPath::new_const("_treehouse/b");
}

struct Server {
    sources: Arc<Sources>,
    target: AsyncDir,
}

#[instrument(skip(sources, target))]
pub async fn serve(sources: Arc<Sources>, target: AsyncDir, port: u16) -> anyhow::Result<()> {
    let app = Router::new()
        .route("/", get(index)) // needed explicitly because * does not match empty paths
        .route("/*path", get(vfs_entry))
        .route("/b", get(branch))
        .fallback(get(four_oh_four))
        .with_state(Arc::new(Server { sources, target }));

    #[cfg(debug_assertions)]
    let app = app.nest("/dev/live-reload", live_reload::router());

    info!("serving on port {port}");
    let listener = TcpListener::bind((Ipv4Addr::from([0u8, 0, 0, 0]), port)).await?;
    Ok(axum::serve(listener, app).await?)
}

fn get_content_type(extension: &str) -> Option<&'static str> {
    match extension {
        "html" => Some("text/html"),
        "js" => Some("text/javascript"),
        "woff" => Some("font/woff2"),
        "svg" => Some("image/svg+xml"),
        _ => None,
    }
}

#[derive(Debug, Deserialize)]
struct VfsQuery {
    #[serde(rename = "v")]
    content_version: Option<String>,
}

#[instrument(skip(state))]
async fn get_static_file(path: &str, query: &VfsQuery, state: &Server) -> Option<Response> {
    let vpath = VPath::try_new(path).ok()?;
    let content = state.target.content(vpath).await?;
    let mut response = content.into_response();

    if let Some(content_type) = vpath.extension().and_then(get_content_type) {
        response
            .headers_mut()
            .insert(CONTENT_TYPE, HeaderValue::from_static(content_type));
    } else {
        response.headers_mut().remove(CONTENT_TYPE);
    }

    if query.content_version.is_some() {
        response.headers_mut().insert(
            CACHE_CONTROL,
            HeaderValue::from_static("public, max-age=31536000, immutable"),
        );
    }

    Some(response)
}

async fn vfs_entry(
    Path(path): Path<String>,
    Query(query): Query<VfsQuery>,
    State(state): State<Arc<Server>>,
) -> Response {
    if let Some(response) = get_static_file(&path, &query, &state).await {
        response
    } else {
        four_oh_four(State(state)).await
    }
}

async fn system_page(target: &AsyncDir, path: &VPath, status_code: StatusCode) -> Response {
    if let Some(content) = target.content(path).await {
        (status_code, Html(content)).into_response()
    } else {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("500 Internal Server Error: system page {path} is not available"),
        )
            .into_response()
    }
}

async fn index(State(state): State<Arc<Server>>) -> Response {
    system_page(&state.target, system::INDEX, StatusCode::OK).await
}

async fn four_oh_four(State(state): State<Arc<Server>>) -> Response {
    system_page(&state.target, system::FOUR_OH_FOUR, StatusCode::NOT_FOUND).await
}

#[instrument(skip(state))]
async fn branch(RawQuery(named_id): RawQuery, State(state): State<Arc<Server>>) -> Response {
    if let Some(named_id) = named_id {
        let branch_id = state
            .sources
            .treehouse
            .branches_by_named_id
            .get(&named_id)
            .copied()
            .or_else(|| {
                state
                    .sources
                    .treehouse
                    .branch_redirects
                    .get(&named_id)
                    .copied()
            });
        if let Some(branch_id) = branch_id {
            let branch = state.sources.treehouse.tree.branch(branch_id);
            if let Source::Tree { input, tree_path } =
                state.sources.treehouse.source(branch.file_id)
            {
                if let Some(content) = state
                    .target
                    .content(tree_path)
                    .await
                    .and_then(|s| String::from_utf8(s).ok())
                {
                    let branch_markup = input[branch.content.clone()].trim();
                    let mut per_page_metadata =
                        String::from("<meta property=\"og:description\" content=\"");
                    write!(per_page_metadata, "{}", EscapeHtml(branch_markup)).unwrap();
                    per_page_metadata.push_str("\">");
                    per_page_metadata.push_str(r#"<meta name="robots" content="noindex">"#);

                    const PER_PAGE_METADATA_REPLACEMENT_STRING: &str = "<!-- treehouse-ca37057a-cff5-45b3-8415-3b02dbf6c799-per-branch-metadata -->";
                    return Html(content.replacen(
                        PER_PAGE_METADATA_REPLACEMENT_STRING,
                        &per_page_metadata,
                        // Replace one under the assumption that it appears in all pages.
                        1,
                    ))
                    .into_response();
                } else {
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("500 Internal Server Error: branch metadata points to entry {tree_path} which does not have readable content")
                    )
                        .into_response();
                }
            }
        }

        system_page(&state.target, system::FOUR_OH_FOUR, StatusCode::NOT_FOUND).await
    } else {
        system_page(&state.target, system::B_DOCS, StatusCode::OK).await
    }
}