route.rs

//! ルーティングとミドルウェアの組み立て

use axum::Router;
use axum::extract::{DefaultBodyLimit, Request, State};
use axum::http::{HeaderName, HeaderValue, Method, StatusCode, header};
use axum::middleware::{Next, from_fn_with_state};
use axum::response::Response;
use crate::conf::Conf;
use crate::error;
use crate::ip;
use crate::session;
use crate::extract::Json;
use crate::state::AppState;
use std::any::Any;
use tower::ServiceBuilder;
use tower::util::option_layer;
use tower_http::catch_panic::CatchPanicLayer;
use tower_http::compression::CompressionLayer;
use tower_http::cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer, ExposeHeaders};
use tower_http::limit::RequestBodyLimitLayer;
use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
use tower_http::services::ServeDir;
use tower_http::set_header::{SetResponseHeader, SetResponseHeaderLayer};
use tower_http::trace::TraceLayer;


/// credentials を使うときに許すメソッド。`*` を返せないので一覧にする
const CORS_DEFAULT_METHOD: &[&str] = &["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"];

/// credentials を使うときに許すリクエストヘッダ
const CORS_DEFAULT_HEADER: &[&str] = &["content-type", "authorization"];

/// 渡されたルータにミドルウェアを掛け、状態を入れて仕上げる。
/// ミドルウェアは上に書いたものほど外側で動く。
///
/// **パスは 1 つも置かない。** `/health` も `/api/session` も `/relay` も `/static` も
/// fallback も、置くかどうかとどこに置くかは使う側が決める。
/// mon が勝手に取ると、そのパスが使えなくなるためである。
/// 部品は [`health`]、[`assets`]、[`crate::session::router`]、[`crate::relay::router`]、
/// [`crate::page::not_found`] にある。
pub fn wrap(routes: Router<AppState>, state: AppState) -> Router {
    let conf = state.conf.clone();
    let request_id = request_id_header(&conf);
    let traced = request_id.clone();
    let traced_conf = conf.clone();

    let middleware = ServiceBuilder::new()
        .layer(option_layer(
            request_id
                .clone()
                .map(|name| SetRequestIdLayer::new(name, MakeRequestUuid)),
        ))
        .layer(option_layer(request_id.map(PropagateRequestIdLayer::new)))
        // ハンドラが自分で付けていればそちらを残す。ページごとに変えられるようにする
        .layer(SetResponseHeaderLayer::if_not_present(
            header::CONTENT_SECURITY_POLICY,
            header_value(&conf.content_security_policy),
        ))
        .layer(TraceLayer::new_for_http().make_span_with(move |request: &Request| {
            let request_id = traced
                .as_ref()
                .and_then(|name| request.headers().get(name))
                .and_then(|value| value.to_str().ok())
                .unwrap_or("-");
            let client_ip = ip::resolve(request.headers(), request.extensions(), &traced_conf)
                .map(|address| address.to_string())
                .unwrap_or_else(|| "-".to_owned());
            tracing::info_span!(
                "request",
                method = %request.method(),
                uri = %request.uri(),
                request_id,
                client_ip = %client_ip
            )
        }))
        .layer(CatchPanicLayer::custom(panicked))
        .layer(from_fn_with_state(state.clone(), timeout))
        .layer(cors(&conf))
        .layer(CompressionLayer::new())
        .layer(from_fn_with_state(state.clone(), session::handle))
        // DefaultBodyLimit は extractor に、RequestBodyLimitLayer は素のボディに効く。
        // ボディの型を包むので、これより内側に from_fn は置けない
        .layer(DefaultBodyLimit::max(conf.body_limit))
        .layer(RequestBodyLimitLayer::new(conf.body_limit));

    routes.layer(middleware).with_state(state)
}


/// リクエスト 1 本にかける上限。時間切れも他と同じ JSON で返す。
/// tower_http の TimeoutLayer は body の空な応答を返すので、ここは自前で組み立てる
async fn timeout(State(state): State<AppState>, request: Request, next: Next) -> Response {
    let Ok(response) = tokio::time::timeout(state.conf.timeout, next.run(request)).await else {
        return error::json(StatusCode::REQUEST_TIMEOUT, "request timeout");
    };

    response
}


