quicproquo/crates/noiseml-client/tests/auth_service.rs

//! Integration test: M2 Authentication Service — KeyPackage upload + fetch.
//!
//! All tests run inside a single `tokio::task::LocalSet` so that `spawn_local`
//! can be used for capnp-rpc tasks (which are `!Send` due to internal `Rc` use).

use std::{collections::VecDeque, sync::Arc};

use capnp::capability::Promise;
use capnp_rpc::{RpcSystem, rpc_twoparty_capnp::Side, twoparty};
use dashmap::DashMap;
use noiseml_core::{
    IdentityKeypair, NoiseKeypair, generate_key_package, handshake_initiator,
    handshake_responder,
};
use noiseml_proto::auth_capnp::authentication_service;
use sha2::{Digest, Sha256};
use tokio::net::{TcpListener, TcpStream};
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};

// ── Types ─────────────────────────────────────────────────────────────────────

type Store = Arc<DashMap<Vec<u8>, VecDeque<Vec<u8>>>>;

// ── Inline AS server implementation ──────────────────────────────────────────

struct TestAuthService {
    store: Store,
}

impl authentication_service::Server for TestAuthService {
    fn upload_key_package(
        &mut self,
        params: authentication_service::UploadKeyPackageParams,
        mut results: authentication_service::UploadKeyPackageResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(v) => v,
            Err(e) => return Promise::err(e),
        };
        let ik = match p.get_identity_key() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
        };
        let pkg = match p.get_package() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
        };
        let fp: Vec<u8> = Sha256::digest(&pkg).to_vec();
        self.store.entry(ik).or_default().push_back(pkg);
        results.get().set_fingerprint(&fp);
        Promise::ok(())
    }

    fn fetch_key_package(
        &mut self,
        params: authentication_service::FetchKeyPackageParams,
        mut results: authentication_service::FetchKeyPackageResults,
    ) -> Promise<(), capnp::Error> {
        let ik = match params.get() {
            Ok(p) => match p.get_identity_key() {
                Ok(v) => v.to_vec(),
                Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
            },
            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
        };
        let pkg = self
            .store
            .get_mut(&ik)
            .and_then(|mut q| q.pop_front())
            .unwrap_or_default();
        results.get().set_package(&pkg);
        Promise::ok(())
    }
}

// ── Test helpers ──────────────────────────────────────────────────────────────

/// Spawn a server that accepts `n_connections` and returns the bound address.
///
/// Must be called from within a `LocalSet` context so that the internal
/// `spawn_local` calls are associated with the correct LocalSet.
async fn spawn_server(
    n_connections: usize,
    keypair: Arc<NoiseKeypair>,
    store: Store,
) -> std::net::SocketAddr {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();

    tokio::task::spawn_local(async move {
        for _ in 0..n_connections {
            let (stream, _) = listener.accept().await.unwrap();
            let kp = Arc::clone(&keypair);
            let st = Arc::clone(&store);
            tokio::task::spawn_local(async move {
                serve_one(stream, kp, st).await;
            });
        }
    });

    addr
}

/// Handle a single Noise + capnp-rpc server connection.
async fn serve_one(stream: TcpStream, keypair: Arc<NoiseKeypair>, store: Store) {
    let transport = handshake_responder(stream, &keypair).await.unwrap();
    let (reader, writer) = transport.into_capnp_io();
    let network = twoparty::VatNetwork::new(
        reader.compat(),
        writer.compat_write(),
        Side::Server,
        Default::default(),
    );
    let svc: authentication_service::Client =
        capnp_rpc::new_client(TestAuthService { store });
    let rpc = RpcSystem::new(Box::new(network), Some(svc.client));
    tokio::task::spawn_local(rpc).await.ok();
}

/// Connect and return a client stub. Must run inside a LocalSet.
async fn connect_client(addr: std::net::SocketAddr) -> authentication_service::Client {
    let kp = NoiseKeypair::generate();
    let stream = TcpStream::connect(addr).await.unwrap();
    let transport = handshake_initiator(stream, &kp).await.unwrap();
    let (reader, writer) = transport.into_capnp_io();
    let network = twoparty::VatNetwork::new(
        reader.compat(),
        writer.compat_write(),
        Side::Client,
        Default::default(),
    );
    let mut rpc = RpcSystem::new(Box::new(network), None);
    let client: authentication_service::Client = rpc.bootstrap(Side::Server);
    tokio::task::spawn_local(rpc);
    client
}

