tsoracle-standalone 0.1.1

//
//  ░▀█▀░█▀▀░█▀█░█▀▄░█▀█░█▀▀░█░░░█▀▀
//  ░░█░░▀▀█░█░█░█▀▄░█▀█░█░░░█░░░█▀▀
//  ░░▀░░▀▀▀░▀▀▀░▀░▀░▀░▀░▀▀▀░▀▀▀░▀▀▀
//
//  tsoracle — Distributed Timestamp Oracle
//  https://www.tsoracle.rs
//
//  Copyright (c) 2026 Prisma Risk
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

//! tonic-based peer transport for openraft 0.10.
//!
//! Uses `RaftNetworkV2` (the modern API in 0.10). `RaftNetwork` (V1) was
//! removed in 0.10.0-alpha.20 and redirects to `openraft-legacy`.
//!
//! Wire format:
//!   - `AppendEntries` / `Vote`: a `RaftMessage { bytes payload }` where
//!     `payload` is a postcard-encoded openraft request or response.
//!   - `Snapshot`: a *client-streaming* RPC of `SnapshotChunk` messages.
//!     One `header` chunk (postcard-encoded vote + meta), then `SNAPSHOT_CHUNK_SIZE`
//!     byte `data` chunks until end-of-stream. The receiver reassembles the
//!     data buffer and calls `Raft::install_full_snapshot`.
//!
//! The chunked snapshot path is what makes this transport safe to use with a
//! state machine that grows past the default gRPC unary frame limit (4 MiB).
//! See `proto/raft_peer.proto` for the framing rules.

use std::collections::HashMap;
use std::future::Future;
use std::io::Cursor;
use std::sync::Arc;
use std::time::Duration;

use futures::StreamExt;
use openraft::error::{NetworkError, RPCError, StreamingError, Unreachable};
use openraft::errors::ReplicationClosed;
use openraft::network::{RPCOption, RaftNetworkFactory, RaftNetworkV2};
use openraft::raft::{
    AppendEntriesRequest, AppendEntriesResponse, SnapshotResponse, TransferLeaderRequest,
    VoteRequest, VoteResponse,
};
use openraft::type_config::alias::{SnapshotOf, VoteOf};
use tokio::sync::Mutex;
use tonic::transport::{Channel, ClientTlsConfig};

use tsoracle_driver_openraft::{OpenraftPeer as Node, TypeConfig};
type NodeId = u64;

pub mod proto {
    tonic::include_proto!("tsoracle.raft.peer.v1");
}

use proto::RaftMessage;
use proto::SnapshotChunk;
use proto::SnapshotHeader;
use proto::raft_peer_service_client::RaftPeerServiceClient;
use proto::raft_peer_service_server::{RaftPeerService, RaftPeerServiceServer};
use proto::snapshot_chunk::Kind as ChunkKind;

/// Snapshot data is shipped in chunks of this many bytes. Sized to fit
/// comfortably inside the default gRPC max-frame limit (4 MiB) with room
/// for proto overhead and to keep per-RPC memory bounded on both sides.
/// The header chunk is sent separately and is small (a Vote + SnapshotMeta).
pub const SNAPSHOT_CHUNK_SIZE: usize = 1024 * 1024;

/// Upper bound on the *total* bytes the snapshot handler will reassemble from a
/// single client-streaming RPC. The handler refuses (with `ResourceExhausted`)
/// any stream whose cumulative `data` chunks would cross this line, so a peer
/// that can reach the raft port cannot drive the receiver to OOM by sending an
/// endless run of chunks. Sized generously for the small high-water state
/// machine; real deployments should size this against the largest realistic
/// state-machine snapshot.
pub const MAX_SNAPSHOT_BYTES: usize = 64 * 1024 * 1024;

/// Per-message decode/encode cap applied to the peer server. It must stay
/// strictly above `SNAPSHOT_CHUNK_SIZE`: every snapshot `Data` chunk is one
/// decoded protobuf message of up to that size, so a smaller cap would reject
/// legitimate chunks. Deriving it from the chunk size (plus headroom for the
/// proto and postcard framing) keeps that invariant impossible to break by
/// accident. `append_entries`/`vote` messages are far smaller, so this bounds
/// them too.
pub const MAX_PEER_MESSAGE_BYTES: usize = SNAPSHOT_CHUNK_SIZE + 256 * 1024;

/// Wall-clock ceiling on a single snapshot-install stream. Bounds a slow-loris
/// peer that opens a stream and dribbles bytes to keep the reassembly buffer
/// alive; sized against `MAX_SNAPSHOT_BYTES` arriving over a slow-but-legitimate
/// link.
const SNAPSHOT_STREAM_TIMEOUT: Duration = Duration::from_secs(60);