/// panic も他と同じ JSON で返す。既定の応答は text/plain の `Service panicked` である。
/// 差し替えると既定の記録もなくなるので、ここで残す
fn panicked(cause: Box<dyn Any + Send + 'static>) -> Response {
    let detail = cause
        .downcast_ref::<String>()
        .map(String::as_str)
        .or_else(|| cause.downcast_ref::<&str>().copied())
        .unwrap_or("詳細を取れません");
    tracing::error!(detail, "handler panicked");

    error::json(StatusCode::INTERNAL_SERVER_ERROR, "internal server error")
}


/// 静的ファイルを配る service。どこに置くかは使う側が決める。
///
/// `STATIC_CACHE_CONTROL` の値を付ける。空なら付けない。
/// 何も付けないと、ブラウザは Last-Modified からの経過時間で勝手に期限を決め、
/// その間は聞きにこない。名前の変わらないファイルを置き換えても古いものが出る。
/// 既定の `no-cache` は保存しないという意味ではなく、使う前に必ず聞きにこさせる、である。
pub fn assets(conf: &Conf) -> SetResponseHeader<ServeDir, Option<HeaderValue>> {
    SetResponseHeader::overriding(
        ServeDir::new(&conf.static_directory),
        header::CACHE_CONTROL,
        header_value(&conf.static_cache_control),
    )
}


/// 空なら付けない。`None` を渡すと、その層はヘッダを足さない。
/// 読めない値も付けないが、そちらは [`crate::conf::load`] が先に弾く
fn header_value(source: &str) -> Option<HeaderValue> {
    match source.is_empty() {
        true => None,
        false => HeaderValue::from_str(source).ok(),
    }
}


/// CORS。credentials を使うときは `*` を返せないので、そのときだけ一覧を明示する
fn cors(conf: &Conf) -> CorsLayer {
    let credentials = conf.cors_credentials && !conf.cors_origin.is_empty();
    if conf.cors_credentials && !credentials {
        tracing::warn!(
            "CORS_CREDENTIALS には CORS_ORIGIN が要ります。credentials は無効にしました"
        );
    }

    CorsLayer::new()
        .allow_credentials(credentials)
        .allow_origin(cors_origin(conf))
        .allow_methods(cors_method(conf, credentials))
        .allow_headers(cors_header(conf, credentials))
        .expose_headers(cors_expose_header(conf, credentials))
        .max_age(conf.cors_max_age)
}


fn cors_origin(conf: &Conf) -> AllowOrigin {
    if conf.cors_origin.is_empty() {
        return AllowOrigin::any();
    }

    let origin: Vec<HeaderValue> = conf
        .cors_origin
        .iter()
        .filter_map(|value| value.parse().ok())
        .collect();
    AllowOrigin::list(origin)
}


fn cors_method(conf: &Conf, credentials: bool) -> AllowMethods {
    if !conf.cors_method.is_empty() {
        return AllowMethods::list(to_method(&conf.cors_method));
    }

    if credentials {
        return AllowMethods::list(to_method(CORS_DEFAULT_METHOD));
    }

    AllowMethods::any()
}


fn cors_header(conf: &Conf, credentials: bool) -> AllowHeaders {
    if !conf.cors_header.is_empty() {
        return AllowHeaders::list(to_header(&conf.cors_header));
    }

    if credentials {
        return AllowHeaders::list(to_header(CORS_DEFAULT_HEADER));
    }

    AllowHeaders::any()
}


fn cors_expose_header(conf: &Conf, credentials: bool) -> ExposeHeaders {
    if !conf.cors_expose_header.is_empty() {
        return ExposeHeaders::list(to_header(&conf.cors_expose_header));
    }

    // credentials を使うと `*` を返せないので、せめてリクエスト ID は読ませる
    if credentials {
        return ExposeHeaders::list(request_id_header(conf));
    }

    ExposeHeaders::any()
}