// ── Tests ─────────────────────────────────────────────────────────────────────

/// Alice uploads a KeyPackage; Bob fetches it. Fingerprints must match.
#[tokio::test]
async fn upload_then_fetch_fingerprints_match() {
    let local = tokio::task::LocalSet::new();
    local
        .run_until(async move {
            let store: Store = Arc::new(DashMap::new());
            let server_kp = Arc::new(NoiseKeypair::generate());

            // Server accepts 2 connections: one for Alice (upload), one for Bob (fetch).
            let addr = spawn_server(2, Arc::clone(&server_kp), Arc::clone(&store)).await;
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;

            // Alice: generate KeyPackage and upload it.
            let alice_identity = IdentityKeypair::generate();
            let (tls_bytes, local_fp) = generate_key_package(&alice_identity).unwrap();

            let alice = connect_client(addr).await;
            let mut req = alice.upload_key_package_request();
            req.get().set_identity_key(&alice_identity.public_key_bytes());
            req.get().set_package(&tls_bytes);
            let resp = req.send().promise.await.unwrap();
            let server_fp = resp.get().unwrap().get_fingerprint().unwrap().to_vec();

            assert_eq!(local_fp, server_fp, "server fingerprint must match local");

            // Bob: fetch Alice's package by her identity key.
            let bob = connect_client(addr).await;
            let mut req2 = bob.fetch_key_package_request();
            req2.get().set_identity_key(&alice_identity.public_key_bytes());
            let resp2 = req2.send().promise.await.unwrap();
            let fetched = resp2.get().unwrap().get_package().unwrap().to_vec();

            assert!(!fetched.is_empty(), "fetched package must not be empty");
            assert_eq!(fetched, tls_bytes, "fetched bytes must match uploaded bytes");

            let fetched_fp: Vec<u8> = Sha256::digest(&fetched).to_vec();
            assert_eq!(fetched_fp, local_fp, "fetched fingerprint must match uploaded");
        })
        .await;
}

/// Fetching a non-existent key returns empty bytes.
#[tokio::test]
async fn fetch_nonexistent_key_returns_empty() {
    let local = tokio::task::LocalSet::new();
    local
        .run_until(async move {
            let store: Store = Arc::new(DashMap::new());
            let server_kp = Arc::new(NoiseKeypair::generate());
            let addr = spawn_server(1, server_kp, store).await;
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;

            let client = connect_client(addr).await;
            let mut req = client.fetch_key_package_request();
            req.get().set_identity_key(&[0xAAu8; 32]);
            let resp = req.send().promise.await.unwrap();
            let pkg = resp.get().unwrap().get_package().unwrap().to_vec();

            assert!(pkg.is_empty(), "unknown identity must return empty package");
        })
        .await;
}

/// Uploading two packages and fetching twice returns them in FIFO order.
#[tokio::test]
async fn packages_consumed_in_fifo_order() {
    let local = tokio::task::LocalSet::new();
    local
        .run_until(async move {
            let store: Store = Arc::new(DashMap::new());

            // Pre-populate the store directly.
            let key = vec![0x01u8; 32];
            store
                .entry(key.clone())
                .or_default()
                .extend([vec![1u8, 2, 3], vec![4u8, 5, 6]]);

            let server_kp = Arc::new(NoiseKeypair::generate());
            // Server accepts 2 connections for the 2 fetches.
            let addr = spawn_server(2, server_kp, Arc::clone(&store)).await;
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;

            let client1 = connect_client(addr).await;
            let mut req1 = client1.fetch_key_package_request();
            req1.get().set_identity_key(&key);
            let pkg1 = req1
                .send()
                .promise
                .await
                .unwrap()
                .get()
                .unwrap()
                .get_package()
                .unwrap()
                .to_vec();
            assert_eq!(pkg1, vec![1u8, 2, 3], "first fetch must return first package");

            let client2 = connect_client(addr).await;
            let mut req2 = client2.fetch_key_package_request();
            req2.get().set_identity_key(&key);
            let pkg2 = req2
                .send()
                .promise
                .await
                .unwrap()
                .get()
                .unwrap()
                .get_package()
                .unwrap()
                .to_vec();
            assert_eq!(pkg2, vec![4u8, 5, 6], "second fetch must return second package");
        })
        .await;
}