// ---------------------------------------------------------------------------
// Pool type + eviction helper.
// ---------------------------------------------------------------------------

type Pool = Arc<Mutex<HashMap<(NodeId, String), RaftPeerServiceClient<Channel>>>>;

// Generic over the value type so the keying/eviction logic is unit-testable
// without constructing a live RaftPeerServiceClient.
async fn evict<V>(pool: &Arc<Mutex<HashMap<(NodeId, String), V>>>, target: NodeId, addr: &str) {
    pool.lock().await.remove(&(target, addr.to_string()));
}

/// Drive a unary peer RPC under the caller's `RPCOption` hard-TTL deadline,
/// returning the decoded response body.
///
/// Why this exists: openraft only enforces `RPCOption` itself for some call
/// sites. `vote` and `transfer_leader` are wrapped in openraft's own
/// `C::timeout` (see `raft_core::broadcast_*`), but the replication
/// `append_entries` path is not — `stream_append_sequential` simply awaits
/// `network.append_entries(req, option)` and relies on the *transport* to honor
/// `option.hard_ttl()`. A transport that ignored the option left append on a
/// silently black-holed connection (no RST — NAT/firewall drop) wedged until TCP
/// keepalive eventually tripped (~2h by default), stalling replication to that
/// follower. Applying the deadline here closes that gap and keeps all three
/// unary RPCs uniformly bounded; for `vote`/`transfer_leader` it is harmless
/// belt-and-suspenders that simply fires no earlier than openraft's own timeout.
///
/// On any failure the pooled client for `(target, addr)` is evicted so the next
/// attempt reconnects fresh. A deadline elapse is surfaced as `Unreachable`
/// (openraft backs off before retrying) rather than `Network` (retry at once),
/// which is the right posture for a connection we just gave up on. Generic over
/// the pool value and response types so it is unit-testable without a live
/// `RaftPeerServiceClient`, mirroring [`evict`].
async fn unary_call<ClientHandle, Body>(
    pool: &Arc<Mutex<HashMap<(NodeId, String), ClientHandle>>>,
    target: NodeId,
    addr: &str,
    deadline: Duration,
    call: impl Future<Output = Result<tonic::Response<Body>, tonic::Status>>,
) -> Result<Body, RPCError<TypeConfig>> {
    match tokio::time::timeout(deadline, call).await {
        Ok(Ok(resp)) => Ok(resp.into_inner()),
        Ok(Err(status)) => {
            evict(pool, target, addr).await;
            Err(RPCError::Network(NetworkError::new(&status)))
        }
        Err(_elapsed) => {
            evict(pool, target, addr).await;
            let timed_out = std::io::Error::new(
                std::io::ErrorKind::TimedOut,
                format!("peer RPC exceeded {deadline:?} deadline"),
            );
            Err(RPCError::Unreachable(Unreachable::new(&timed_out)))
        }
    }
}

// ---------------------------------------------------------------------------
// PeerFactory — constructs PeerNetwork instances for each target node.
// ---------------------------------------------------------------------------

pub struct PeerFactory {
    pool: Pool,
    tls: Option<ClientTlsConfig>,
}

impl PeerFactory {
    pub fn new(tls: Option<ClientTlsConfig>) -> Self {
        Self {
            pool: Arc::new(Mutex::new(HashMap::new())),
            tls,
        }
    }
}

impl RaftNetworkFactory<TypeConfig> for PeerFactory {
    type Network = PeerNetwork;

    async fn new_client(&mut self, target: NodeId, node: &Node) -> Self::Network {
        PeerNetwork {
            target,
            addr: node.addr.clone(),
            pool: self.pool.clone(),
            tls: self.tls.clone(),
        }
    }
}

// ---------------------------------------------------------------------------
// PeerNetwork — implements RaftNetworkV2<TypeConfig> for one target peer.
// ---------------------------------------------------------------------------

pub struct PeerNetwork {
    target: NodeId,
    addr: String,
    pool: Pool,
    tls: Option<ClientTlsConfig>,
}

