rustpad/src/app.rs

use axum::{
    extract::{DefaultBodyLimit, Path, State},
    http::{header, HeaderValue, StatusCode},
    response::{Html, IntoResponse, Response},
    routing::{get, post},
    Router,
};
use tower_http::{services::ServeDir, trace::TraceLayer};

use crate::{api, db, state::SharedState, websocket};

pub fn router(state: SharedState, static_dir: &str, upload_max_size_bytes: usize) -> Router {
    Router::new()
        .route("/", get(home))
        .route("/p/{slug}", get(pad))
        .route("/s/{token}", get(public_page))
        .route("/w/{workspace_slug}", get(workspace))
        .route("/w/{workspace_slug}/n/{note_slug}", get(note))
        .route("/health", get(health))
        .route("/f/{token}/{filename}", get(api::download_file))
        .route("/files/{directory}/{filename}", get(api::download_legacy_file))
        .route("/api/public/{token}", get(api::public_page))
        .route("/api/pads", post(api::create_pad))
        .route("/api/pads/{slug}", get(api::pad_info))
        .route("/api/pads/{slug}/history", post(api::pad_history))
        .route("/api/pads/{slug}/publish", post(api::publish_pad_page))
        .route("/api/pads/{slug}/restore", post(api::pad_restore))
        .route("/api/pads/{slug}/files", post(api::upload_pad_file))
        .route("/api/workspaces", post(api::create_workspace))
        .route("/api/workspaces/{workspace_slug}", get(api::workspace_info))
        .route("/api/workspaces/{workspace_slug}/open", post(api::open_workspace))
        .route("/api/workspaces/{workspace_slug}/notes", post(api::create_note))
        .route(
            "/api/workspaces/{workspace_slug}/notes/{note_slug}",
            get(api::note_info),
        )
        .route(
            "/api/workspaces/{workspace_slug}/notes/{note_slug}/publish",
            post(api::publish_note_page),
        )
        .route(
            "/api/workspaces/{workspace_slug}/notes/{note_slug}/history",
            post(api::history),
        )
        .route(
            "/api/workspaces/{workspace_slug}/notes/{note_slug}/restore",
            post(api::restore),
        )
        .route(
            "/api/workspaces/{workspace_slug}/notes/{note_slug}/files",
            post(api::upload_note_file),
        )
        .route("/ws/p/{slug}", get(websocket::upgrade_pad))
        .route(
            "/ws/{workspace_slug}/{note_slug}",
            get(websocket::upgrade),
        )
        .route("/static", get(static_not_found))
        .route("/static/{*path}", get(static_not_found))
        .nest_service("/assets", ServeDir::new(static_dir))
        .fallback(not_found)
        .layer(DefaultBodyLimit::max(upload_max_size_bytes.saturating_add(1024 * 1024)))
        .layer(TraceLayer::new_for_http())
        .with_state(state)
}

async fn health() -> &'static str {
    "ok"
}

async fn home(State(state): State<SharedState>) -> Response {
    versioned_html(include_str!("../static/home.html"), &state.asset_version)
}

async fn pad(
    State(state): State<SharedState>,
    Path(slug): Path<String>,
) -> Response {
    match db::find_pad(&state.db, &slug).await {
        Ok(Some(pad)) => {
            let html = include_str!("../static/pad.html")
                .replace("__PAD_TITLE__", &escape_html(&pad.title));
            versioned_html(&html, &state.asset_version)
        },
        Ok(None) => error_response(
            StatusCode::NOT_FOUND,
            "404",
            "Note not found",
            "This note does not exist or has been deleted.",
            "/",
            "Home page",
            &state.asset_version,
        ),
        Err(error) => {
            tracing::error!(%error, %slug, "failed to load standalone pad");
            internal_error(&state.asset_version)
        }
    }
}

async fn public_page(
    State(state): State<SharedState>,
    Path(token): Path<String>,
) -> Response {
    match db::find_published_page(&state.db, &token).await {
        Ok(Some(_)) => versioned_html(include_str!("../static/public.html"), &state.asset_version),
        Ok(None) => error_response(
            StatusCode::NOT_FOUND,
            "404",
            "Published page not found",
            "The link is invalid or the published page has been removed.",
            "/",
            "Home page",
            &state.asset_version,
        ),
        Err(error) => {
            tracing::error!(%error, %token, "failed to load published page");
            internal_error(&state.asset_version)
        }
    }
}