/// リクエスト ID を入れるヘッダ。空なら付けない
fn request_id_header(conf: &Conf) -> Option<HeaderName> {
    if conf.request_id_header.is_empty() {
        return None;
    }

    Some(
        conf.request_id_header
            .parse()
            .expect("REQUEST_ID_HEADER はヘッダ名として使えません"),
    )
}


/// 読めなかったものは落とす。設定の書き間違いで起動そのものを止めない
fn to_method<T: AsRef<str>>(source: &[T]) -> Vec<Method> {
    source
        .iter()
        .filter_map(|value| value.as_ref().parse().ok())
        .collect()
}


fn to_header<T: AsRef<str>>(source: &[T]) -> Vec<HeaderName> {
    source
        .iter()
        .filter_map(|value| value.as_ref().parse().ok())
        .collect()
}


/// ヘルスチェック。`{"status": "ok"}` を返すだけで、DB は見ない。
/// どこに置くか、そもそも置くかは使う側が決める
pub async fn health() -> Json<serde_json::Value> {
    Json(serde_json::json!({"status": "ok"}))
}


#[cfg(test)]
mod tests {
    use crate::conf::{self, Conf};
    use crate::route::{assets, health, wrap};
    use crate::state::AppState;
    use crate::{page, relay, session};
    use axum::Router;
    use axum::body::{Body, to_bytes};
    use axum::http::{Request, StatusCode};
    use axum::response::IntoResponse;
    use axum::routing::get;
    use sqlx::sqlite::SqlitePoolOptions;
    use std::sync::Arc;
    use std::time::Duration;
    use tower::ServiceExt;

    /// 使う側が組むルータ。mon の部品を自分で好きな場所に置く
    fn routes(conf: &Conf) -> Router<AppState> {
        Router::new()
            .route("/health", get(health))
            .route("/mine", get(|| async { "mine" }))
            .route("/panic", get(panics))
            .route("/slow", get(takes_too_long))
            .route("/renew", axum::routing::post(renew_session))
            .route("/own-policy", get(own_policy))
            .nest("/api/session", session::router())
            .nest("/relay", relay::router())
            .nest_service("/static", assets(conf))
            .fallback(page::not_found)
    }

    /// 中身は残して id だけ作り直す
    async fn renew_session(session: session::Session) -> StatusCode {
        session.renew().await;
        StatusCode::NO_CONTENT
    }

    /// 自分で CSP を付けるハンドラ
    async fn own_policy() -> axum::response::Response {
        ([("content-security-policy", "default-src 'none'")], "自前").into_response()
    }