impl PeerNetwork {
    /// Return a cached (or freshly connected) tonic client to the target.
    async fn client(&self) -> Result<RaftPeerServiceClient<Channel>, RPCError<TypeConfig>> {
        let key = (self.target, self.addr.clone());
        {
            let pool = self.pool.lock().await;
            if let Some(client) = pool.get(&key) {
                return Ok(client.clone());
            }
        }
        let channel = match &self.tls {
            Some(tls) => Channel::from_shared(format!("https://{}", self.addr))
                .map_err(|err| RPCError::Unreachable(Unreachable::new(&err)))?
                .tls_config(tls.clone())
                .map_err(|err| RPCError::Unreachable(Unreachable::new(&err)))?
                .connect()
                .await
                .map_err(|err| RPCError::Unreachable(Unreachable::new(&err)))?,
            None => Channel::from_shared(format!("http://{}", self.addr))
                .map_err(|err| RPCError::Unreachable(Unreachable::new(&err)))?
                .connect()
                .await
                .map_err(|err| RPCError::Unreachable(Unreachable::new(&err)))?,
        };
        let client = RaftPeerServiceClient::new(channel);
        self.pool.lock().await.insert(key, client.clone());
        Ok(client)
    }
}

impl RaftNetworkV2<TypeConfig> for PeerNetwork {
    async fn append_entries(
        &mut self,
        rpc: AppendEntriesRequest<TypeConfig>,
        option: RPCOption,
    ) -> Result<AppendEntriesResponse<TypeConfig>, RPCError<TypeConfig>> {
        let mut c = self.client().await?;
        let payload =
            postcard::to_stdvec(&rpc).map_err(|err| RPCError::Network(NetworkError::new(&err)))?;
        let reply = unary_call(
            &self.pool,
            self.target,
            &self.addr,
            option.hard_ttl(),
            c.append_entries(RaftMessage { payload }),
        )
        .await?;
        let body: AppendEntriesResponse<TypeConfig> = postcard::from_bytes(&reply.payload)
            .map_err(|err| RPCError::Network(NetworkError::new(&err)))?;
        Ok(body)
    }

    /// Forward a leadership-transfer request to the target peer. openraft calls
    /// this on the outgoing leader when `trigger().transfer_leader` fires; the
    /// receiver hands the request to `Raft::handle_transfer_leader`. Without
    /// this override the default no-op drops the request and leadership only
    /// moves on the next election timeout.
    async fn transfer_leader(
        &mut self,
        req: TransferLeaderRequest<TypeConfig>,
        option: RPCOption,
    ) -> Result<(), RPCError<TypeConfig>> {
        let mut c = self.client().await?;
        let payload =
            postcard::to_stdvec(&req).map_err(|err| RPCError::Network(NetworkError::new(&err)))?;
        // The reply payload is empty by contract; we only care that the RPC
        // landed within the deadline.
        unary_call(
            &self.pool,
            self.target,
            &self.addr,
            option.hard_ttl(),
            c.transfer_leader(RaftMessage { payload }),
        )
        .await?;
        Ok(())
    }

    async fn vote(
        &mut self,
        rpc: VoteRequest<TypeConfig>,
        option: RPCOption,
    ) -> Result<VoteResponse<TypeConfig>, RPCError<TypeConfig>> {
        let mut c = self.client().await?;
        let payload =
            postcard::to_stdvec(&rpc).map_err(|err| RPCError::Network(NetworkError::new(&err)))?;
        let reply = unary_call(
            &self.pool,
            self.target,
            &self.addr,
            option.hard_ttl(),
            c.vote(RaftMessage { payload }),
        )
        .await?;
        let body: VoteResponse<TypeConfig> = postcard::from_bytes(&reply.payload)
            .map_err(|err| RPCError::Network(NetworkError::new(&err)))?;
        Ok(body)
    }

