nordisk/edda

crates/nornir-flight/tests/tls_roundtrip.rs

//! Acceptance for the TLS-terminating transport path (design §f1, `tls` feature):
//! a `do_put` → `do_get` round-trip driven **over a self-signed TLS channel**
//! returns the batches byte-identical, and the same holds with bearer auth
//! layered under TLS.
//!
//! This whole file compiles only with `--features tls` (it exercises
//! `FlightServer::serve_with_tls[_and_auth]` + `FlightConn::connect_tls`, which
//! are `#[cfg(feature = "tls")]`). It is RED-when-broken: if the TLS serve/connect
//! seam drops or mangles a batch, or fails to establish the encrypted channel,
//! these assertions fail. The cert is generated in-process with `rcgen` — no
//! `openssl` subprocess (LAW #9, zero-shell).

#![cfg(feature = "tls")]

use std::net::{SocketAddr, TcpListener};
use std::sync::Arc;
use std::time::Duration;

use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use nornir_flight::{FlightConn, FlightRoutes, FlightServer, MemFlightStore};
use tonic::transport::{Certificate, ClientTlsConfig, Identity, ServerTlsConfig};

/// Grab an OS-assigned free port (bind, read, drop — the server rebinds it).
fn free_addr() -> SocketAddr {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap()
}

/// `n` distinct two-column batches (Int64 + nullable Utf8).
fn sample_batches(n: usize) -> Vec<RecordBatch> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, true),
    ]));
    (0..n)
        .map(|i| {
            let base = i as i64 * 10;
            let ids = Int64Array::from(vec![base, base + 1, base + 2]);
            let names = StringArray::from(vec![Some(format!("r{i}-a")), None, Some(format!("r{i}-c"))]);
            RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(names)]).unwrap()
        })
        .collect()
}

/// A fresh self-signed cert/key for `localhost`, returned as (cert_pem, key_pem).
/// The one cert is BOTH the server identity AND (self-signed) the client's trust
/// root — a hermetic loopback PKI, generated purely in Rust.
fn self_signed_localhost() -> (String, String) {
    let ck = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).unwrap();
    // rcgen 0.14 renamed `CertifiedKey::key_pair` to `signing_key` (same `KeyPair`).
    (ck.cert.pem(), ck.signing_key.serialize_pem())
}

/// Start a TLS Flight server (fresh `MemFlightStore`) on a free port with `token`
/// auth (empty ⇒ open) using `cert_pem`/`key_pem` as its identity, and return
/// (`https://localhost:port`, client-trust-config) once it accepts connects.
async fn start_tls_server(token: &str, cert_pem: &str, key_pem: &str) -> (String, ClientTlsConfig) {
    let addr = free_addr();
    let routes = FlightRoutes::shared(Arc::new(MemFlightStore::new()));
    let tok = Arc::new(token.to_string());
    let server_tls = ServerTlsConfig::new().identity(Identity::from_pem(cert_pem, key_pem));
    tokio::spawn(async move {
        let _ = FlightServer::serve_with_tls_and_auth(routes, addr, tok, server_tls).await;
    });
    for _ in 0..200 {
        if tokio::net::TcpStream::connect(addr).await.is_ok() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    // The self-signed server cert IS the client's CA; SNI/verification host is
    // `localhost`, matching the cert's SAN.
    let client_tls = ClientTlsConfig::new()
        .ca_certificate(Certificate::from_pem(cert_pem))
        .domain_name("localhost");
    (format!("https://localhost:{}", addr.port()), client_tls)
}

#[tokio::test]
async fn tls_put_then_get_roundtrip_byte_identical() {
    let (cert, key) = self_signed_localhost();
    let (ep, client_tls) = start_tls_server("", &cert, &key).await;
    let mut c = FlightConn::connect_tls(ep, client_tls).await.unwrap();
    let batches = sample_batches(3);
    c.put("build/42", batches.clone()).await.unwrap();
    let got = c.get("build/42").await.unwrap();
    assert_eq!(
        got, batches,
        "do_get over TLS must return the do_put batches byte-identical"
    );
}

#[tokio::test]
async fn tls_exchange_roundtrip_byte_identical() {
    let (cert, key) = self_signed_localhost();
    let (ep, client_tls) = start_tls_server("", &cert, &key).await;
    let mut c = FlightConn::connect_tls(ep, client_tls).await.unwrap();
    let batches = sample_batches(4);
    let got = c.exchange("x/1", batches.clone()).await.unwrap();
    assert_eq!(got, batches, "do_exchange over TLS must echo the batches byte-identical");
}

#[tokio::test]
async fn tls_with_bearer_auth_roundtrip() {
    // TLS layered UNDER bearer auth: correct token + trusted cert round-trips.
    let token = "a-long-enough-secret-token-1234";
    let (cert, key) = self_signed_localhost();
    let (ep, client_tls) = start_tls_server(token, &cert, &key).await;
    let mut c = FlightConn::connect_tls(ep, client_tls)
        .await
        .unwrap()
        .with_bearer(token)
        .unwrap();
    let batches = sample_batches(2);
    c.put("k", batches.clone()).await.unwrap();
    assert_eq!(
        c.get("k").await.unwrap(),
        batches,
        "correct bearer over TLS must be accepted and round-trip"
    );
}

#[tokio::test]
async fn tls_missing_token_rejected() {
    // Even over a trusted TLS channel, a missing bearer is UNAUTHENTICATED — TLS
    // secures the pipe, auth guards the verb (RED if TLS silently drops auth).
    let token = "a-long-enough-secret-token-1234";
    let (cert, key) = self_signed_localhost();
    let (ep, client_tls) = start_tls_server(token, &cert, &key).await;
    let mut c = FlightConn::connect_tls(ep, client_tls).await.unwrap(); // no bearer
    let err = c
        .put("k", sample_batches(1))
        .await
        .expect_err("unauthenticated put over TLS MUST be rejected");
    let msg = err.to_string().to_lowercase();
    assert!(
        msg.contains("bearer") || msg.contains("unauthenticated"),
        "expected an auth rejection over TLS, got: {msg}"
    );
}