Rename project to quicnprotochat

crates/quicnprotochat-client/tests/mls_group.rs

@@ -0,0 +1,433 @@
+//! Integration test: full MLS group flow via Authentication Service + Delivery Service.
+//!
+//! Steps:
+//! - Start in-process AS and DS (Noise_XX + capnp-rpc) on a LocalSet.
+//! - Alice and Bob generate KeyPackages and upload to AS.
+//! - Alice fetches Bob's KeyPackage, creates a group, and invites Bob.
+//! - Welcome + application messages traverse the Delivery Service.
+//! - Both sides decrypt and confirm plaintext payloads.
+
+use std::{collections::VecDeque, sync::Arc, time::Duration};
+
+use anyhow::Context;
+use capnp::capability::Promise;
+use capnp_rpc::{rpc_twoparty_capnp::Side, twoparty, RpcSystem};
+use dashmap::DashMap;
+use quicnprotochat_core::{
+    handshake_initiator, handshake_responder, GroupMember, IdentityKeypair, NoiseKeypair,
+};
+use quicnprotochat_proto::{auth_capnp::authentication_service, delivery_capnp::delivery_service};
+use sha2::{Digest, Sha256};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
+
+type KeyPackageStore = Arc<DashMap<Vec<u8>, VecDeque<Vec<u8>>>>;
+type DeliveryStore = Arc<DashMap<Vec<u8>, VecDeque<Vec<u8>>>>;
+
+/// Full Alice↔Bob MLS round-trip against live AS + DS.
+#[tokio::test]
+async fn mls_group_end_to_end_round_trip() -> anyhow::Result<()> {
+    let local = tokio::task::LocalSet::new();
+
+    local
+        .run_until(async move {
+            let server_keypair = Arc::new(NoiseKeypair::generate());
+            let kp_store: KeyPackageStore = Arc::new(DashMap::new());
+            let ds_store: DeliveryStore = Arc::new(DashMap::new());
+
+            let as_addr =
+                spawn_as_server(2, Arc::clone(&server_keypair), Arc::clone(&kp_store)).await;
+            let ds_addr =
+                spawn_ds_server(2, Arc::clone(&server_keypair), Arc::clone(&ds_store)).await;
+
+            tokio::time::sleep(Duration::from_millis(10)).await;
+
+            let alice_id = Arc::new(IdentityKeypair::generate());
+            let bob_id = Arc::new(IdentityKeypair::generate());
+
+            let mut alice = GroupMember::new(Arc::clone(&alice_id));
+            let mut bob = GroupMember::new(Arc::clone(&bob_id));
+
+            let alice_kp = alice.generate_key_package()?;
+            let bob_kp = bob.generate_key_package()?;
+
+            let alice_as = connect_as(as_addr, &NoiseKeypair::generate()).await?;
+            let bob_as = connect_as(as_addr, &NoiseKeypair::generate()).await?;
+
+            upload_key_package(&alice_as, &alice_id.public_key_bytes(), &alice_kp).await?;
+            upload_key_package(&bob_as, &bob_id.public_key_bytes(), &bob_kp).await?;
+
+            let fetched_bob_kp = fetch_key_package(&alice_as, &bob_id.public_key_bytes()).await?;
+            anyhow::ensure!(
+                !fetched_bob_kp.is_empty(),
+                "AS must return Bob's KeyPackage"
+            );
+
+            alice.create_group(b"m3-integration")?;
+            let (_commit, welcome) = alice.add_member(&fetched_bob_kp)?;
+
+            let alice_ds = connect_ds(ds_addr, &NoiseKeypair::generate()).await?;
+            let bob_ds = connect_ds(ds_addr, &NoiseKeypair::generate()).await?;
+
+            enqueue(&alice_ds, &bob_id.public_key_bytes(), &welcome).await?;
+
+            let welcome_payloads = fetch_all(&bob_ds, &bob_id.public_key_bytes()).await?;
+            let welcome_bytes = welcome_payloads
+                .first()
+                .cloned()
+                .context("welcome must be present")?;
+            bob.join_group(&welcome_bytes)?;
+
+            let ct_ab = alice.send_message(b"hello bob")?;
+            enqueue(&alice_ds, &bob_id.public_key_bytes(), &ct_ab).await?;
+            let bob_msgs = fetch_all(&bob_ds, &bob_id.public_key_bytes()).await?;
+            let ab_plaintext = bob
+                .receive_message(bob_msgs.first().context("missing alice→bob payload")?)?
+                .context("alice→bob must be application message")?;
+            assert_eq!(ab_plaintext, b"hello bob");
+
+            let ct_ba = bob.send_message(b"hello alice")?;
+            enqueue(&bob_ds, &alice_id.public_key_bytes(), &ct_ba).await?;
+            let alice_msgs = fetch_all(&alice_ds, &alice_id.public_key_bytes()).await?;
+            let ba_plaintext = alice
+                .receive_message(alice_msgs.first().context("missing bob→alice payload")?)?
+                .context("bob→alice must be application message")?;
+            assert_eq!(ba_plaintext, b"hello alice");
+
+            Ok(())
+        })
+        .await
+}
+
+// ── Test helpers ────────────────────────────────────────────────────────────
+
+async fn spawn_as_server(
+    n_connections: usize,
+    keypair: Arc<NoiseKeypair>,
+    store: KeyPackageStore,
+) -> 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_as_connection(stream, kp, st).await;
+            });
+        }
+    });
+
+    addr
+}
+
+async fn serve_as_connection(
+    stream: TcpStream,
+    keypair: Arc<NoiseKeypair>,
+    store: KeyPackageStore,
+) {
+    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 service: authentication_service::Client = capnp_rpc::new_client(AuthService { store });
+
+    RpcSystem::new(Box::new(network), Some(service.client))
+        .await
+        .ok();
+}
+
+async fn spawn_ds_server(
+    n_connections: usize,
+    keypair: Arc<NoiseKeypair>,
+    store: DeliveryStore,
+) -> 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_ds_connection(stream, kp, st).await;
+            });
+        }
+    });
+
+    addr
+}
+
+async fn serve_ds_connection(stream: TcpStream, keypair: Arc<NoiseKeypair>, store: DeliveryStore) {
+    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 service: delivery_service::Client = capnp_rpc::new_client(DeliveryService { store });
+
+    RpcSystem::new(Box::new(network), Some(service.client))
+        .await
+        .ok();
+}
+
+async fn connect_as(
+    addr: std::net::SocketAddr,
+    noise_keypair: &NoiseKeypair,
+) -> anyhow::Result<authentication_service::Client> {
+    let stream = TcpStream::connect(addr)
+        .await
+        .with_context(|| format!("could not connect to AS at {addr}"))?;
+
+    let transport = handshake_initiator(stream, noise_keypair)
+        .await
+        .context("Noise handshake to AS failed")?;
+    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);
+    Ok(client)
+}
+
+async fn connect_ds(
+    addr: std::net::SocketAddr,
+    noise_keypair: &NoiseKeypair,
+) -> anyhow::Result<delivery_service::Client> {
+    let stream = TcpStream::connect(addr)
+        .await
+        .with_context(|| format!("could not connect to DS at {addr}"))?;
+
+    let transport = handshake_initiator(stream, noise_keypair)
+        .await
+        .context("Noise handshake to DS failed")?;
+    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: delivery_service::Client = rpc.bootstrap(Side::Server);
+    tokio::task::spawn_local(rpc);
+    Ok(client)
+}
+
+async fn upload_key_package(
+    as_client: &authentication_service::Client,
+    identity_key: &[u8],
+    package: &[u8],
+) -> anyhow::Result<()> {
+    let mut req = as_client.upload_key_package_request();
+    req.get().set_identity_key(identity_key);
+    req.get().set_package(package);
+
+    let resp = req
+        .send()
+        .promise
+        .await
+        .context("upload_key_package RPC failed")?;
+
+    let server_fp = resp
+        .get()
+        .context("upload_key_package: bad response")?
+        .get_fingerprint()
+        .context("upload_key_package: missing fingerprint")?
+        .to_vec();
+
+    let local_fp: Vec<u8> = Sha256::digest(package).to_vec();
+    anyhow::ensure!(server_fp == local_fp, "fingerprint mismatch");
+    Ok(())
+}
+
+async fn fetch_key_package(
+    as_client: &authentication_service::Client,
+    identity_key: &[u8],
+) -> anyhow::Result<Vec<u8>> {
+    let mut req = as_client.fetch_key_package_request();
+    req.get().set_identity_key(identity_key);
+
+    let resp = req
+        .send()
+        .promise
+        .await
+        .context("fetch_key_package RPC failed")?;
+
+    let pkg = resp
+        .get()
+        .context("fetch_key_package: bad response")?
+        .get_package()
+        .context("fetch_key_package: missing package")?
+        .to_vec();
+
+    Ok(pkg)
+}
+
+async fn enqueue(
+    ds_client: &delivery_service::Client,
+    recipient_key: &[u8],
+    payload: &[u8],
+) -> anyhow::Result<()> {
+    let mut req = ds_client.enqueue_request();
+    req.get().set_recipient_key(recipient_key);
+    req.get().set_payload(payload);
+    req.send().promise.await.context("enqueue RPC failed")?;
+    Ok(())
+}
+
+async fn fetch_all(
+    ds_client: &delivery_service::Client,
+    recipient_key: &[u8],
+) -> anyhow::Result<Vec<Vec<u8>>> {
+    let mut req = ds_client.fetch_request();
+    req.get().set_recipient_key(recipient_key);
+
+    let resp = req.send().promise.await.context("fetch RPC failed")?;
+
+    let list = resp
+        .get()
+        .context("fetch: bad response")?
+        .get_payloads()
+        .context("fetch: missing payloads")?;
+
+    let mut payloads = Vec::with_capacity(list.len() as usize);
+    for i in 0..list.len() {
+        payloads.push(list.get(i).context("fetch: payload read failed")?.to_vec());
+    }
+
+    Ok(payloads)
+}
+
+// ── Inline service implementations ─────────────────────────────────────────-
+
+struct AuthService {
+    store: KeyPackageStore,
+}
+
+impl authentication_service::Server for AuthService {
+    fn upload_key_package(
+        &mut self,
+        params: authentication_service::UploadKeyPackageParams,
+        mut results: authentication_service::UploadKeyPackageResults,
+    ) -> Promise<(), capnp::Error> {
+        let params = match params.get() {
+            Ok(p) => p,
+            Err(e) => return Promise::err(e),
+        };
+
+        let ik = match params.get_identity_key() {
+            Ok(v) => v.to_vec(),
+            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
+        };
+        let pkg = match params.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(())
+    }
+}
+
+struct DeliveryService {
+    store: DeliveryStore,
+}
+
+impl delivery_service::Server for DeliveryService {
+    fn enqueue(
+        &mut self,
+        params: delivery_service::EnqueueParams,
+        _results: delivery_service::EnqueueResults,
+    ) -> Promise<(), capnp::Error> {
+        let params = match params.get() {
+            Ok(p) => p,
+            Err(e) => return Promise::err(e),
+        };
+
+        let recipient = match params.get_recipient_key() {
+            Ok(v) => v.to_vec(),
+            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
+        };
+        let payload = match params.get_payload() {
+            Ok(v) => v.to_vec(),
+            Err(e) => return Promise::err(capnp::Error::failed(format!("{e}"))),
+        };
+
+        self.store.entry(recipient).or_default().push_back(payload);
+        Promise::ok(())
+    }
+
+    fn fetch(
+        &mut self,
+        params: delivery_service::FetchParams,
+        mut results: delivery_service::FetchResults,
+    ) -> Promise<(), capnp::Error> {
+        let recipient = match params.get() {
+            Ok(p) => match p.get_recipient_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 messages: Vec<Vec<u8>> = self
+            .store
+            .get_mut(&recipient)
+            .map(|mut q| q.drain(..).collect())
+            .unwrap_or_default();
+
+        let mut list = results.get().init_payloads(messages.len() as u32);
+        for (i, msg) in messages.iter().enumerate() {
+            list.set(i as u32, msg);
+        }
+
+        Promise::ok(())
+    }
+}