quicnprotochat-lang-clients/crates/quicnprotochat-bindings/src/lib.rs

//! `quicnprotochat-bindings` — language-agnostic Rust API.
//!
//! This crate provides a `Client` struct whose methods return structured data
//! instead of printing to stdout, making it straightforward to wrap with PyO3
//! (Python) or Magnus (Ruby).
//!
//! Each method opens a fresh QUIC connection, performs the operation, and
//! returns. State is persisted to disk between calls (same QPCE format as the
//! CLI, so state files are fully interoperable).

mod rpc;
mod state;

use std::path::{Path, PathBuf};

use anyhow::Context;
use opaque_ke::{
    ClientLogin, ClientLoginFinishParameters, ClientRegistration,
    ClientRegistrationFinishParameters, CredentialResponse, RegistrationResponse,
};
use quicnprotochat_core::{hybrid_encrypt, opaque_auth::OpaqueSuite, IdentityKeypair};

use rpc::{
    connect_node, current_timestamp_ms, enqueue, fetch_all, fetch_hybrid_key, fetch_key_package,
    fetch_wait, try_hybrid_decrypt, upload_hybrid_key, upload_key_package,
};
use state::{decode_identity_key, load_existing_state, load_or_init_state, save_state, sha256};

// ── Public return types ──────────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub struct WhoamiInfo {
    pub identity_key: String,
    pub fingerprint: String,
    pub hybrid_key: bool,
    pub has_group: bool,
}

#[derive(Debug, Clone)]
pub struct HealthInfo {
    pub status: String,
    pub rtt_ms: u64,
}

#[derive(Debug, Clone)]
pub struct ReceivedMsg {
    pub plaintext: String,
}

// ── Client ───────────────────────────────────────────────────────────────────

/// A configured client handle.
///
/// # Notes
/// - One tokio runtime is created per `Client` instance.
/// - State is loaded from / persisted to `state_path` on every call.
/// - `access_token` must be a valid bearer token accepted by the server.
///   After calling `login()` you should construct a new `Client` with the
///   returned session token as the `access_token`.
/// - Each method opens a new QUIC connection (same behaviour as the CLI).
pub struct Client {
    server: String,
    ca_cert: PathBuf,
    server_name: String,
    state_path: PathBuf,
    state_password: Option<String>,
    access_token: Vec<u8>,
    device_id: Vec<u8>,
    rt: tokio::runtime::Runtime,
}

impl Client {
    pub fn new(
        server: impl Into<String>,
        ca_cert: impl AsRef<Path>,
        server_name: impl Into<String>,
        state_path: impl AsRef<Path>,
        access_token: impl Into<String>,
        state_password: Option<String>,
        device_id: Option<String>,
    ) -> anyhow::Result<Self> {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .context("build tokio runtime")?;
        Ok(Self {
            server: server.into(),
            ca_cert: ca_cert.as_ref().to_path_buf(),
            server_name: server_name.into(),
            state_path: state_path.as_ref().to_path_buf(),
            access_token: access_token.into().into_bytes(),
            device_id: device_id.unwrap_or_default().into_bytes(),
            state_password,
            rt,
        })
    }

    /// Run an async future inside a `LocalSet` (required for `capnp-rpc`'s `!Send` types).
    fn run<F, T>(&self, fut: F) -> anyhow::Result<T>
    where
        F: std::future::Future<Output = anyhow::Result<T>>,
    {
        let local = tokio::task::LocalSet::new();
        self.rt.block_on(local.run_until(fut))
    }

    // ── Identity ─────────────────────────────────────────────────────────────

    /// Return local identity information from the state file (no network call).
    pub fn whoami(&self) -> anyhow::Result<WhoamiInfo> {
        let state = load_existing_state(&self.state_path, self.state_password.as_deref())?;
        let identity = IdentityKeypair::from_seed(state.identity_seed);
        let pk = identity.public_key_bytes();
        let fp = sha256(&pk);
        Ok(WhoamiInfo {
            identity_key: hex::encode(pk),
            fingerprint: hex::encode(fp),
            hybrid_key: state.hybrid_key.is_some(),
            has_group: state.group.is_some(),
        })
    }