    /// Send a snapshot to the target as a stream of `SnapshotChunk`s.
    ///
    /// Stream layout: one `header` chunk (postcard vote + postcard meta),
    /// followed by `ceil(data.len() / SNAPSHOT_CHUNK_SIZE)` `data` chunks.
    /// An empty data buffer is permitted and results in zero `data` chunks.
    ///
    /// Cancellation: if `cancel` resolves before the server responds, the
    /// in-flight stream is dropped (closing the client side of the stream)
    /// and we return `StreamingError::Closed`.
    async fn full_snapshot(
        &mut self,
        vote: VoteOf<TypeConfig>,
        snapshot: SnapshotOf<TypeConfig>,
        cancel: impl Future<Output = ReplicationClosed> + openraft::OptionalSend + 'static,
        _option: RPCOption,
    ) -> Result<SnapshotResponse<TypeConfig>, StreamingError<TypeConfig>> {
        // Pre-encode the header fields. Both are small (Vote + SnapshotMeta).
        let vote_bytes = postcard::to_stdvec(&vote)
            .map_err(|e| StreamingError::Network(NetworkError::new(&e)))?;
        let meta_bytes = postcard::to_stdvec(&snapshot.meta)
            .map_err(|e| StreamingError::Network(NetworkError::new(&e)))?;
        let data_bytes = snapshot.snapshot.into_inner();

        // Build the chunk stream lazily. We materialize chunks into owned
        // Vecs (a copy per chunk) because prost generates `bytes` fields as
        // `Vec<u8>` by default. The original `data_bytes` is freed when the
        // iterator is exhausted; total transient memory is ≈ data_bytes plus
        // one in-flight chunk.
        let header_chunk = SnapshotChunk {
            kind: Some(ChunkKind::Header(SnapshotHeader {
                vote: vote_bytes,
                meta: meta_bytes,
            })),
        };

        let data_chunks = data_bytes
            .chunks(SNAPSHOT_CHUNK_SIZE)
            .map(|c| SnapshotChunk {
                kind: Some(ChunkKind::Data(c.to_vec())),
            })
            .collect::<Vec<_>>();

        let outbound =
            futures::stream::iter(std::iter::once(header_chunk).chain(data_chunks.into_iter()));

        let mut c = self.client().await.map_err(|e| match e {
            RPCError::Network(n) => StreamingError::Network(n),
            RPCError::Unreachable(u) => StreamingError::Unreachable(u),
            RPCError::Timeout(t) => StreamingError::Timeout(t),
            // RPCError::RemoteError is Infallible in the default type param
        })?;

        // Drive the streaming RPC concurrently with the cancellation future.
        // If `cancel` fires first, dropping the in-flight future closes the
        // outbound stream and aborts on the server side.
        tokio::select! {
            result = c.snapshot(outbound) => {
                let raw = match result {
                    Ok(resp) => resp,
                    Err(err) => {
                        evict(&self.pool, self.target, &self.addr).await;
                        return Err(StreamingError::Network(NetworkError::new(&err)));
                    }
                };
                let inner = raw.into_inner();
                let resp: SnapshotResponse<TypeConfig> =
                    postcard::from_bytes(&inner.payload)
                        .map_err(|err| StreamingError::Network(NetworkError::new(&err)))?;
                Ok(resp)
            }
            closed = cancel => {
                Err(StreamingError::Closed(closed))
            }
        }
    }
}

// ---------------------------------------------------------------------------
// PeerServiceImpl — tonic server, demuxes RaftMessage → Raft API calls.
// ---------------------------------------------------------------------------

pub struct PeerServiceImpl<SM = ()> {
    pub raft: openraft::Raft<TypeConfig, SM>,
}

#[tonic::async_trait]
impl<SM: Send + Sync + 'static> RaftPeerService for PeerServiceImpl<SM> {
    async fn append_entries(
        &self,
        request: tonic::Request<RaftMessage>,
    ) -> Result<tonic::Response<RaftMessage>, tonic::Status> {
        let body: AppendEntriesRequest<TypeConfig> =
            postcard::from_bytes(&request.into_inner().payload)
                .map_err(|e| tonic::Status::invalid_argument(e.to_string()))?;
        let resp = self
            .raft
            .append_entries(body)
            .await
            .map_err(|e| tonic::Status::internal(e.to_string()))?;
        let payload =
            postcard::to_stdvec(&resp).map_err(|e| tonic::Status::internal(e.to_string()))?;
        Ok(tonic::Response::new(RaftMessage { payload }))
    }

    async fn vote(
        &self,
        request: tonic::Request<RaftMessage>,
    ) -> Result<tonic::Response<RaftMessage>, tonic::Status> {
        let body: VoteRequest<TypeConfig> = postcard::from_bytes(&request.into_inner().payload)
            .map_err(|e| tonic::Status::invalid_argument(e.to_string()))?;
        let resp = self
            .raft
            .vote(body)
            .await
            .map_err(|e| tonic::Status::internal(e.to_string()))?;
        let payload =
            postcard::to_stdvec(&resp).map_err(|e| tonic::Status::internal(e.to_string()))?;
        Ok(tonic::Response::new(RaftMessage { payload }))
    }

    async fn transfer_leader(
        &self,
        request: tonic::Request<RaftMessage>,
    ) -> Result<tonic::Response<RaftMessage>, tonic::Status> {
        let body: TransferLeaderRequest<TypeConfig> =
            postcard::from_bytes(&request.into_inner().payload)
                .map_err(|e| tonic::Status::invalid_argument(e.to_string()))?;
        self.raft
            .handle_transfer_leader(body)
            .await
            .map_err(|e| tonic::Status::internal(e.to_string()))?;
        Ok(tonic::Response::new(RaftMessage {
            payload: Vec::new(),
        }))
    }

    /// Reassemble a streamed snapshot and hand it to `install_full_snapshot`.
    ///
    /// Framing and the byte/time bounds live in [`reassemble_snapshot`]; this
    /// handler just enforces the per-stream wall-clock limit and forwards the
    /// result to the local `Raft`.
    async fn snapshot(
        &self,
        request: tonic::Request<tonic::Streaming<SnapshotChunk>>,
    ) -> Result<tonic::Response<RaftMessage>, tonic::Status> {
        let assembled = tokio::time::timeout(
            SNAPSHOT_STREAM_TIMEOUT,
            reassemble_snapshot(request.into_inner(), MAX_SNAPSHOT_BYTES),
        )
        .await
        .map_err(|_| tonic::Status::deadline_exceeded("snapshot stream timed out"))??;

        let snapshot = openraft::storage::Snapshot {
            meta: assembled.meta,
            snapshot: Cursor::new(assembled.data),
        };

        let resp = self
            .raft
            .install_full_snapshot(assembled.vote, snapshot)
            .await
            .map_err(|e| tonic::Status::internal(e.to_string()))?;

        let payload =
            postcard::to_stdvec(&resp).map_err(|e| tonic::Status::internal(e.to_string()))?;
        Ok(tonic::Response::new(RaftMessage { payload }))
    }
}

