trz-gateway-server 0.2.10

use std::error::Error;
use std::io::ErrorKind;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use std::time::Instant;

use mime::APPLICATION_JSON;
use openssl::asn1::Asn1Time;
use openssl::pkey::HasPublic;
use openssl::pkey::PKeyRef;
use reqwest::Response;
use reqwest::StatusCode;
use reqwest::header::CONTENT_TYPE;
use tempfile::TempDir;
use terrazzo_fixture::Fixture;
use tokio::io::AsyncReadExt as _;
use tokio::io::AsyncWriteExt as _;
use tokio::net::TcpStream;
use tracing::debug;
use tracing::info;
use trz_gateway_common::api::tunnel::GetCertificateRequest;
use trz_gateway_common::certificate_info::CertificateInfo;
use trz_gateway_common::dynamic_config::DynamicConfig;
use trz_gateway_common::security_configuration::SecurityConfig;
use trz_gateway_common::security_configuration::certificate::CertificateConfig;
use trz_gateway_common::security_configuration::certificate::pem::PemCertificate;
use trz_gateway_common::security_configuration::trusted_store::pem::PemTrustedStore;
use trz_gateway_common::tracing::test_utils::enable_tracing_for_tests;
use trz_gateway_common::x509::PemString as _;
use trz_gateway_common::x509::ca::make_intermediate;
use trz_gateway_common::x509::cert::make_cert;
use trz_gateway_common::x509::key::make_key;
use trz_gateway_common::x509::name::CertitficateName;
use trz_gateway_common::x509::validity::Validity;

use super::Server;
use super::gateway_config::GatewayConfig;
use super::gateway_config::Ports;
use super::root_ca_configuration;
use super::root_ca_configuration::RootCaConfigError;
use crate::auth_code::AuthCode;
use crate::server::HTTP_TIMEOUT;

const ROOT_CA_FILENAME: CertificateInfo<&str> = CertificateInfo {
    certificate: "root-ca-cert.pem",
    private_key: "root-ca-key.pem",
};

#[tokio::test]
async fn status() -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let _client = make_client(&config).await?;

    let () = handle.stop("End of test").await?;
    Ok(())
}

#[tokio::test]
async fn certificate_http() -> Result<(), Box<dyn Error>> {
    certificate("http").await
}

#[tokio::test]
async fn certificate_https() -> Result<(), Box<dyn Error>> {
    certificate("https").await
}

async fn certificate(scheme: &str) -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let client = make_client(&config).await?;

    let private_key = make_key()?;
    let response = send_certificate_request(
        &config,
        client,
        scheme,
        GetCertificateRequest {
            auth_code: AuthCode::current(),
            public_key: &private_key,
            name: "Test client ID".into(),
        },
    )
    .await?;
    assert_eq!(StatusCode::OK, response.status());

    let pem = response.text().await?;
    let (rest, certificate) = x509_parser::pem::parse_x509_pem(pem.as_bytes())?;
    assert_eq!([0; 0], rest);
    let certificate = certificate.parse_x509()?;
    assert_eq!("CN=Test Root CA", certificate.issuer().to_string());
    assert_eq!("CN=Test client ID", certificate.subject().to_string());

    let () = handle.stop("End of test").await?;
    Ok(())
}

#[tokio::test]
async fn invalid_auth_code() -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let client = make_client(&config).await?;

    let private_key = make_key()?;
    let response = send_certificate_request(
        &config,
        client,
        "https",
        GetCertificateRequest {
            auth_code: AuthCode::from("invalid-code"),
            public_key: &private_key,
            name: "Test client ID".into(),
        },
    )
    .await?;
    assert_eq!(StatusCode::FORBIDDEN, response.status());

    let body = response.text().await?;
    assert_eq!("[InvalidAuthCode] AuthCode is invalid", body);

    let () = handle.stop("End of test").await?;
    Ok(())
}

#[tokio::test]
async fn tunnel() -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let client = make_client(&config).await?;

    let private_key = make_key()?;
    let response = send_certificate_request(
        &config,
        client,
        "https",
        GetCertificateRequest {
            auth_code: AuthCode::current(),
            public_key: &private_key,
            name: "Test client ID".into(),
        },
    )
    .await?;
    assert_eq!(StatusCode::OK, response.status());

    let _pem = response.text().await?;

    let () = handle.stop("End of test").await?;
    Ok(())
}

#[tokio::test]
async fn idle_tcp_connection_times_out() -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let _client = make_client(&config).await?;

    let mut stream =
        TcpStream::connect((config.host().as_str(), *config.ports().first().unwrap())).await?;
    let start = Instant::now();
    let mut buffer = [0; 1];
    let read_result = tokio::time::timeout(Duration::from_secs(3), stream.read(&mut buffer)).await;
    let elapsed = start.elapsed();

    debug!("Read result: {read_result:?}, elapsed time: {elapsed:?}");
    match read_result {
        Ok(Err(error))
            if matches!(
                error.kind(),
                ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted
            ) => {}
        other => {
            panic!("Expected the idle TCP connection to fail with RST after timeout, got {other:?}")
        }
    }
    assert!(
        elapsed >= Duration::from_millis(900),
        "Connection closed too early after {elapsed:?}",
    );

    let () = handle.stop("End of test").await?;
    Ok(())
}

