kvitt/server/crates/kvitt-server/tests/sync_api/serve.rs

//! `serve` as deployed: the router with `ConnectInfo<SocketAddr>` (so the rate limiter sees each client's address,
//! SPEC §8.6), `/i/` with an empty code as the HTML 404 page, and `rotate-epoch`, `migrate`, `healthcheck` and the
//! SIGTERM drain of `serve` as the real binary (SPEC §10.4).

use std::net::SocketAddr;
use std::process::Command;

use axum::body::Body;
use axum::http::{Method, Request};
use kvitt_server::cli::make_service;
use tokio::net::TcpListener;

use crate::common::*;

#[tokio::test]
async fn two_client_addresses_get_separate_rate_limit_buckets_end_to_end() {
    let app = TestApp::empty().await;
    // One service, so one rate limiter, reached from two client addresses: 127.0.0.1 and ::1.
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let v4 = listener.local_addr().unwrap();
    let listener6 = TcpListener::bind("[::1]:0").await.unwrap();
    let v6 = listener6.local_addr().unwrap();
    let service = make_service(app.state.clone());
    let service6 = service.clone();
    tokio::spawn(async move { axum::serve(listener, service).await.unwrap() });
    tokio::spawn(async move { axum::serve(listener6, service6).await.unwrap() });

    let client = reqwest::Client::new();
    let get = |addr: SocketAddr| {
        let client = client.clone();
        async move {
            client
                .get(format!("http://{addr}/health"))
                .send()
                .await
                .unwrap()
                .status()
                .as_u16()
        }
    };
    for _ in 0..60 {
        assert_eq!(get(v4).await, 200);
    }
    assert_eq!(
        get(v4).await,
        429,
        "the 61st request from 127.0.0.1 in the minute"
    );
    assert_eq!(get(v6).await, 200, "::1 has its own bucket");
}

#[tokio::test]
async fn empty_invite_code_is_the_html_404_page() {
    let app = TestApp::empty().await;
    let reply = app
        .send(
            Request::builder()
                .method(Method::GET)
                .uri("/i/")
                .body(Body::empty())
                .unwrap(),
        )
        .await;
    assert_eq!(reply.status, 404);
    assert_eq!(reply.headers["content-type"], "text/html; charset=utf-8");
    assert!(reply.text().contains("<p>404</p>"));
}

#[tokio::test]
async fn rotate_epoch_sets_a_new_random_epoch_and_exits_0() {
    let app = TestApp::empty().await;
    let epoch = || async {
        sqlx::query_scalar::<_, String>("SELECT value FROM meta WHERE key = 'sync_epoch'")
            .fetch_one(&app.state.db)
            .await
            .unwrap()
    };
    let before = epoch().await;
    let output = Command::new(env!("CARGO_BIN_EXE_kvitt-server"))
        .arg("rotate-epoch")
        .env("KVITT_DATA_DIR", app.data_dir())
        .env("KVITT_PUBLIC_URL", PUBLIC_URL)
        .env("RUST_LOG", "warn")
        .output()
        .unwrap();
    assert_eq!(output.status.code(), Some(0), "{output:?}");
    let after = epoch().await;
    assert_ne!(before, after);
    assert_eq!(after.len(), 32);
    assert!(after.chars().all(|c| matches!(c, '0'..='9' | 'a'..='f')));
    assert!(String::from_utf8_lossy(&output.stdout).contains(&after));
}

/// The real binary with a data directory and listen address of the test's own; push disabled (no credentials).
fn binary(dir: &std::path::Path, bind: SocketAddr) -> Command {
    let mut command = Command::new(env!("CARGO_BIN_EXE_kvitt-server"));
    command
        .env("KVITT_DATA_DIR", dir)
        .env("KVITT_BIND", bind.to_string())
        .env("KVITT_PUBLIC_URL", PUBLIC_URL)
        .env("KVITT_FCM_CREDENTIALS", dir.join("no-such-key.json"))
        .env("RUST_LOG", "warn");
    command
}

/// A loopback address with a port nothing listens on (bound once, then released).
fn free_addr() -> SocketAddr {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    listener.local_addr().unwrap()
}

#[tokio::test]
async fn migrate_creates_the_database_at_the_current_schema_and_is_idempotent() {
    let dir = TempDir::new();
    let data = dir.0.join("data");
    for _ in 0..2 {
        let output = binary(&data, free_addr()).arg("migrate").output().unwrap();
        assert_eq!(output.status.code(), Some(0), "{output:?}");
    }
    // Opened without migrating, so what is read is what `migrate` wrote.
    let pool = kvitt_server::db::connect(&data.join("kvitt.sqlite"))
        .await
        .unwrap();
    let version: String = sqlx::query_scalar("SELECT value FROM meta WHERE key = 'schema_version'")
        .fetch_one(&pool)
        .await
        .unwrap();
    assert_eq!(version, kvitt_server::db::SCHEMA_VERSION.to_string());
    let users: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM users")
        .fetch_one(&pool)
        .await
        .unwrap();
    assert_eq!(users, 0);
    pool.close().await;
}

