quicproquo/crates/quicprochat-rpc/src/client.rs

//! QUIC RPC client — connect to server, send requests, receive push events.

use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;

use bytes::{Bytes, BytesMut};
use quinn::{Connection, Endpoint};
use tokio::sync::mpsc;
use tracing::{debug, warn};

use crate::auth_handshake;
use crate::error::{RpcError, RpcStatus};
use crate::framing::{PushFrame, RequestFrame, ResponseFrame};

/// Configuration for the RPC client.
pub struct RpcClientConfig {
    /// Server address to connect to.
    pub server_addr: std::net::SocketAddr,
    /// Server name for TLS verification.
    pub server_name: String,
    /// TLS client config (rustls).
    pub tls_config: Arc<rustls::ClientConfig>,
    /// ALPN protocol.
    pub alpn: Vec<u8>,
    /// Session token to send during auth handshake.
    pub session_token: Option<Vec<u8>>,
}

/// A QUIC RPC client connection.
pub struct RpcClient {
    connection: Connection,
    next_request_id: AtomicU32,
}

impl RpcClient {
    /// Connect to the RPC server.
    pub async fn connect(config: RpcClientConfig) -> Result<Self, RpcError> {
        let mut tls = (*config.tls_config).clone();
        tls.alpn_protocols = vec![config.alpn];
        let quic_tls = quinn::crypto::rustls::QuicClientConfig::try_from(tls)
            .map_err(|e| RpcError::Connection(format!("TLS config: {e}")))?;

        let bind_addr = std::net::SocketAddr::new(
            std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
            0,
        );
        let mut endpoint = Endpoint::client(bind_addr)
            .map_err(|e| RpcError::Connection(e.to_string()))?;
        endpoint.set_default_client_config(quinn::ClientConfig::new(Arc::new(quic_tls)));

        let connection = endpoint
            .connect(config.server_addr, &config.server_name)
            .map_err(|e| RpcError::Connection(e.to_string()))?
            .await
            .map_err(|e| RpcError::Connection(e.to_string()))?;

        debug!(remote = %connection.remote_address(), "connected to RPC server");

        // Perform auth handshake if a session token was provided.
        if let Some(ref token) = config.session_token {
            let (mut send, mut recv) = connection
                .open_bi()
                .await
                .map_err(|e| RpcError::Connection(format!("open auth stream: {e}")))?;

            auth_handshake::send_auth_init(&mut send, token).await?;
            send.finish()
                .map_err(|e| RpcError::Connection(format!("finish auth send: {e}")))?;
            auth_handshake::recv_auth_ack(&mut recv).await?;
            debug!("auth handshake complete");
        }

        Ok(Self {
            connection,
            next_request_id: AtomicU32::new(1),
        })
    }

    /// Send an RPC request and wait for the response.
    pub async fn call(
        &self,
        method_id: u16,
        payload: Bytes,
    ) -> Result<Bytes, RpcError> {
        let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);

        let (mut send, mut recv) = self
            .connection
            .open_bi()
            .await
            .map_err(|e| RpcError::Connection(e.to_string()))?;

        // Send request.
        let frame = RequestFrame {
            method_id,
            request_id,
            payload,
        };
        let encoded = frame.encode();
        send.write_all(&encoded)
            .await
            .map_err(|e| RpcError::Connection(e.to_string()))?;
        send.finish().map_err(|e| RpcError::Connection(e.to_string()))?;

        // Read response.
        let mut buf = BytesMut::new();
        while let Some(chunk) = recv
            .read_chunk(65536, true)
            .await
            .map_err(|e| RpcError::Connection(e.to_string()))?
        {
            buf.extend_from_slice(&chunk.bytes);
            if buf.len() > crate::framing::MAX_PAYLOAD_SIZE + crate::framing::RESPONSE_HEADER_SIZE {
                return Err(RpcError::PayloadTooLarge {
                    size: buf.len(),
                    max: crate::framing::MAX_PAYLOAD_SIZE,
                });
            }
        }

        let response = ResponseFrame::decode(&mut buf)?
            .ok_or_else(|| RpcError::Decode("incomplete response frame".into()))?;

        if response.request_id != request_id {
            return Err(RpcError::Decode(format!(
                "request_id mismatch: sent {request_id}, got {}",
                response.request_id
            )));
        }

        match RpcStatus::from_u8(response.status) {
            Some(RpcStatus::Ok) => Ok(response.payload),
            Some(status) => Err(RpcError::Server {
                status,
                message: String::from_utf8_lossy(&response.payload).into_owned(),
            }),
            None => Err(RpcError::Decode(format!(
                "unknown status byte: {}",
                response.status
            ))),
        }
    }

    /// Subscribe to server-push events. Returns a receiver channel.
    /// Spawns a background task that reads uni-streams.
    pub fn subscribe_push(&self) -> mpsc::UnboundedReceiver<PushFrame> {
        let (tx, rx) = mpsc::unbounded_channel();
        let conn = self.connection.clone();

        tokio::spawn(async move {
            loop {
                match conn.accept_uni().await {
                    Ok(mut recv) => {
                        let mut buf = BytesMut::new();
                        loop {
                            match recv.read_chunk(65536, true).await {
                                Ok(Some(chunk)) => buf.extend_from_slice(&chunk.bytes),
                                Ok(None) => break,
                                Err(e) => {
                                    debug!("push stream read error: {e}");
                                    break;
                                }
                            }
                        }
                        match PushFrame::decode(&mut buf) {
                            Ok(Some(frame)) => {
                                if tx.send(frame).is_err() {
                                    return; // receiver dropped
                                }
                            }
                            Ok(None) => debug!("incomplete push frame"),
                            Err(e) => debug!("push decode error: {e}"),
                        }
                    }
                    Err(quinn::ConnectionError::ApplicationClosed(_)) => break,
                    Err(e) => {
                        warn!("accept_uni error: {e}");
                        break;
                    }
                }
            }
        });

        rx
    }

    /// Close the connection gracefully.
    pub fn close(&self) {
        self.connection.close(0u32.into(), b"bye");
    }

    /// Get the underlying QUIC connection (for advanced use).
    pub fn connection(&self) -> &Connection {
        &self.connection
    }
}