#[tokio::test]
async fn http_connection_times_out() -> Result<(), Box<dyn Error>> {
    let _use_temp_dir = use_temp_dir();
    let config = TestConfig::new();
    let (_server, handle, _crash) = Server::run(config.clone()).await?;

    let _client = make_client(&config).await?;

    let mut stream =
        TcpStream::connect((config.host().as_str(), *config.ports().first().unwrap())).await?;
    send_plaintext_keep_alive_request(&mut stream, &config).await?;
    assert_eq!(
        StatusCode::NOT_FOUND,
        read_http_response_status(&mut stream).await?
    );

    let start = Instant::now();
    let mut buffer = [0; 1];
    let idle_read_result = tokio::time::timeout(
        HTTP_TIMEOUT + Duration::from_secs(1),
        stream.read(&mut buffer),
    )
    .await;
    let elapsed = start.elapsed();
    info!("Timeout: {idle_read_result:?}, elapsed: {elapsed:?}");

    assert!(
        elapsed >= HTTP_TIMEOUT - Duration::from_millis(100),
        "Connection closed too early after {elapsed:?}",
    );

    if let Ok(Ok(0)) = idle_read_result {
    } else {
        panic!("Unexpected read result: {idle_read_result:?}")
    }

    let () = handle.stop("End of test").await?;
    Ok(())
}

async fn make_client(config: &TestConfig) -> Result<reqwest::Client, Box<dyn Error>> {
    let client = {
        use reqwest::tls::Certificate;
        let trusted_root = Certificate::from_pem(
            config
                .tls_config
                .trusted_store
                .root_certificates_pem
                .as_bytes(),
        )?;
        reqwest::ClientBuilder::new()
            .add_root_certificate(trusted_root)
            .build()?
    };
    let mut wait = Duration::from_millis(1);
    while wait < Duration::from_secs(5) {
        let t = Instant::now();
        let request = client.get(format!("https://{}:{}/status", config.host(), config.port));
        match request.send().await {
            Ok(response) => match response.text().await.as_deref() {
                Ok("UP") => return Ok(client),
                response => debug!("Unexpected response: {response:?}"),
            },
            Err(error) => debug!("Failed: {error:?}"),
        }
        tokio::time::sleep(wait).await;
        wait = Duration::max(t.elapsed(), wait) * 2;
    }
    panic!("Failed to connect")
}

async fn send_certificate_request(
    config: &TestConfig,
    client: reqwest::Client,
    scheme: &str,
    request: GetCertificateRequest<AuthCode, &PKeyRef<impl HasPublic>>,
) -> Result<Response, Box<dyn Error>> {
    let public_key = request.public_key.public_key_to_pem().pem_string()?;
    let request = client
        .get(format!(
            "{scheme}://{}:{}/remote/certificate",
            config.host(),
            config.port
        ))
        .header(CONTENT_TYPE, APPLICATION_JSON.as_ref())
        .body(serde_json::to_string(&GetCertificateRequest {
            auth_code: request.auth_code,
            public_key,
            name: request.name,
        })?);
    Ok(request.send().await?)
}

async fn send_plaintext_keep_alive_request(
    stream: &mut TcpStream,
    config: &TestConfig,
) -> Result<(), Box<dyn Error>> {
    stream
        .write_all(
            format!(
                "GET /404-not-found HTTP/1.1\r\nHost: {}\r\nConnection: keep-alive\r\n\r\n",
                config.host()
            )
            .as_bytes(),
        )
        .await?;
    stream.flush().await?;
    Ok(())
}

async fn read_http_response_status(stream: &mut TcpStream) -> Result<StatusCode, Box<dyn Error>> {
    let mut response = Vec::new();
    let header_end = loop {
        let mut chunk = [0; 1024];
        let bytes_read =
            tokio::time::timeout(Duration::from_secs(3), stream.read(&mut chunk)).await??;
        if bytes_read == 0 {
            return Err("Connection closed before complete HTTP response".into());
        }
        response.extend_from_slice(&chunk[..bytes_read]);
        if let Some(header_end) = response.windows(4).position(|window| window == b"\r\n\r\n") {
            break header_end + 4;
        }
    };

    let header_text = std::str::from_utf8(&response[..header_end])?;
    let mut lines = header_text.split("\r\n");
    let status_line = lines.next().ok_or("Missing HTTP status line")?;
    let status = status_line
        .split_whitespace()
        .nth(1)
        .ok_or("Missing HTTP status code")?
        .parse::<u16>()?;
    let content_length = lines
        .filter_map(|line| line.split_once(':'))
        .find_map(|(name, value)| {
            name.eq_ignore_ascii_case("content-length")
                .then(|| value.trim().parse::<usize>())
        })
        .transpose()?
        .unwrap_or(0);

    while response.len() < header_end + content_length {
        let mut chunk = [0; 1024];
        let bytes_read =
            tokio::time::timeout(Duration::from_secs(3), stream.read(&mut chunk)).await??;
        if bytes_read == 0 {
            return Err("Connection closed before complete HTTP body".into());
        }
        response.extend_from_slice(&chunk[..bytes_read]);
    }

    Ok(StatusCode::from_u16(status)?)
}