    /// わざと panic する
    async fn panics() -> &'static str {
        panic!("わざと")
    }

    /// 上限を必ず越える
    async fn takes_too_long() -> &'static str {
        tokio::time::sleep(Duration::from_secs(30)).await;
        "ここへは来ない"
    }

    async fn app() -> Router {
        app_with(conf::load().unwrap()).await
    }

    /// メモリ上の SQLite を使うので、接続は 1 本に絞る（複数だと別々の DB になる）
    async fn app_with(conf: Conf) -> Router {
        let database = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .unwrap();
        let store = session::store(&conf, &database);
        session::prepare(&store).await.unwrap();

        let routes = routes(&conf);
        wrap(
            routes,
            AppState {
                database,
                client: relay::client(&conf).unwrap(),
                session: store,
                conf: Arc::new(conf),
            },
        )
    }

    /// オリジンを 1 つだけ許し、Cookie も通す設定
    fn strict_conf() -> Conf {
        Conf {
            cors_origin: vec!["https://rollpie.com".to_owned()],
            cors_credentials: true,
            cors_max_age: Duration::from_secs(3600),
            ..Conf::default()
        }
    }

    /// CSP を書いた設定
    fn policy_conf() -> Conf {
        Conf {
            content_security_policy: "default-src 'self'".to_owned(),
            ..Conf::default()
        }
    }

    fn get_uri(uri: &str) -> Request<Body> {
        Request::builder().uri(uri).body(Body::empty()).unwrap()
    }

    fn get_with_cookie(uri: &str, id: &str) -> Request<Body> {
        Request::builder()
            .uri(uri)
            .header("cookie", format!("sid={id}"))
            .body(Body::empty())
            .unwrap()
    }

    fn post_language(language: &str, id: Option<&str>) -> Request<Body> {
        let mut builder = Request::builder()
            .method("POST")
            .uri("/api/session")
            .header("content-type", "application/json");

        if let Some(id) = id {
            builder = builder.header("cookie", format!("sid={id}"));
        }

        builder
            .body(Body::from(format!(r#"{{"language":"{language}"}}"#)))
            .unwrap()
    }

    fn preflight(origin: &str) -> Request<Body> {
        Request::builder()
            .method("OPTIONS")
            .uri("/api/session")
            .header("origin", origin)
            .header("access-control-request-method", "POST")
            .header("access-control-request-headers", "content-type")
            .body(Body::empty())
            .unwrap()
    }

    async fn read(response: axum::response::Response) -> String {
        let content = to_bytes(response.into_body(), 65536).await.unwrap();
        String::from_utf8_lossy(&content).into_owned()
    }

    fn header(response: &axum::response::Response, name: &str) -> Option<String> {
        let value = response.headers().get(name)?;
        Some(value.to_str().unwrap().to_owned())
    }

    fn cookie(response: &axum::response::Response) -> Option<String> {
        let line = header(response, "set-cookie")?;
        let value = line.strip_prefix("sid=")?;
        value.split(';').next().map(|id| id.to_owned())
    }

    /// 置き場を差し替えて往復させる
    async fn round_trip(kind: &str) {
        let conf = Conf {
            session_store: kind.to_owned(),
            session_directory: format!("target/test-session-{kind}"),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let written = app.clone().oneshot(post_language("ja", None)).await.unwrap();
        let id = cookie(&written).unwrap_or_else(|| panic!("{kind}: Set-Cookie がありません"));

        let answer = app.oneshot(get_with_cookie("/api/session", &id)).await.unwrap();
        assert_eq!(answer.status(), StatusCode::OK);
        assert!(read(answer).await.contains("\"ja\""), "{kind}");
    }

    #[tokio::test]
    async fn the_given_routes_are_mounted() {
        let response = app().await.oneshot(get_uri("/mine")).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(read(response).await, "mine");
    }

    #[tokio::test]
    async fn health_answers_ok() {
        let response = app().await.oneshot(get_uri("/health")).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn an_unknown_page_answers_html() {
        let response = app().await.oneshot(get_uri("/nothing")).await.unwrap();
        assert_eq!(response.status(), StatusCode::NOT_FOUND);
        let kind = header(&response, "content-type").unwrap();
        assert!(kind.contains("text/html"), "{kind}");
    }

    #[tokio::test]
    async fn a_static_file_is_served() {
        let response = app().await.oneshot(get_uri("/static/icon.svg")).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert!(read(response).await.contains("<svg"));
    }

    #[tokio::test]
    async fn a_static_file_carries_cache_control() {
        let response = app().await.oneshot(get_uri("/static/icon.svg")).await.unwrap();
        assert_eq!(header(&response, "cache-control").unwrap(), "no-cache");
    }

    #[tokio::test]
    async fn an_empty_cache_control_adds_no_header() {
        let conf = Conf {
            static_cache_control: String::new(),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let response = app.oneshot(get_uri("/static/icon.svg")).await.unwrap();
        assert_eq!(header(&response, "cache-control"), None);
    }

    /// 既定は空なので、mon は何も決めない
    #[tokio::test]
    async fn no_policy_adds_no_header() {
        let response = app().await.oneshot(get_uri("/health")).await.unwrap();
        assert_eq!(header(&response, "content-security-policy"), None);
    }

    #[tokio::test]
    async fn the_policy_is_added_to_every_answer() {
        let app = app_with(policy_conf()).await;
        let response = app.oneshot(get_uri("/static/icon.svg")).await.unwrap();

        let policy = header(&response, "content-security-policy").unwrap();
        assert_eq!(policy, "default-src 'self'");
    }

    /// ハンドラが自分で付けたものは残す。ページごとに変えられる
    #[tokio::test]
    async fn a_handler_can_set_its_own_policy() {
        let app = app_with(policy_conf()).await;
        let response = app.oneshot(get_uri("/own-policy")).await.unwrap();

        let policy = header(&response, "content-security-policy").unwrap();
        assert_eq!(policy, "default-src 'none'");
    }

    #[tokio::test]
    async fn the_preflight_carries_every_cors_header() {
        let app = app_with(strict_conf()).await;
        let response = app.oneshot(preflight("https://rollpie.com")).await.unwrap();

        let origin = header(&response, "access-control-allow-origin").unwrap();
        assert_eq!(origin, "https://rollpie.com");
        assert_eq!(header(&response, "access-control-allow-credentials").unwrap(), "true");
        assert_eq!(header(&response, "access-control-max-age").unwrap(), "3600");

        let method = header(&response, "access-control-allow-methods").unwrap();
        assert!(method.contains("POST"), "{method}");
        let allowed = header(&response, "access-control-allow-headers").unwrap();
        assert!(allowed.contains("content-type"), "{allowed}");
    }

    #[tokio::test]
    async fn an_unlisted_origin_gets_no_cors_header() {
        let app = app_with(strict_conf()).await;
        let response = app.oneshot(preflight("https://evil.example")).await.unwrap();
        assert_eq!(header(&response, "access-control-allow-origin"), None);
    }

    #[tokio::test]
    async fn credentials_need_an_origin_list() {
        let conf = Conf {
            cors_origin: Vec::new(),
            ..strict_conf()
        };

        // オリジンを絞らずに credentials を立てても、`*` との組み合わせで落ちないこと
        let app = app_with(conf).await;
        let response = app.oneshot(preflight("https://any.example")).await.unwrap();
        assert_eq!(header(&response, "access-control-allow-origin").unwrap(), "*");
        assert_eq!(header(&response, "access-control-allow-credentials"), None);
    }

    #[tokio::test]
    async fn the_request_id_is_readable_from_javascript() {
        let request = Request::builder()
            .uri("/health")
            .header("origin", "https://rollpie.com")
            .body(Body::empty())
            .unwrap();
        let response = app_with(strict_conf()).await.oneshot(request).await.unwrap();

        let expose = header(&response, "access-control-expose-headers").unwrap();
        assert!(expose.contains("x-request-id"), "{expose}");
    }

    #[tokio::test]
    async fn the_request_id_header_can_be_renamed() {
        let conf = Conf {
            request_id_header: "x-trace".to_owned(),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let response = app.oneshot(get_uri("/health")).await.unwrap();
        assert!(header(&response, "x-trace").is_some());
        assert_eq!(header(&response, "x-request-id"), None);
    }

    #[tokio::test]
    async fn the_request_id_can_be_turned_off() {
        let conf = Conf {
            request_id_header: String::new(),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let response = app.oneshot(get_uri("/health")).await.unwrap();
        assert_eq!(header(&response, "x-request-id"), None);
        assert_eq!(header(&response, "x-trace"), None);
    }

    /// panic は 500 になる。中身も他のエラーと同じ JSON である
    #[tokio::test]
    async fn a_panic_answers_json() {
        let response = app().await.oneshot(get_uri("/panic")).await.unwrap();
        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);

        let kind = header(&response, "content-type").unwrap();
        assert!(kind.contains("application/json"), "{kind}");
        assert!(read(response).await.contains(r#""error""#));
    }


    /// 時間切れは 408 になる。body の空な応答ではなく JSON である
    #[tokio::test]
    async fn a_timeout_answers_json() {
        let conf = Conf {
            timeout: Duration::from_millis(100),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let response = app.oneshot(get_uri("/slow")).await.unwrap();
        assert_eq!(response.status(), StatusCode::REQUEST_TIMEOUT);

        let kind = header(&response, "content-type").unwrap();
        assert!(kind.contains("application/json"), "{kind}");
        assert!(read(response).await.contains(r#""error""#));
    }


    #[tokio::test]
    async fn only_reading_issues_no_cookie() {
        let response = app().await.oneshot(get_uri("/health")).await.unwrap();
        assert_eq!(header(&response, "set-cookie"), None);
    }

    #[tokio::test]
    async fn writing_issues_a_cookie() {
        let response = app().await.oneshot(post_language("ja", None)).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        let line = header(&response, "set-cookie").unwrap();
        assert_eq!(cookie(&response).unwrap().len(), 64);
        assert!(line.contains("HttpOnly"), "{line}");
        assert!(line.contains("SameSite=Lax"), "{line}");
        assert!(!line.contains("Secure"), "{line}");
    }

    /// 保存し直さないなら、セッションがあっても Set-Cookie は出さない。
    /// 毎回付けると、静的ファイルまで共有キャッシュに載らなくなる
    #[tokio::test]
    async fn reading_an_existing_session_issues_no_cookie() {
        let app = app().await;
        let written = app.clone().oneshot(post_language("ja", None)).await.unwrap();
        let id = cookie(&written).unwrap();

        let answer = app.oneshot(get_with_cookie("/api/session", &id)).await.unwrap();
        assert_eq!(answer.status(), StatusCode::OK);
        assert_eq!(header(&answer, "set-cookie"), None);
    }


    #[tokio::test]
    async fn a_memory_store_keeps_it() {
        round_trip("memory").await;
    }

    #[tokio::test]
    async fn a_file_store_keeps_it() {
        round_trip("file").await;
    }

    #[tokio::test]
    async fn a_sqlite_store_keeps_it() {
        round_trip("sqlite").await;
    }

    #[tokio::test]
    async fn the_none_store_issues_nothing() {
        let conf = Conf {
            session_store: "none".to_owned(),
            ..Conf::default()
        };

        let app = app_with(conf).await;
        let response = app.oneshot(post_language("ja", None)).await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(header(&response, "set-cookie"), None);
    }

    /// 相手が決めた id をそのまま採ると、あとから乗っ取られる
    #[tokio::test]
    async fn a_made_up_id_is_not_adopted() {
        let made_up = "a".repeat(64);
        let response = app()
            .await
            .oneshot(post_language("ja", Some(&made_up)))
            .await
            .unwrap();

        assert_ne!(cookie(&response).unwrap(), made_up);
    }

    /// 立場が変わったら id を替える。中身は残り、古い id は通らなくなる
    #[tokio::test]
    async fn renewing_changes_the_id_and_keeps_the_data() {
        let app = app().await;
        let written = app.clone().oneshot(post_language("ja", None)).await.unwrap();
        let first = cookie(&written).unwrap();

        let request = Request::builder()
            .method("POST")
            .uri("/renew")
            .header("cookie", format!("sid={first}"))
            .body(Body::empty())
            .unwrap();
        let response = app.clone().oneshot(request).await.unwrap();
        let second = cookie(&response).unwrap();
        assert_ne!(second, first);

        let answer = app
            .clone()
            .oneshot(get_with_cookie("/api/session", &second))
            .await
            .unwrap();
        assert!(read(answer).await.contains("\"ja\""));

        let stale = app.oneshot(get_with_cookie("/api/session", &first)).await.unwrap();
        assert!(read(stale).await.contains("null"));
    }


    #[tokio::test]
    async fn destroying_clears_the_cookie() {
        let app = app().await;
        let written = app.clone().oneshot(post_language("ja", None)).await.unwrap();
        let id = cookie(&written).unwrap();

        let request = Request::builder()
            .method("DELETE")
            .uri("/api/session")
            .header("cookie", format!("sid={id}"))
            .body(Body::empty())
            .unwrap();
        let response = app.clone().oneshot(request).await.unwrap();
        assert_eq!(response.status(), StatusCode::NO_CONTENT);

        let line = header(&response, "set-cookie").unwrap();
        assert!(line.contains("Max-Age=0"), "{line}");

        let after = app.oneshot(get_with_cookie("/api/session", &id)).await.unwrap();
        assert!(read(after).await.contains("null"));
    }
}