relay.rs

//! 中継のテスト。上流をテストの中で立てて、行きと帰りを確かめる。
//!
//! 中継だけは mon 単体で試せない。転送先が動いていないと、何も通らないためである。

use axum::Router;
use axum::body::{Body, Bytes, to_bytes};
use axum::extract::{Path, Request};
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::{any, get};
use futures_util::StreamExt;
use mon::state::AppState;
use mon::{Conf, database, relay, route, session};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpListener;
use tower::ServiceExt;


/// 上流を 1 つ起こして、その起点を返す
async fn upstream() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();

    tokio::spawn(async move {
        axum::serve(listener, upstream_router()).await.unwrap();
    });

    format!("http://127.0.0.1:{port}")
}


fn upstream_router() -> Router {
    Router::new()
        .route("/echo", any(echo))
        .route("/status/{code}", get(status))
        .route("/header", get(with_header))
        .route("/big", get(big))
        .route("/slow", get(slow))
        .route("/go", get(go))
        .route("/forwarded", get(forwarded))
        .route("/cookie", get(cookie))
        .route("/set-cookie", get(set_cookie))
}


/// 受け取った Cookie をそのまま書き出す
async fn cookie(request: Request) -> String {
    header_text(request.headers(), "cookie")
}


/// mon のセッションと同じ名前のクッキーと、別のクッキーを書く
async fn set_cookie() -> Response {
    let mut response = "ok".into_response();
    let headers = response.headers_mut();
    headers.append("set-cookie", HeaderValue::from_static("sid=planted; Path=/"));
    headers.append("set-cookie", HeaderValue::from_static("theme=dark; Path=/"));

    response
}


/// 受け取った forwarded 系をそのまま書き出す
async fn forwarded(request: Request) -> String {
    let headers = request.headers();

    format!(
        "{} {} {}",
        header_text(headers, "x-forwarded-for"),
        header_text(headers, "x-forwarded-proto"),
        header_text(headers, "x-forwarded-host")
    )
}


fn header_text(headers: &HeaderMap, name: &str) -> String {
    headers
        .get(name)
        .and_then(|value| value.to_str().ok())
        .unwrap_or("-")
        .to_owned()
}


/// 受け取ったものをそのまま書き出す
async fn echo(request: Request) -> String {
    let (parts, body) = request.into_parts();
    let sent = parts
        .headers
        .get("x-sent")
        .and_then(|value| value.to_str().ok())
        .unwrap_or("-")
        .to_owned();
    let content = to_bytes(body, 1024 * 1024).await.unwrap();

    format!(
        "{} {} {} {}",
        parts.method,
        parts.uri,
        sent,
        String::from_utf8_lossy(&content)
    )
}


async fn status(Path(code): Path<u16>) -> StatusCode {
    StatusCode::from_u16(code).unwrap()
}


async fn with_header() -> Response {
    ([("x-from-upstream", "yes")], "ok").into_response()
}


/// 302 を返す。中継はこれを追わず、そのまま返すのが正しい
async fn go() -> Response {
    (StatusCode::FOUND, [("location", "/echo")], "").into_response()
}


async fn big() -> Vec<u8> {
    vec![b'a'; 4 * 1024 * 1024]
}


/// 先頭を出したあと 1 秒おいて続きを出す。読み切ってから返す作りだと 1 秒待たされる
async fn slow() -> Response {
    let stream = futures_util::stream::unfold(0u8, |step| async move {
        match step {
            0 => Some((Ok::<Bytes, std::io::Error>(Bytes::from_static(b"first")), 1)),

            1 => {
                tokio::time::sleep(Duration::from_secs(1)).await;
                Some((Ok(Bytes::from_static(b"second")), 2))
            }

            _ => None,
        }
    });

    Body::from_stream(stream).into_response()
}


/// 中継だけを置いた設定。セッションは使わない
fn relay_conf(relay_origin: &str, relay_timeout: Duration) -> Conf {
    Conf {
        relay_origin: relay_origin.to_owned(),
        relay_timeout,
        session_store: "none".to_owned(),
        ..Conf::default()
    }
}


async fn mon_app(relay_origin: &str, relay_timeout: Duration) -> Router {
    mon_app_with(relay_conf(relay_origin, relay_timeout)).await
}


async fn mon_app_with(conf: Conf) -> Router {
    let database = database::connect("sqlite::memory:").await.unwrap();
    let store = session::store(&conf, &database);
    session::prepare(&store).await.unwrap();
    let routes = Router::new().nest("/relay", relay::router());

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


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


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


#[tokio::test]
async fn the_method_path_query_header_and_body_all_arrive() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let request = Request::builder()
        .method("POST")
        .uri("/relay/echo?a=1&b=2")
        .header("x-sent", "hello")
        .body(Body::from("honbun"))
        .unwrap();

    let response = app.oneshot(request).await.unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(read(response).await, "POST /echo?a=1&b=2 hello honbun");
}


#[tokio::test]
async fn the_status_comes_back_as_it_is() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;

    let response = app.clone().oneshot(get_uri("/relay/status/404")).await.unwrap();
    assert_eq!(response.status(), StatusCode::NOT_FOUND);

    let response = app.oneshot(get_uri("/relay/status/201")).await.unwrap();
    assert_eq!(response.status(), StatusCode::CREATED);
}


#[tokio::test]
async fn the_upstream_headers_come_back() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let response = app.oneshot(get_uri("/relay/header")).await.unwrap();

    let value = response.headers().get("x-from-upstream").unwrap();
    assert_eq!(value, "yes");
}