#[tokio::test]
async fn healthcheck_exits_0_on_200_and_1_on_another_status_or_no_server() {
    let dir = TempDir::new();
    let healthcheck = |addr: SocketAddr| binary(&dir.0, addr).arg("healthcheck").output().unwrap();

    // The real router: 200.
    let app = TestApp::empty().await;
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let up = listener.local_addr().unwrap();
    let service = make_service(app.state.clone());
    tokio::spawn(async move { axum::serve(listener, service).await.unwrap() });
    let ok = tokio::task::spawn_blocking({
        let dir = dir.0.clone();
        move || binary(&dir, up).arg("healthcheck").output().unwrap()
    })
    .await
    .unwrap();
    assert_eq!(ok.status.code(), Some(0), "{ok:?}");

    // A server answering 503 (as E-1 does when the database is unreachable).
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let unhealthy = listener.local_addr().unwrap();
    tokio::spawn(async move {
        loop {
            let (mut stream, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 1024];
            let _ = tokio::io::AsyncReadExt::read(&mut stream, &mut buf).await;
            let _ = tokio::io::AsyncWriteExt::write_all(
                &mut stream,
                b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
            )
            .await;
        }
    });
    let failed = tokio::task::spawn_blocking({
        let dir = dir.0.clone();
        move || binary(&dir, unhealthy).arg("healthcheck").output().unwrap()
    })
    .await
    .unwrap();
    assert_eq!(failed.status.code(), Some(1), "{failed:?}");
    assert!(String::from_utf8_lossy(&failed.stderr).contains("returned 503"));

    // Nothing listening.
    let down = healthcheck(free_addr());
    assert_eq!(down.status.code(), Some(1), "{down:?}");
}

#[cfg(unix)]
mod drain {
    use std::net::SocketAddr;
    use std::process::{Child, ExitStatus, Stdio};
    use std::time::{Duration, Instant};

    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpStream;

    use super::{binary, free_addr};
    use crate::common::TempDir;
    use kvitt_server::cli::SHUTDOWN_DRAIN;

    /// The body of an E-3 request whose last byte the test holds back.
    const BODY: &[u8] = br#"{"code":"0000-0000-0000","device_name":"x"}"#;

    /// A running `serve`, killed and reaped on drop if the test did not see it exit.
    struct Serve(Child);

    impl Drop for Serve {
        fn drop(&mut self) {
            if self.0.try_wait().ok().flatten().is_none() {
                let _ = self.0.kill();
                let _ = self.0.wait();
            }
        }
    }

    /// `serve` on a free port, once `/health` answers 200.
    async fn start(dir: &TempDir) -> (Serve, SocketAddr) {
        let addr = free_addr();
        let child = Serve(
            binary(&dir.0, addr)
                .arg("serve")
                .stdout(Stdio::null())
                .stderr(Stdio::null())
                .spawn()
                .unwrap(),
        );
        let deadline = Instant::now() + Duration::from_secs(30);
        loop {
            if let Ok(reply) = reqwest::get(format!("http://{addr}/health")).await
                && reply.status() == 200
            {
                return (child, addr);
            }
            assert!(Instant::now() < deadline, "serve did not come up");
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    /// An E-3 request in flight: the head and all but the last byte of the body sent.
    async fn in_flight(addr: SocketAddr) -> TcpStream {
        let mut stream = TcpStream::connect(addr).await.unwrap();
        let head = format!(
            "POST /invites/accept HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/json\r\n\
             Content-Length: {}\r\nConnection: close\r\n\r\n",
            BODY.len()
        );
        stream.write_all(head.as_bytes()).await.unwrap();
        stream.write_all(&BODY[..BODY.len() - 1]).await.unwrap();
        tokio::time::sleep(Duration::from_millis(200)).await;
        stream
    }

    fn sigterm(child: &Serve) {
        let status = std::process::Command::new("kill")
            .args(["-TERM", &child.0.id().to_string()])
            .status()
            .unwrap();
        assert!(status.success());
    }

    async fn wait(child: &mut Serve, within: Duration) -> ExitStatus {
        let deadline = Instant::now() + within;
        loop {
            if let Some(status) = child.0.try_wait().unwrap() {
                return status;
            }
            assert!(Instant::now() < deadline, "serve did not exit");
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    }

    #[tokio::test]
    async fn sigterm_lets_an_in_flight_request_finish_then_exits_0() {
        let dir = TempDir::new();
        let (mut child, addr) = start(&dir).await;
        let mut stream = in_flight(addr).await;

        sigterm(&child);
        tokio::time::sleep(Duration::from_millis(500)).await;
        assert!(
            child.0.try_wait().unwrap().is_none(),
            "exited with a request in flight"
        );
        assert!(
            TcpStream::connect(addr).await.is_err(),
            "new connections are no longer accepted"
        );

        stream.write_all(&BODY[BODY.len() - 1..]).await.unwrap();
        let mut reply = Vec::new();
        stream.read_to_end(&mut reply).await.unwrap();
        let reply = String::from_utf8_lossy(&reply);
        assert!(reply.starts_with("HTTP/1.1 400"), "{reply}");
        assert!(reply.contains("invite_invalid"), "{reply}");

        assert_eq!(
            wait(&mut child, Duration::from_secs(5)).await.code(),
            Some(0)
        );
    }

    #[tokio::test]
    async fn sigterm_stops_waiting_after_the_drain_period_and_exits_0() {
        let dir = TempDir::new();
        let (mut child, addr) = start(&dir).await;
        let _stuck = in_flight(addr).await;

        let signalled = Instant::now();
        sigterm(&child);
        let status = wait(&mut child, SHUTDOWN_DRAIN + Duration::from_secs(5)).await;
        assert_eq!(status.code(), Some(0));
        assert!(
            signalled.elapsed() >= SHUTDOWN_DRAIN - Duration::from_millis(500),
            "exited after {:?}, before the drain period",
            signalled.elapsed()
        );
    }
}