relay.rs

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

use axum::Router;
use axum::body::{Body, Bytes, to_bytes};
use axum::extract::{Path, Request};
use axum::http::{HeaderMap, 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))
}


/// 受け取った 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()
}


async fn mon_app(relay_origin: &str, relay_timeout: Duration) -> Router {
    let conf = Conf {
        relay_origin: relay_origin.to_owned(),
        relay_timeout,
        session_store: "none".to_owned(),
        ..Conf::default()
    };

    let database = database::connect("sqlite::memory:").await.unwrap();
    let store = session::store(&conf, &database);
    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");
}


/// すでにあるものには触らない。前段が付けたものを書き換えると、数える位置がずれる
#[tokio::test]
async fn the_existing_forwarded_headers_are_kept() {
    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", "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");
}


/// 中継先の 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:?} かかった。読み切ってから返している"
    );
}