async fn workspace(
    State(state): State<SharedState>,
    Path(workspace_slug): Path<String>,
) -> Response {
    match db::find_workspace(&state.db, &workspace_slug).await {
        Ok(Some(workspace)) => {
            let html = include_str!("../static/workspace.html")
                .replace("__WORKSPACE_TITLE__", &escape_html(&workspace.title));
            versioned_html(&html, &state.asset_version)
        },
        Ok(None) => error_response(
            StatusCode::NOT_FOUND,
            "404",
            "Workspace not found",
            "This workspace does not exist or has been deleted.",
            "/",
            "Home page",
            &state.asset_version,
        ),
        Err(error) => {
            tracing::error!(%error, %workspace_slug, "failed to load workspace page");
            internal_error(&state.asset_version)
        }
    }
}

async fn note(
    State(state): State<SharedState>,
    Path((workspace_slug, note_slug)): Path<(String, String)>,
) -> Response {
    let workspace = match db::find_workspace(&state.db, &workspace_slug).await {
        Ok(Some(workspace)) => workspace,
        Ok(None) => {
            return error_response(
                StatusCode::NOT_FOUND,
                "404",
                "Workspace not found",
                "The workspace for this note does not exist or has been deleted.",
                "/",
                "Home page",
                &state.asset_version,
            );
        }
        Err(error) => {
            tracing::error!(%error, %workspace_slug, "failed to load note workspace");
            return internal_error(&state.asset_version);
        }
    };

    match db::find_note(&state.db, workspace.id, &note_slug).await {
        Ok(Some(note)) => {
            let html = include_str!("../static/note.html")
                .replace("__NOTE_TITLE__", &escape_html(&note.title))
                .replace("__WORKSPACE_TITLE__", &escape_html(&workspace.title))
                .replace("__WORKSPACE_SLUG__", &escape_html(&workspace_slug));
            versioned_html(&html, &state.asset_version)
        },
        Ok(None) => error_response(
            StatusCode::NOT_FOUND,
            "404",
            "Note not found",
            "This note does not exist or has been deleted.",
            &format!("/w/{workspace_slug}"),
            "Back do workspace",
            &state.asset_version,
        ),
        Err(error) => {
            tracing::error!(%error, %workspace_slug, %note_slug, "failed to load note page");
            internal_error(&state.asset_version)
        }
    }
}

async fn static_not_found() -> Response {
    let mut response = (StatusCode::NOT_FOUND, "404").into_response();
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("text/plain; charset=utf-8"),
    );
    response
}

async fn not_found(State(state): State<SharedState>) -> Response {
    error_response(
        StatusCode::NOT_FOUND,
        "404",
        "Page not found",
        "Check the address or return to the home page.",
        "/",
        "Home page",
        &state.asset_version,
    )
}

fn internal_error(asset_version: &str) -> Response {
    error_response(
        StatusCode::INTERNAL_SERVER_ERROR,
        "500",
        "Server error",
        "The page could not be loaded. Please try again shortly.",
        "/",
        "Home page",
        asset_version,
    )
}

fn error_response(
    status: StatusCode,
    code: &str,
    title: &str,
    message: &str,
    primary_url: &str,
    primary_label: &str,
    asset_version: &str,
) -> Response {
    let html = include_str!("../static/error.html")
        .replace("__ASSET_VERSION__", &escape_html(asset_version))
        .replace("__ERROR_CODE__", &escape_html(code))
        .replace("__ERROR_TITLE__", &escape_html(title))
        .replace("__ERROR_MESSAGE__", &escape_html(message))
        .replace("__PRIMARY_URL__", &escape_html(primary_url))
        .replace("__PRIMARY_LABEL__", &escape_html(primary_label));

    let mut response = (status, Html(html)).into_response();
    no_store(&mut response);
    response
}

fn versioned_html(template: &str, asset_version: &str) -> Response {
    let html = template.replace("__ASSET_VERSION__", asset_version);
    let mut response = Html(html).into_response();
    no_store(&mut response);
    response
}

fn no_store(response: &mut Response) {
    response.headers_mut().insert(
        header::CACHE_CONTROL,
        HeaderValue::from_static("private, no-store"),
    );
}

fn escape_html(value: &str) -> String {
    value
        .replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&#39;")
}