#[tokio::test]
async fn a_large_body_is_not_broken() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let response = app.oneshot(get_uri("/relay/big")).await.unwrap();
    let content = to_bytes(response.into_body(), 8 * 1024 * 1024).await.unwrap();

    assert_eq!(content.len(), 4 * 1024 * 1024);
    assert!(content.iter().all(|byte| *byte == b'a'));
}


/// mon が最前段のときは、forwarded 系を自分で足す。
/// これがないと、中継先は公開しているホストも scheme も知れない
#[tokio::test]
async fn the_forwarded_headers_are_added() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let request = Request::builder()
        .uri("/relay/forwarded")
        .header("host", "rollpie.com")
        .body(Body::empty())
        .unwrap();

    let response = app.oneshot(request).await.unwrap();
    assert_eq!(read(response).await, "- http rollpie.com");
}


/// 前段がいる（CLIENT_IP_HEADER を書いた）なら、すでにあるものには触らない。
/// 前段が付けたものを書き換えると、数える位置がずれる
#[tokio::test]
async fn the_existing_forwarded_headers_are_kept() {
    let conf = Conf {
        client_ip_header: "x-forwarded-for".to_owned(),
        ..relay_conf(&upstream().await, Duration::from_secs(3))
    };

    let app = mon_app_with(conf).await;
    let request = Request::builder()
        .uri("/relay/forwarded")
        .header("host", "127.0.0.1")
        .header("x-forwarded-for", "1.2.3.4")
        .header("x-forwarded-proto", "https")
        .header("x-forwarded-host", "rollpie.com")
        .body(Body::empty())
        .unwrap();

    let response = app.oneshot(request).await.unwrap();
    assert_eq!(read(response).await, "1.2.3.4 https rollpie.com");
}


/// 前段がいない（CLIENT_IP_HEADER が空の）なら、相手が書いた forwarded 系は信用せず、mon が書き直す
#[tokio::test]
async fn forged_forwarded_headers_are_replaced() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let request = Request::builder()
        .uri("/relay/forwarded")
        .header("host", "127.0.0.1")
        .header("x-forwarded-for", "1.2.3.4")
        .header("x-forwarded-proto", "https")
        .header("x-forwarded-host", "evil.example")
        .body(Body::empty())
        .unwrap();

    // oneshot では TCP の相手がないので、x-forwarded-for は付かない
    let response = app.oneshot(request).await.unwrap();
    assert_eq!(read(response).await, "- http 127.0.0.1");
}


/// mon のセッションの id は中継先へ渡さない。ほかのクッキーは渡す
#[tokio::test]
async fn the_session_cookie_is_not_relayed() {
    let conf = Conf {
        session_store: "memory".to_owned(),
        ..relay_conf(&upstream().await, Duration::from_secs(3))
    };

    let app = mon_app_with(conf).await;
    let request = Request::builder()
        .uri("/relay/cookie")
        .header("cookie", format!("sid={}; theme=dark", "a".repeat(64)))
        .body(Body::empty())
        .unwrap();

    let response = app.oneshot(request).await.unwrap();
    assert_eq!(read(response).await, "theme=dark");
}


/// 中継先が mon のセッションのクッキーを書こうとしても通さない
#[tokio::test]
async fn the_upstream_cannot_set_the_session_cookie() {
    let conf = Conf {
        session_store: "memory".to_owned(),
        ..relay_conf(&upstream().await, Duration::from_secs(3))
    };

    let app = mon_app_with(conf).await;
    let response = app.oneshot(get_uri("/relay/set-cookie")).await.unwrap();
    let line: Vec<&str> = response
        .headers()
        .get_all("set-cookie")
        .iter()
        .map(|value| value.to_str().unwrap())
        .collect();

    assert_eq!(line, vec!["theme=dark; Path=/"]);
}


/// 中継先の 302 は追わない。追うと相手には最終の応答だけが届き、行き先が消える
#[tokio::test]
async fn a_redirect_is_not_followed() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;
    let response = app.oneshot(get_uri("/relay/go")).await.unwrap();

    assert_eq!(response.status(), StatusCode::FOUND);
    assert_eq!(response.headers().get("location").unwrap(), "/echo");
}


/// 上流がいない
#[tokio::test]
async fn no_upstream_is_502() {
    let app = mon_app("http://127.0.0.1:1", Duration::from_secs(3)).await;
    let response = app.oneshot(get_uri("/relay/echo")).await.unwrap();

    assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
}


/// つながるが返事をしない。受けるだけで応じない相手を立てる
#[tokio::test]
async fn a_silent_upstream_is_504() {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    // accept しないので、つながったまま何も返らない
    let _held = listener;

    let app = mon_app(
        &format!("http://127.0.0.1:{port}"),
        Duration::from_millis(300),
    )
    .await;
    let response = app.oneshot(get_uri("/relay/echo")).await.unwrap();

    assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT);
}


/// 読み切ってから返す作りに戻ると、ここで 1 秒待たされる
#[tokio::test]
async fn the_body_flows_before_the_upstream_finishes() {
    let app = mon_app(&upstream().await, Duration::from_secs(3)).await;

    let started = Instant::now();
    let response = app.oneshot(get_uri("/relay/slow")).await.unwrap();
    let mut chunks = response.into_body().into_data_stream();
    let first = chunks.next().await.unwrap().unwrap();
    let waited = started.elapsed();

    assert_eq!(first, Bytes::from_static(b"first"));
    assert!(
        waited < Duration::from_millis(500),
        "先頭が届くまで {waited:?} かかった。読み切ってから返している"
    );
}