    // ── Server connectivity ───────────────────────────────────────────────────

    /// Probe server health. Returns status string and round-trip time in ms.
    pub fn health(&self) -> anyhow::Result<HealthInfo> {
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        self.run(async move {
            let sent_at = current_timestamp_ms();
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            let req = node.health_request();
            let resp = req.send().promise.await.context("health RPC failed")?;
            let status = resp
                .get()
                .context("health: bad response")?
                .get_status()
                .context("health: missing status")?
                .to_str()
                .unwrap_or("unknown")
                .to_string();
            let rtt_ms = current_timestamp_ms().saturating_sub(sent_at);
            Ok(HealthInfo { status, rtt_ms })
        })
    }

    // ── OPAQUE authentication ─────────────────────────────────────────────────

    /// Register a new user account via OPAQUE.
    pub fn register_user(&self, username: &str, password: &str) -> anyhow::Result<()> {
        let username = username.to_string();
        let password = password.to_string();
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        self.run(async move {
            let mut rng = rand::rngs::OsRng;
            let node = connect_node(&server, &ca_cert, &server_name).await?;

            let reg_start = ClientRegistration::<OpaqueSuite>::start(&mut rng, password.as_bytes())
                .map_err(|e| anyhow::anyhow!("OPAQUE register start: {e}"))?;

            let mut req = node.opaque_register_start_request();
            {
                let mut p = req.get();
                p.set_username(&username);
                p.set_request(&reg_start.message.serialize());
            }
            let resp = req
                .send()
                .promise
                .await
                .context("opaque_register_start RPC failed")?;
            let response_bytes = resp
                .get()
                .context("register_start: bad response")?
                .get_response()
                .context("register_start: missing response")?
                .to_vec();

            let reg_response = RegistrationResponse::<OpaqueSuite>::deserialize(&response_bytes)
                .map_err(|e| anyhow::anyhow!("invalid registration response: {e}"))?;

            let reg_finish = reg_start
                .state
                .finish(
                    &mut rng,
                    password.as_bytes(),
                    reg_response,
                    ClientRegistrationFinishParameters::<OpaqueSuite>::default(),
                )
                .map_err(|e| anyhow::anyhow!("OPAQUE register finish: {e}"))?;

            let mut req = node.opaque_register_finish_request();
            {
                let mut p = req.get();
                p.set_username(&username);
                p.set_upload(&reg_finish.message.serialize());
                p.set_identity_key(&[]);
            }
            let resp = req
                .send()
                .promise
                .await
                .context("opaque_register_finish RPC failed")?;
            let success = resp
                .get()
                .context("register_finish: bad response")?
                .get_success();
            anyhow::ensure!(success, "server rejected registration");
            Ok(())
        })
    }