#[derive(Debug)]
struct TestConfig {
    port: u16,
    root_ca: Arc<PemCertificate>,
    tls_config: <TestConfig as GatewayConfig>::TlsConfig,
}

impl TestConfig {
    fn new() -> Arc<Self> {
        enable_tracing_for_tests();
        let root_ca = make_root_ca().expect("root_ca_config()");
        let tls_config = make_tls_config().expect("tls_config()");
        Arc::new(Self {
            port: portpicker::pick_unused_port().expect("pick_unused_port()"),
            root_ca,
            tls_config,
        })
    }
}

impl GatewayConfig for TestConfig {
    fn enable_tracing(&self) -> bool {
        false
    }

    fn host(&self) -> String {
        "localhost".into()
    }

    fn ports(&self) -> impl Ports + 'static {
        vec![self.port]
    }

    type RootCaConfig = Arc<PemCertificate>;
    fn root_ca(&self) -> Self::RootCaConfig {
        self.root_ca.clone()
    }

    type TlsConfig = Arc<SecurityConfig<PemTrustedStore, PemCertificate>>;
    fn tls(&self) -> Self::TlsConfig {
        self.tls_config.clone()
    }

    type ClientCertificateIssuerConfig = Arc<DynamicConfig<Self::TlsConfig>>;
    fn client_certificate_issuer(&self) -> Self::ClientCertificateIssuerConfig {
        Arc::new(DynamicConfig::from(self.tls_config.clone()))
    }
}

fn make_root_ca() -> Result<Arc<PemCertificate>, RootCaConfigError> {
    let temp_dir = TEMP_DIR.get();

    static MUTEX: std::sync::Mutex<()> = Mutex::new(());
    let _lock = MUTEX.lock().unwrap();
    let root_ca = root_ca_configuration::load_root_ca(
        CertitficateName {
            common_name: Some("Test Root CA"),
            ..CertitficateName::default()
        },
        ROOT_CA_FILENAME.map(|filename| temp_dir.path().join(filename)),
        Validity { from: 0, to: 365 }
            .try_map(Asn1Time::days_from_now)
            .expect("Asn1Time::days_from_now")
            .as_deref()
            .try_into()
            .expect("Asn1Time to SystemTime"),
    )?;
    Ok(Arc::new(root_ca))
}

fn make_tls_config() -> Result<<TestConfig as GatewayConfig>::TlsConfig, Box<dyn Error>> {
    let root_ca = make_root_ca()?;
    let root_certificate = root_ca.certificate()?;
    let root_certificate_pem = root_ca.certificate_pem.clone();
    let validity = root_certificate.certificate.as_ref().try_into()?;

    let intermediate = make_intermediate(
        (*root_certificate).as_ref(),
        CertitficateName {
            organization: Some("Terrazzo Test"),
            common_name: Some("Intermediate CA"),
            ..CertitficateName::default()
        },
        validity,
    )?;

    let certificate_key = make_key()?;
    let certificate = make_cert(
        intermediate.as_ref(),
        CertitficateName {
            organization: Some("Terrazzo Test"),
            common_name: Some("localhost"),
            ..CertitficateName::default()
        },
        validity,
        &certificate_key.public_key_to_pem().pem_string()?,
        vec![],
    )?;

    Ok(Arc::new(SecurityConfig {
        trusted_store: PemTrustedStore {
            root_certificates_pem: root_certificate_pem,
        },
        certificate: PemCertificate {
            intermediates_pem: intermediate.certificate.to_pem()?.pem_string()?,
            certificate_pem: certificate.to_pem()?.pem_string()?,
            private_key_pem: certificate_key.private_key_to_pem_pkcs8()?.pem_string()?,
        },
    }))
}

static TEMP_DIR: Fixture<TempDir> = Fixture::new();

fn use_temp_dir() -> Arc<TempDir> {
    use std::sync::atomic::AtomicI32;
    use std::sync::atomic::Ordering::SeqCst;
    static NEXT: AtomicI32 = AtomicI32::new(0);
    TEMP_DIR.get_or_init(|| {
        tempfile::Builder::new()
            .suffix(&NEXT.fetch_add(1, SeqCst).to_string())
            .tempdir()
            .inspect(|temp_dir| debug!("Using tempprary folder {}", temp_dir.path().display()))
            .expect("TempDir::new()")
    })
}