/// A snapshot reassembled from a peer's client-streaming RPC.
#[derive(Debug)]
struct AssembledSnapshot {
    vote: VoteOf<TypeConfig>,
    meta: openraft::type_config::alias::SnapshotMetaOf<TypeConfig>,
    data: Vec<u8>,
}

/// Parse the leading header chunk, then concatenate trailing `data` chunks into
/// a single buffer, bounding the total at `max_bytes`.
///
/// Framing contract (mirrors `proto/raft_peer.proto`):
///   - exactly one `header` chunk at the start;
///   - zero or more `data` chunks afterwards;
///   - any other ordering is rejected as `InvalidArgument`.
///
/// The running total is checked *before* each chunk is buffered, so a stream
/// that would cross `max_bytes` is refused with `ResourceExhausted` without the
/// buffer ever exceeding the limit — this is what keeps a reachable peer from
/// driving the receiver to OOM by sending an unbounded run of chunks. Generic
/// over the stream so the bound is unit-testable: `tonic::Streaming` is not
/// constructible in a test, but a `futures::stream::iter` is.
async fn reassemble_snapshot<S>(
    mut stream: S,
    max_bytes: usize,
) -> Result<AssembledSnapshot, tonic::Status>
where
    S: futures::Stream<Item = Result<SnapshotChunk, tonic::Status>> + Unpin,
{
    // The first chunk must be the header.
    let first = stream
        .next()
        .await
        .ok_or_else(|| tonic::Status::invalid_argument("snapshot stream ended before header"))?
        .map_err(|e| tonic::Status::internal(format!("snapshot stream error: {e}")))?;
    let header = match first.kind {
        Some(ChunkKind::Header(h)) => h,
        Some(ChunkKind::Data(_)) => {
            return Err(tonic::Status::invalid_argument(
                "first snapshot chunk must be a header",
            ));
        }
        None => {
            return Err(tonic::Status::invalid_argument(
                "snapshot chunk missing kind",
            ));
        }
    };

    let vote: VoteOf<TypeConfig> = postcard::from_bytes(&header.vote)
        .map_err(|e| tonic::Status::invalid_argument(format!("bad vote: {e}")))?;
    let meta: openraft::type_config::alias::SnapshotMetaOf<TypeConfig> =
        postcard::from_bytes(&header.meta)
            .map_err(|e| tonic::Status::invalid_argument(format!("bad meta: {e}")))?;

    // Reassemble subsequent data chunks, refusing the stream the moment its
    // cumulative size would exceed `max_bytes`.
    let mut data: Vec<u8> = Vec::new();
    while let Some(chunk) = stream.next().await {
        let chunk =
            chunk.map_err(|e| tonic::Status::internal(format!("snapshot stream error: {e}")))?;
        match chunk.kind {
            Some(ChunkKind::Data(bytes)) => {
                if data.len() + bytes.len() > max_bytes {
                    return Err(tonic::Status::resource_exhausted(format!(
                        "snapshot exceeds {max_bytes}-byte reassembly limit"
                    )));
                }
                data.extend_from_slice(&bytes);
            }
            Some(ChunkKind::Header(_)) => {
                return Err(tonic::Status::invalid_argument(
                    "unexpected header chunk after first",
                ));
            }
            None => {
                return Err(tonic::Status::invalid_argument(
                    "snapshot chunk missing kind",
                ));
            }
        }
    }

    Ok(AssembledSnapshot { vote, meta, data })
}