    /// Log in via OPAQUE. Returns the session token as a hex string.
    ///
    /// Construct a new `Client` with this token as `access_token` for
    /// authenticated operations.
    pub fn login(&self, username: &str, password: &str) -> anyhow::Result<String> {
        let username = username.to_string();
        let password = password.to_string();
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        self.run(async move {
            let mut rng = rand::rngs::OsRng;
            let node = connect_node(&server, &ca_cert, &server_name).await?;

            let login_start = ClientLogin::<OpaqueSuite>::start(&mut rng, password.as_bytes())
                .map_err(|e| anyhow::anyhow!("OPAQUE login start: {e}"))?;

            let mut req = node.opaque_login_start_request();
            {
                let mut p = req.get();
                p.set_username(&username);
                p.set_request(&login_start.message.serialize());
            }
            let resp = req
                .send()
                .promise
                .await
                .context("opaque_login_start RPC failed")?;
            let response_bytes = resp
                .get()
                .context("login_start: bad response")?
                .get_response()
                .context("login_start: missing response")?
                .to_vec();

            let credential_response =
                CredentialResponse::<OpaqueSuite>::deserialize(&response_bytes)
                    .map_err(|e| anyhow::anyhow!("invalid credential response: {e}"))?;

            let login_finish = login_start
                .state
                .finish(
                    &mut rng,
                    password.as_bytes(),
                    credential_response,
                    ClientLoginFinishParameters::<OpaqueSuite>::default(),
                )
                .map_err(|e| anyhow::anyhow!("OPAQUE login finish (bad password?): {e}"))?;

            // Derive identity key from local state file.
            let state = load_existing_state(&state_path, state_password.as_deref())
                .context("load state to get identity key for login")?;
            let identity = IdentityKeypair::from_seed(state.identity_seed);
            let identity_key = identity.public_key_bytes().to_vec();

            let mut req = node.opaque_login_finish_request();
            {
                let mut p = req.get();
                p.set_username(&username);
                p.set_finalization(&login_finish.message.serialize());
                p.set_identity_key(&identity_key);
            }
            let resp = req
                .send()
                .promise
                .await
                .context("opaque_login_finish RPC failed")?;
            let token = resp
                .get()
                .context("login_finish: bad response")?
                .get_session_token()
                .context("login_finish: missing session token")?
                .to_vec();
            anyhow::ensure!(!token.is_empty(), "server returned empty session token");
            Ok(hex::encode(token))
        })
    }

    // ── Key management ────────────────────────────────────────────────────────

    /// Upload the stored identity's KeyPackage and hybrid key. Returns the
    /// fingerprint as a hex string.
    pub fn register_state(&self) -> anyhow::Result<String> {
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let state = load_or_init_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            let tls_bytes = member
                .generate_key_package()
                .context("KeyPackage generation failed")?;
            let fingerprint = sha256(&tls_bytes);
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            upload_key_package(
                &node,
                &member.identity().public_key_bytes(),
                &tls_bytes,
                &token,
                &device,
            )
            .await?;
            if let Some(ref hkp) = hybrid_kp {
                upload_hybrid_key(
                    &node,
                    &member.identity().public_key_bytes(),
                    &hkp.public_key(),
                    &token,
                    &device,
                )
                .await?;
            }
            save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            Ok(hex::encode(fingerprint))
        })
    }

    /// Returns `true` if the given peer has uploaded a hybrid public key.
    pub fn check_key(&self, peer_hex: &str) -> anyhow::Result<bool> {
        let peer_key = decode_identity_key(peer_hex)?;
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            let pk = fetch_hybrid_key(&node, &peer_key, &token, &device).await?;
            Ok(pk.is_some())
        })
    }

    // ── MLS group operations ──────────────────────────────────────────────────

    /// Create a new MLS group and persist state.
    pub fn create_group(&self, group_id: &str) -> anyhow::Result<()> {
        let group_id = group_id.to_string();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        self.run(async move {
            let state = load_or_init_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            anyhow::ensure!(member.group_ref().is_none(), "group already exists in state");
            member
                .create_group(group_id.as_bytes())
                .context("create_group failed")?;
            save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            Ok(())
        })
    }

    /// Invite a peer: fetch their KeyPackage, add to group, enqueue Welcome.
    pub fn invite(&self, peer_hex: &str) -> anyhow::Result<()> {
        let peer_key = decode_identity_key(peer_hex)?;
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let state = load_existing_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            let _ = member
                .group_ref()
                .context("no active group; call create_group first")?;

            let node = connect_node(&server, &ca_cert, &server_name).await?;

            let existing_members: Vec<Vec<u8>> = member
                .member_identities()
                .into_iter()
                .filter(|k| k.as_slice() != member.identity().public_key_bytes())
                .collect();

            let peer_kp = fetch_key_package(&node, &peer_key, &token, &device).await?;
            anyhow::ensure!(!peer_kp.is_empty(), "server returned empty KeyPackage for peer");

            let (commit, welcome) = member.add_member(&peer_kp).context("add_member failed")?;

            // Forward commit to existing members.
            for mk in &existing_members {
                if mk.as_slice() == peer_key.as_slice() {
                    continue;
                }
                let peer_hpk = fetch_hybrid_key(&node, mk, &token, &device).await?;
                let payload = if let Some(ref pk) = peer_hpk {
                    hybrid_encrypt(pk, &commit).context("hybrid encrypt commit")?
                } else {
                    commit.clone()
                };
                enqueue(&node, mk, &payload, &token, &device).await?;
            }

            // Enqueue Welcome for new peer (hybrid-encrypted if they support it).
            let peer_hybrid_pk = fetch_hybrid_key(&node, &peer_key, &token, &device).await?;
            let payload = if let Some(ref pk) = peer_hybrid_pk {
                hybrid_encrypt(pk, &welcome).context("hybrid encrypt welcome")?
            } else {
                welcome
            };
            enqueue(&node, &peer_key, &payload, &token, &device).await?;

            save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            Ok(())
        })
    }

    /// Join a group by consuming the first Welcome from the server queue.
    pub fn join(&self) -> anyhow::Result<()> {
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let state = load_existing_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            anyhow::ensure!(member.group_ref().is_none(), "group already active in state");
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            let welcomes =
                fetch_all(&node, &member.identity().public_key_bytes(), &token, &device).await?;
            let raw = welcomes
                .first()
                .cloned()
                .context("no Welcome found in queue for this identity")?;
            let welcome_bytes =
                try_hybrid_decrypt(hybrid_kp.as_ref(), &raw).context("decrypt Welcome failed")?;
            member
                .join_group(&welcome_bytes)
                .context("join_group failed")?;
            save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            Ok(())
        })
    }

    // ── Messaging ─────────────────────────────────────────────────────────────

    /// Send a plaintext message to a peer via the Delivery Service.
    pub fn send_message(&self, peer_hex: &str, text: &str) -> anyhow::Result<()> {
        let peer_key = decode_identity_key(peer_hex)?;
        let text = text.to_string();
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let state = load_existing_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            let ct = member
                .send_message(text.as_bytes())
                .context("send_message failed")?;
            let peer_hybrid_pk = fetch_hybrid_key(&node, &peer_key, &token, &device).await?;
            let payload = if let Some(ref pk) = peer_hybrid_pk {
                hybrid_encrypt(pk, &ct).context("hybrid encrypt failed")?
            } else {
                ct
            };
            enqueue(&node, &peer_key, &payload, &token, &device).await?;
            save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            Ok(())
        })
    }

    /// Poll for incoming messages. `wait_ms = 0` returns immediately;
    /// `wait_ms > 0` long-polls up to that many milliseconds.
    pub fn recv(&self, wait_ms: u64) -> anyhow::Result<Vec<ReceivedMsg>> {
        let server = self.server.clone();
        let ca_cert = self.ca_cert.clone();
        let server_name = self.server_name.clone();
        let state_path = self.state_path.clone();
        let state_password = self.state_password.clone();
        let token = self.access_token.clone();
        let device = self.device_id.clone();
        self.run(async move {
            let state = load_existing_state(&state_path, state_password.as_deref())?;
            let (mut member, hybrid_kp) = state.into_parts(&state_path)?;
            let node = connect_node(&server, &ca_cert, &server_name).await?;
            let payloads = fetch_wait(
                &node,
                &member.identity().public_key_bytes(),
                wait_ms,
                &token,
                &device,
            )
            .await?;

            let mut msgs = Vec::new();
            for payload in &payloads {
                let mls_payload = match try_hybrid_decrypt(hybrid_kp.as_ref(), payload) {
                    Ok(b) => b,
                    Err(_) => continue,
                };
                match member.receive_message(&mls_payload) {
                    Ok(Some(pt)) => msgs.push(ReceivedMsg {
                        plaintext: String::from_utf8_lossy(&pt).into_owned(),
                    }),
                    Ok(None) => {} // MLS commit — no application message
                    Err(_) => {}
                }
            }
            if !payloads.is_empty() {
                save_state(&state_path, &member, hybrid_kp.as_ref(), state_password.as_deref())?;
            }
            Ok(msgs)
        })
    }
}