/// Construct the tonic server-side handler for the RaftPeerService.
pub fn server<SM: Send + Sync + 'static>(
    raft: openraft::Raft<TypeConfig, SM>,
) -> RaftPeerServiceServer<PeerServiceImpl<SM>> {
    RaftPeerServiceServer::new(PeerServiceImpl { raft })
}

#[cfg(test)]
mod tls_tests {
    use super::*;
    use crate::config::PeerTlsConfig;
    use crate::peer_tls::build_peer_tls;
    use std::sync::Arc;
    use tonic::transport::{Certificate, ClientTlsConfig, Identity};

    // --- cert helpers (rcgen 0.13) ---
    struct Certs {
        ca_pem: String,
        node_cert: String,
        node_key: String,
        other_leaf_cert: String,
        other_leaf_key: String,
    }

    fn mint() -> Certs {
        use rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair};
        let mk_ca = |name: &str| {
            let key = KeyPair::generate().unwrap();
            let mut p = CertificateParams::new(vec![name.to_string()]).unwrap();
            p.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
            let cert = p.self_signed(&key).unwrap();
            (cert, key)
        };
        let (ca, ca_key) = mk_ca("tso-ca");
        let leaf_key = KeyPair::generate().unwrap();
        let leaf_params =
            CertificateParams::new(vec!["localhost".to_string(), "127.0.0.1".to_string()]).unwrap();
        let leaf = leaf_params.signed_by(&leaf_key, &ca, &ca_key).unwrap();
        let (other_ca, other_ca_key) = mk_ca("other-ca");
        let other_key = KeyPair::generate().unwrap();
        let other_params = CertificateParams::new(vec!["127.0.0.1".to_string()]).unwrap();
        let other_leaf = other_params
            .signed_by(&other_key, &other_ca, &other_ca_key)
            .unwrap();
        Certs {
            ca_pem: ca.pem(),
            node_cert: leaf.pem(),
            node_key: leaf_key.serialize_pem(),
            other_leaf_cert: other_leaf.pem(),
            other_leaf_key: other_key.serialize_pem(),
        }
    }

    fn node_material(c: &Certs, dir: &std::path::Path) -> crate::peer_tls::PeerTlsMaterial {
        let cert = dir.join("n.crt");
        let key = dir.join("n.key");
        let ca = dir.join("ca.crt");
        std::fs::write(&cert, &c.node_cert).unwrap();
        std::fs::write(&key, &c.node_key).unwrap();
        std::fs::write(&ca, &c.ca_pem).unwrap();
        build_peer_tls(&PeerTlsConfig { cert, key, ca }).unwrap()
    }

    // Minimal stub server (handlers never called — only the TLS handshake is).
    #[derive(Clone)]
    struct Stub;

    #[tonic::async_trait]
    impl proto::raft_peer_service_server::RaftPeerService for Stub {
        async fn append_entries(
            &self,
            _: tonic::Request<proto::RaftMessage>,
        ) -> Result<tonic::Response<proto::RaftMessage>, tonic::Status> {
            Err(tonic::Status::unimplemented("stub"))
        }
        async fn vote(
            &self,
            _: tonic::Request<proto::RaftMessage>,
        ) -> Result<tonic::Response<proto::RaftMessage>, tonic::Status> {
            Err(tonic::Status::unimplemented("stub"))
        }
        async fn snapshot(
            &self,
            _: tonic::Request<tonic::Streaming<proto::SnapshotChunk>>,
        ) -> Result<tonic::Response<proto::RaftMessage>, tonic::Status> {
            Err(tonic::Status::unimplemented("stub"))
        }
        async fn transfer_leader(
            &self,
            _: tonic::Request<proto::RaftMessage>,
        ) -> Result<tonic::Response<proto::RaftMessage>, tonic::Status> {
            Err(tonic::Status::unimplemented("stub"))
        }
    }

    async fn spawn_stub(server_tls: tonic::transport::ServerTlsConfig) -> std::net::SocketAddr {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            tonic::transport::Server::builder()
                .tls_config(server_tls)
                .unwrap()
                .add_service(proto::raft_peer_service_server::RaftPeerServiceServer::new(
                    Stub,
                ))
                .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
                .await
                .ok();
        });
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        addr
    }

    fn make_net(addr: std::net::SocketAddr, tls: Option<ClientTlsConfig>) -> PeerNetwork {
        PeerNetwork {
            target: 2,
            addr: addr.to_string(),
            pool: Arc::new(Mutex::new(HashMap::new())),
            tls,
        }
    }

    // `Channel::connect()` is lazy — the TLS handshake happens on the first RPC.
    // We therefore attempt a real RPC (append_entries) and check whether it fails
    // at the transport level (tonic Status) vs. at the stub handler (Unimplemented).
    // A successful mTLS handshake produces Unimplemented; a rejected one produces
    // a connection-level error (Unavailable / Unknown).
    async fn probe(net: PeerNetwork) -> tonic::Code {
        match net.client().await {
            Err(_) => tonic::Code::Unavailable,
            Ok(mut c) => {
                match c
                    .append_entries(tonic::Request::new(proto::RaftMessage {
                        payload: Vec::new(),
                    }))
                    .await
                {
                    Ok(_) => tonic::Code::Ok,
                    Err(s) => s.code(),
                }
            }
        }
    }

    #[tokio::test]
    async fn valid_node_cert_connects() {
        let dir = tempfile::tempdir().unwrap();
        let c = mint();
        let m = node_material(&c, dir.path());
        let addr = spawn_stub(m.server.clone()).await;
        // Stub returns Unimplemented — handshake succeeded.
        assert_eq!(
            probe(make_net(addr, Some(m.client.clone()))).await,
            tonic::Code::Unimplemented
        );
    }

    #[tokio::test]
    async fn no_client_cert_rejected() {
        let dir = tempfile::tempdir().unwrap();
        let c = mint();
        let m = node_material(&c, dir.path());
        let addr = spawn_stub(m.server.clone()).await;
        let no_id = ClientTlsConfig::new()
            .ca_certificate(Certificate::from_pem(&c.ca_pem))
            .domain_name("localhost");
        let code = probe(make_net(addr, Some(no_id))).await;
        assert_ne!(
            code,
            tonic::Code::Unimplemented,
            "server must reject a client with no cert"
        );
    }

    #[tokio::test]
    async fn wrong_ca_client_rejected() {
        let dir = tempfile::tempdir().unwrap();
        let c = mint();
        let m = node_material(&c, dir.path());
        let addr = spawn_stub(m.server.clone()).await;
        let wrong = ClientTlsConfig::new()
            .ca_certificate(Certificate::from_pem(&c.ca_pem))
            .identity(Identity::from_pem(&c.other_leaf_cert, &c.other_leaf_key))
            .domain_name("localhost");
        let code = probe(make_net(addr, Some(wrong))).await;
        assert_ne!(
            code,
            tonic::Code::Unimplemented,
            "server must reject a cert from a foreign CA"
        );
    }

    #[tokio::test]
    async fn plaintext_against_tls_fails() {
        let dir = tempfile::tempdir().unwrap();
        let c = mint();
        let m = node_material(&c, dir.path());
        let addr = spawn_stub(m.server.clone()).await;
        let code = probe(make_net(addr, None)).await;
        assert_ne!(
            code,
            tonic::Code::Unimplemented,
            "plaintext must not reach a TLS-only server"
        );
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::sync::Arc;

    use openraft::Vote;
    use openraft::type_config::alias::SnapshotMetaOf;
    use tokio::sync::Mutex;

    use super::*;

    #[tokio::test]
    async fn pool_key_distinguishes_addr_changes() {
        // Same NodeId, different addr, must not collide in the pool.
        let mut map: HashMap<(u64, String), u8> = HashMap::new();
        map.insert((1, "old:1".to_string()), 0);
        assert!(!map.contains_key(&(1, "new:1".to_string())));
        assert!(map.contains_key(&(1, "old:1".to_string())));
    }

    #[tokio::test]
    async fn evict_removes_only_the_targeted_entry() {
        let pool: Arc<Mutex<HashMap<(u64, String), u8>>> = Arc::new(Mutex::new(HashMap::new()));
        {
            let mut guard = pool.lock().await;
            guard.insert((1, "a:1".to_string()), 0);
            guard.insert((2, "b:2".to_string()), 0);
        }
        evict(&pool, 1, "a:1").await;
        let guard = pool.lock().await;
        assert!(!guard.contains_key(&(1, "a:1".to_string())));
        assert!(guard.contains_key(&(2, "b:2".to_string())));
    }

    fn seeded_pool() -> Arc<Mutex<HashMap<(u64, String), u8>>> {
        let pool = Arc::new(Mutex::new(HashMap::new()));
        pool.try_lock().unwrap().insert((1, "a:1".to_string()), 0);
        pool
    }

    // A call that never resolves must be cut at the hard-TTL deadline and
    // surfaced as `Unreachable` (so openraft backs off), with the pooled client
    // evicted so the next attempt reconnects. A never-resolving future against a
    // short real deadline is deterministic: the timer is the only thing that can
    // complete the call.
    #[tokio::test]
    async fn unary_call_deadline_elapse_evicts_and_reports_unreachable() {
        let pool = seeded_pool();
        let never = std::future::pending::<Result<tonic::Response<u8>, tonic::Status>>();
        let err = unary_call(&pool, 1, "a:1", Duration::from_millis(10), never)
            .await
            .expect_err("a never-resolving call must hit the deadline");
        assert!(matches!(err, RPCError::Unreachable(_)));
        assert!(!pool.lock().await.contains_key(&(1, "a:1".to_string())));
    }

    // A transport error inside the deadline propagates as `Network` (retry at
    // once) and also evicts the client.
    #[tokio::test]
    async fn unary_call_transport_error_evicts_and_reports_network() {
        let pool = seeded_pool();
        let failed = async { Err::<tonic::Response<u8>, _>(tonic::Status::unavailable("down")) };
        let err = unary_call(&pool, 1, "a:1", Duration::from_secs(5), failed)
            .await
            .expect_err("a transport error must propagate");
        assert!(matches!(err, RPCError::Network(_)));
        assert!(!pool.lock().await.contains_key(&(1, "a:1".to_string())));
    }

    // A successful call returns the decoded body and must NOT evict the client.
    #[tokio::test]
    async fn unary_call_success_returns_body_and_keeps_client() {
        let pool = seeded_pool();
        let ok = async { Ok(tonic::Response::new(42u8)) };
        let body = unary_call(&pool, 1, "a:1", Duration::from_secs(5), ok)
            .await
            .expect("a successful call returns its body");
        assert_eq!(body, 42);
        assert!(pool.lock().await.contains_key(&(1, "a:1".to_string())));
    }

    fn header_chunk() -> SnapshotChunk {
        let vote: VoteOf<TypeConfig> = Vote::new(1, 1);
        let meta = SnapshotMetaOf::<TypeConfig> {
            last_log_id: None,
            last_membership: Default::default(),
            snapshot_id: "test-snap".to_string(),
        };
        SnapshotChunk {
            kind: Some(ChunkKind::Header(SnapshotHeader {
                vote: postcard::to_stdvec(&vote).expect("encode vote"),
                meta: postcard::to_stdvec(&meta).expect("encode meta"),
            })),
        }
    }

    fn data_chunk(bytes: &[u8]) -> SnapshotChunk {
        SnapshotChunk {
            kind: Some(ChunkKind::Data(bytes.to_vec())),
        }
    }

    fn ok_stream(
        chunks: Vec<SnapshotChunk>,
    ) -> impl futures::Stream<Item = Result<SnapshotChunk, tonic::Status>> + Unpin {
        futures::stream::iter(chunks.into_iter().map(Ok))
    }

    #[tokio::test]
    async fn snapshot_over_limit_is_resource_exhausted() {
        // Two 600-byte data chunks (1200 bytes total) against a 1 KiB ceiling:
        // the second chunk pushes the running total past the limit and must be
        // refused before it is buffered.
        let chunks = vec![
            header_chunk(),
            data_chunk(&[0u8; 600]),
            data_chunk(&[0u8; 600]),
        ];
        let err = reassemble_snapshot(ok_stream(chunks), 1024)
            .await
            .expect_err("over-limit stream must be rejected");
        assert_eq!(err.code(), tonic::Code::ResourceExhausted);
    }

    #[tokio::test]
    async fn snapshot_under_limit_assembles() {
        let chunks = vec![header_chunk(), data_chunk(b"hello "), data_chunk(b"world")];
        let assembled = reassemble_snapshot(ok_stream(chunks), 1024)
            .await
            .expect("under-limit stream assembles");
        assert_eq!(assembled.data, b"hello world");
    }

    #[tokio::test]
    async fn data_before_header_is_invalid_argument() {
        let chunks = vec![data_chunk(b"premature")];
        let err = reassemble_snapshot(ok_stream(chunks), 1024)
            .await
            .expect_err("data before header must be rejected");
        assert_eq!(err.code(), tonic::Code::InvalidArgument);
    }
}