armdb/armour/

rpc.rs

use std::{collections::HashMap, path::Path, sync::Arc, thread::JoinHandle};

use armour_rpc::RpcError;
use armour_rpc::protocol::{
    OpCode, Request, RequestPayload, read_bound, read_bytes, read_u8, read_u32_be, read_u64_be,
    read_upsert_key,
};
use compio::buf::{IoBuf, IoBufMut};
use compio::io::framed::codec::{Decoder, Encoder};
use compio::{
    io::{
        AsyncRead, AsyncWrite,
        framed::{Framed, frame::LengthDelimited},
    },
    net::{TcpListener, ToSocketAddrsAsync, UnixListener},
    runtime::RuntimeBuilder,
};
use futures_util::{
    SinkExt, StreamExt,
    future::{Either, select},
    pin_mut,
};

use super::db::Db;
use super::handler::RpcHandler;

// ─── Response types (server-side only) ───────────────────────────────────────

#[derive(Debug)]
enum Response {
    Ok(ResponsePayload),
    Err { code: u16, message: String },
}

#[derive(Debug)]
pub(crate) struct CollectionMeta {
    pub name: String,
    pub hashname: u64,
    pub typ_hash: u64,
    pub version: u16,
    pub count: u64,
}

#[derive(Debug)]
enum ResponsePayload {
    Empty,
    Bool(bool),
    OptionalData(Option<Vec<u8>>),
    OptionalLen(Option<u32>),
    OptionalKV(Option<(Vec<u8>, Vec<u8>)>),
    KeyValues(Vec<(Vec<u8>, Vec<u8>)>),
    Keys(Vec<Vec<u8>>),
    Key(Vec<u8>),
    Count(u64),
    Collections(Vec<CollectionMeta>),
    SchemaBytes(Vec<u8>),
}

// ─── Codec ────────────────────────────────────────────────────────────────────

struct RpcCodec;

impl<B: IoBuf> Decoder<Request, B> for RpcCodec {
    type Error = RpcError;

    fn decode(&mut self, buf: &compio::buf::Slice<B>) -> Result<Request, Self::Error> {
        let bytes: &[u8] = buf;
        let mut pos = 0;
        let op_byte = read_u8(bytes, &mut pos)?;
        let op = OpCode::from_repr(op_byte)
            .ok_or_else(|| RpcError::Protocol("unknown opcode".to_string()))?;
        let hashname = read_u64_be(bytes, &mut pos)?;

        let payload = match op {
            OpCode::Get | OpCode::Contains => {
                let key = read_bytes(bytes, &mut pos)?;
                RequestPayload::Key(key)
            }
            OpCode::First | OpCode::Last | OpCode::ListCollections | OpCode::GetSchema => {
                RequestPayload::Empty
            }
            OpCode::Count => {
                let exact = read_u8(bytes, &mut pos)? != 0;
                RequestPayload::Count { exact }
            }
            OpCode::EntryLen => {
                let key = read_bytes(bytes, &mut pos)?;
                RequestPayload::EntryLen { key }
            }
            OpCode::Range | OpCode::RangeKeys => {
                let start = read_bound(bytes, &mut pos)?;
                let end = read_bound(bytes, &mut pos)?;
                let limit = read_u32_be(bytes, &mut pos)?;
                RequestPayload::Range { start, end, limit }
            }
            OpCode::Upsert => {
                let key = read_upsert_key(bytes, &mut pos)?;
                let flag_byte = read_u8(bytes, &mut pos)?;
                let flag = match flag_byte {
                    0 => None,
                    1 => Some(true),
                    2 => Some(false),
                    _ => {
                        return Err(RpcError::Protocol("invalid upsert flag".to_string()));
                    }
                };
                let value = read_bytes(bytes, &mut pos)?;
                RequestPayload::Upsert { key, flag, value }
            }
            OpCode::Remove => {
                let key = read_bytes(bytes, &mut pos)?;
                let soft = read_u8(bytes, &mut pos)? != 0;
                RequestPayload::Remove { key, soft }
            }
            OpCode::Take => {
                let key = read_bytes(bytes, &mut pos)?;
                let soft = read_u8(bytes, &mut pos)? != 0;
                RequestPayload::Take { key, soft }
            }
            OpCode::ApplyBatch => {
                let count = read_u32_be(bytes, &mut pos)? as usize;
                const MAX_PREALLOC: usize = 64 * 1024;
                let mut items = Vec::with_capacity(count.min(MAX_PREALLOC));
                for _ in 0..count {
                    let key = read_bytes(bytes, &mut pos)?;
                    let has_val = read_u8(bytes, &mut pos)?;
                    let val = match has_val {
                        0 => None,
                        1 => Some(read_bytes(bytes, &mut pos)?),
                        _ => {
                            return Err(RpcError::Protocol("invalid batch value tag".to_string()));
                        }
                    };
                    items.push((key, val));
                }
                RequestPayload::Batch(items)
            }
        };

        Ok(Request {
            op,
            hashname,
            payload,
        })
    }
}

impl<B: IoBufMut> Encoder<Response, B> for RpcCodec {
    type Error = RpcError;

    fn encode(&mut self, item: Response, buf: &mut B) -> Result<(), Self::Error> {
        let mut tmp = Vec::new();
        match item {
            Response::Ok(payload) => {
                tmp.push(0x00);
                encode_ok_payload(&mut tmp, payload);
            }
            Response::Err { code, message } => encode_err_frame(&mut tmp, code, &message),
        }
        buf.extend_from_slice(&tmp)
            .map_err(|e| RpcError::Other(e.to_string()))?;
        Ok(())
    }
}

/// Encode an error frame: `[0x01][code: u16 BE][len: u16 BE][msg bytes]`.
///
/// The length prefix is a `u16`, so the message is capped at `u16::MAX` bytes.
/// Truncate the byte slice to match — otherwise a message ≥ 64 KiB writes a
/// prefix wrapped mod 65536 while emitting all bytes, leaving the client to read
/// `len` bytes and misparse the unconsumed remainder of the frame.
fn encode_err_frame(buf: &mut Vec<u8>, code: u16, message: &str) {
    buf.push(0x01);
    buf.extend_from_slice(&code.to_be_bytes());
    let max = message.len().min(u16::MAX as usize);
    let end = message.floor_char_boundary(max);
    let msg = &message.as_bytes()[..end];
    buf.extend_from_slice(&(msg.len() as u16).to_be_bytes());
    buf.extend_from_slice(msg);
}

/// Build a wire error [`Response`] applying the **5xx-opaque policy**.
///
/// - **4xx** (client error): the message is forwarded verbatim — the client must
///   understand what it got wrong so it can fix the request (unknown collection,
///   invalid op, bad key, schema mismatch, …).
/// - **5xx** (server error): the real cause is logged server-side via `tracing`,
///   and the wire message becomes a generic `"internal error"` so internal
///   details (file paths, corruption state, codec internals) don't leak to RPC
///   clients. `original` is the verbatim cause to log.
///
/// **501 is exempt.** [`DbError::NotImplemented`] is raised only by client input
/// the engine does not support (`soft = true` on remove/take), its `Display` is
/// a fixed string that cannot leak anything, and the client needs to see it to
/// know the request is permanently unsupported rather than transiently broken.
/// It is therefore forwarded verbatim and logged at `warn` — an `error` here
/// would raise operational alerts for a plain client mistake.
///
/// Note: `ShardMismatch` is a server-side routing bug (the engine mis-routed a
/// key), so it stays in the `_ => 500` catch-all of [`DbError::status_code`] and
/// is therefore opaque to the client.
fn err_response(code: u16, original: &str) -> Response {
    if code == 501 {
        tracing::warn!("RPC unsupported request ({code}): {original}");
        Response::Err {
            code,
            message: original.into(),
        }
    } else if code >= 500 {
        tracing::error!("RPC server error ({code}): {original}");
        Response::Err {
            code,
            message: "internal error".into(),
        }
    } else {
        Response::Err {
            code,
            message: original.into(),
        }
    }
}

fn encode_ok_payload(buf: &mut Vec<u8>, payload: ResponsePayload) {
    match payload {
        ResponsePayload::Empty => {}
        ResponsePayload::Bool(b) => {
            buf.push(if b { 1 } else { 0 });
        }
        ResponsePayload::OptionalData(opt) => match opt {
            None => buf.push(0),
            Some(data) => {
                buf.push(1);
                buf.extend_from_slice(&(data.len() as u32).to_be_bytes());
                buf.extend_from_slice(&data);
            }
        },
        ResponsePayload::OptionalLen(opt) => match opt {
            None => buf.push(0),
            Some(len) => {
                buf.push(1);
                buf.extend_from_slice(&len.to_be_bytes());
            }
        },
        ResponsePayload::OptionalKV(opt) => match opt {
            None => buf.push(0),
            Some((key, val)) => {
                buf.push(1);
                buf.extend_from_slice(&(key.len() as u32).to_be_bytes());
                buf.extend_from_slice(&key);
                buf.extend_from_slice(&(val.len() as u32).to_be_bytes());
                buf.extend_from_slice(&val);
            }
        },
        ResponsePayload::KeyValues(pairs) => {
            buf.extend_from_slice(&(pairs.len() as u32).to_be_bytes());
            for (key, val) in pairs {
                buf.extend_from_slice(&(key.len() as u32).to_be_bytes());
                buf.extend_from_slice(&key);
                buf.extend_from_slice(&(val.len() as u32).to_be_bytes());
                buf.extend_from_slice(&val);
            }
        }
        ResponsePayload::Keys(keys) => {
            buf.extend_from_slice(&(keys.len() as u32).to_be_bytes());
            for key in keys {
                buf.extend_from_slice(&(key.len() as u32).to_be_bytes());
                buf.extend_from_slice(&key);
            }
        }
        ResponsePayload::Key(key) => {
            buf.extend_from_slice(&(key.len() as u32).to_be_bytes());
            buf.extend_from_slice(&key);
        }
        ResponsePayload::Count(n) => {
            buf.extend_from_slice(&n.to_be_bytes());
        }
        ResponsePayload::SchemaBytes(bytes) => {
            buf.extend_from_slice(&bytes);
        }
        ResponsePayload::Collections(collections) => {
            buf.extend_from_slice(&(collections.len() as u32).to_be_bytes());
            for c in collections {
                let name = c.name.as_bytes();
                buf.extend_from_slice(&(name.len() as u32).to_be_bytes());
                buf.extend_from_slice(name);
                // partition_name = name for armdb
                buf.extend_from_slice(&(name.len() as u32).to_be_bytes());
                buf.extend_from_slice(name);
                buf.extend_from_slice(&c.hashname.to_be_bytes());
                buf.extend_from_slice(&c.typ_hash.to_be_bytes());
                buf.extend_from_slice(&c.version.to_be_bytes());
                buf.extend_from_slice(&c.count.to_be_bytes());
            }
        }
    }
}

// ─── Server ───────────────────────────────────────────────────────────────────

pub type TreeMap = Arc<HashMap<u64, Arc<dyn RpcHandler>>>;

/// Backoff bounds for the accept loops. On a persistent accept failure (e.g.
/// `EMFILE` when the fd table is exhausted) an immediate retry spins a hot error
/// loop; back off between failures and reset on the next successful accept.
const ACCEPT_BACKOFF_MIN: std::time::Duration = std::time::Duration::from_millis(5);
const ACCEPT_BACKOFF_MAX: std::time::Duration = std::time::Duration::from_millis(500);

async fn handle_connection<R, W>(
    reader: R,
    writer: W,
    trees: TreeMap,
    mut stop_rx: async_broadcast::Receiver<()>,
) where
    R: AsyncRead + Unpin + 'static,
    W: AsyncWrite + Unpin + 'static,
{
    let mut framed = Framed::new::<Response, Request>(RpcCodec, LengthDelimited::new())
        .with_reader(reader)
        .with_writer(writer);

    loop {
        let next_request = framed.next();
        let stop_signal = stop_rx.recv();
        pin_mut!(next_request);
        pin_mut!(stop_signal);

        match select(next_request, stop_signal).await {
            Either::Left((Some(result), _)) => {
                let response = match result {
                    Ok(request) => dispatch(&trees, request),
                    // A `Decoder` error (unknown opcode, invalid upsert flag,
                    // invalid batch tag, truncated frame) is always *client*
                    // garbage — the request never reached a collection — so it
                    // is 400, not 500. ("Corrupt data already in the DB" is a
                    // different path and stays 500 via `DbError::status_code`.)
                    Err(e) => err_response(400, &e.to_string()),
                };
                if framed.send(response).await.is_err() {
                    break;
                }
            }
            Either::Left((None, _)) => break,
            Either::Right(_) => break,
        }
    }
}

fn dispatch(trees: &TreeMap, request: Request) -> Response {
    if request.op == OpCode::ListCollections {
        return list_collections(trees);
    }
    if request.op == OpCode::GetSchema {
        return get_schema(trees, request.hashname);
    }

    let handler = match trees.get(&request.hashname) {
        Some(h) => h,
        None => {
            return Response::Err {
                code: 404,
                message: format!("tree not found: {:#X}", request.hashname),
            };
        }
    };

    let result = match (request.op, request.payload) {
        (OpCode::Get, RequestPayload::Key(key)) => handler
            .get(&key)
            .map(|v| Response::Ok(ResponsePayload::OptionalData(v))),
        (OpCode::Contains, RequestPayload::Key(key)) => handler
            .contains(&key)
            .map(|b| Response::Ok(ResponsePayload::Bool(b))),
        (OpCode::First, RequestPayload::Empty) => handler
            .first()
            .map(|v| Response::Ok(ResponsePayload::OptionalKV(v))),
        (OpCode::Last, RequestPayload::Empty) => handler
            .last()
            .map(|v| Response::Ok(ResponsePayload::OptionalKV(v))),
        (OpCode::Range, RequestPayload::Range { start, end, limit }) => handler
            .range(start, end, limit)
            .map(|v| Response::Ok(ResponsePayload::KeyValues(v))),
        (OpCode::RangeKeys, RequestPayload::Range { start, end, limit }) => handler
            .range_keys(start, end, limit)
            .map(|v| Response::Ok(ResponsePayload::Keys(v))),
        (OpCode::Count, RequestPayload::Count { exact }) => handler
            .count(exact)
            .map(|n| Response::Ok(ResponsePayload::Count(n))),
        (OpCode::Upsert, RequestPayload::Upsert { key, flag, value }) => handler
            .upsert(key, flag, value)
            .map(|k| Response::Ok(ResponsePayload::Key(k))),
        (OpCode::Remove, RequestPayload::Remove { key, soft }) => handler
            .remove(&key, soft)
            .map(|_| Response::Ok(ResponsePayload::Empty)),
        (OpCode::Take, RequestPayload::Take { key, soft }) => handler
            .take(&key, soft)
            .map(|v| Response::Ok(ResponsePayload::OptionalData(v))),
        (OpCode::EntryLen, RequestPayload::EntryLen { key }) => handler
            .entry_len(&key)
            .map(|v| Response::Ok(ResponsePayload::OptionalLen(v))),
        (OpCode::ApplyBatch, RequestPayload::Batch(items)) => handler
            .apply_batch(items)
            .map(|_| Response::Ok(ResponsePayload::Empty)),
        _ => {
            return Response::Err {
                code: 400,
                message: "invalid op/payload combination".into(),
            };
        }
    };

    match result {
        Ok(resp) => resp,
        Err(e) => err_response(e.status_code(), &e.to_string()),
    }
}

fn get_schema(trees: &TreeMap, hashname: u64) -> Response {
    let handler = match trees.get(&hashname) {
        Some(h) => h,
        None => {
            return Response::Err {
                code: 404,
                message: format!("tree not found: {hashname:#X}"),
            };
        }
    };
    let schema = handler.schema();
    match serde_json::to_vec(&schema) {
        Ok(bytes) => Response::Ok(ResponsePayload::SchemaBytes(bytes)),
        // Server-side serialization failure → 5xx (opaque to the client).
        Err(e) => err_response(500, &e.to_string()),
    }
}

fn list_collections(trees: &TreeMap) -> Response {
    let collections = trees
        .iter()
        .map(|(&hashname, handler)| {
            let (typ_hash, version) = handler.info();
            let count = handler.count(false).unwrap_or_else(|e| {
                tracing::warn!("count() failed for '{}': {e}", handler.name());
                0
            });
            CollectionMeta {
                name: handler.name().to_string(),
                hashname,
                typ_hash,
                version,
                count,
            }
        })
        .collect();
    Response::Ok(ResponsePayload::Collections(collections))
}

async fn accept_tcp(
    listener: TcpListener,
    trees: TreeMap,
    mut shutdown_rx: async_broadcast::Receiver<()>,
    mut local_rx: async_broadcast::Receiver<()>,
) {
    let mut backoff = ACCEPT_BACKOFF_MIN;
    loop {
        let accept_fut = listener.accept();
        // AI-05: stop on either the Db-global shutdown or this listener's own
        // handle signal (`RpcHandle::stop`), so a standalone handle can terminate
        // its accept loop without a global shutdown.
        let global_stop = shutdown_rx.recv();
        let local_stop = local_rx.recv();
        pin_mut!(accept_fut);
        pin_mut!(global_stop);
        pin_mut!(local_stop);
        let stop_fut = select(global_stop, local_stop);

        match select(accept_fut, stop_fut).await {
            Either::Left((Ok((stream, _addr)), _)) => {
                backoff = ACCEPT_BACKOFF_MIN;
                let trees = trees.clone();
                let (reader, writer) = stream.into_split();
                let rx = shutdown_rx.clone();
                compio::runtime::spawn(handle_connection(reader, writer, trees, rx)).detach();
            }
            Either::Left((Err(e), _)) => {
                tracing::error!("tcp accept error: {e}");
                compio::time::sleep(backoff).await;
                backoff = (backoff * 2).min(ACCEPT_BACKOFF_MAX);
            }
            Either::Right(_) => {
                tracing::info!("shutting down TCP server");
                break;
            }
        }
    }
}

async fn accept_uds(
    listener: UnixListener,
    trees: TreeMap,
    mut shutdown_rx: async_broadcast::Receiver<()>,
    mut local_rx: async_broadcast::Receiver<()>,
) {
    let mut backoff = ACCEPT_BACKOFF_MIN;
    loop {
        let accept_fut = listener.accept();
        let global_stop = shutdown_rx.recv();
        let local_stop = local_rx.recv();
        pin_mut!(accept_fut);
        pin_mut!(global_stop);
        pin_mut!(local_stop);
        let stop_fut = select(global_stop, local_stop);

        match select(accept_fut, stop_fut).await {
            Either::Left((Ok((stream, _addr)), _)) => {
                backoff = ACCEPT_BACKOFF_MIN;
                let trees = trees.clone();
                let (reader, writer) = stream.into_split();
                let rx = shutdown_rx.clone();
                compio::runtime::spawn(handle_connection(reader, writer, trees, rx)).detach();
            }
            Either::Left((Err(e), _)) => {
                tracing::error!("uds accept error: {e}");
                compio::time::sleep(backoff).await;
                backoff = (backoff * 2).min(ACCEPT_BACKOFF_MAX);
            }
            Either::Right(_) => {
                tracing::info!("shutting down UDS server");
                break;
            }
        }
    }
}

// ─── RpcHandle ───────────────────────────────────────────────────────────────

/// Handle for a running RPC listener thread. Dropping it triggers shutdown
/// and joins the thread.
pub struct RpcHandle {
    thread: Option<JoinHandle<()>>,
    shutdown_tx: async_broadcast::Sender<()>,
}

impl RpcHandle {
    pub fn stop(&mut self) {
        let _ = self.shutdown_tx.try_broadcast(());
        if let Some(h) = self.thread.take() {
            let _ = h.join();
        }
    }
}

impl Drop for RpcHandle {
    fn drop(&mut self) {
        self.stop();
    }
}

// ─── Db listen methods ────────────────────────────────────────────────────

impl Db {
    /// Start a TCP RPC listener on `addr`.
    ///
    /// # Trust model
    ///
    /// The RPC server assumes a **trusted network**: peers are not authenticated,
    /// operations are not authorized, and heavy requests are unbudgeted — a `Range`
    /// with unbounded bounds and `limit = 0` (or `Count { exact: true }`) scans the
    /// whole collection, materializes the full result in RAM, and blocks the
    /// single-thread runtime while it runs. `limit = 0` deliberately means
    /// "everything"; there is no default server-side clamp (it would be
    /// behavior-breaking). Bind to loopback / a private interface (or prefer
    /// [`listen_uds`](Self::listen_uds)) and front untrusted networks with an mTLS
    /// tunnel — do not expose this listener directly. See the "Network trust model"
    /// section in `docs/design.md`.
    pub fn listen_tcp(&self, addr: impl ToSocketAddrsAsync + Send + 'static) {
        let trees = self.build_tree_map();
        let shutdown_rx = self.shutdown.subscribe_broadcast();
        // AI-05: keep the receiver and drive it in the accept loop so
        // `RpcHandle::stop` can terminate this listener on its own. The sender
        // clone stored on the handle keeps the channel alive.
        let (local_tx, local_rx) = async_broadcast::broadcast::<()>(1);
        let shutdown_tx = local_tx.clone();

        let thread = std::thread::spawn(move || {
            // AI-04: a build/bind failure must log and return, not panic a
            // detached thread — a silent panic leaves callers believing the RPC
            // listener is up when nothing is accepting connections.
            let runtime = match RuntimeBuilder::new().build() {
                Ok(rt) => rt,
                Err(e) => {
                    tracing::error!("failed to build compio runtime for TCP RPC listener: {e}");
                    return;
                }
            };
            runtime.block_on(async move {
                let listener = match TcpListener::bind(addr).await {
                    Ok(l) => l,
                    Err(e) => {
                        tracing::error!("failed to bind TCP RPC listener: {e}");
                        return;
                    }
                };
                accept_tcp(listener, trees, shutdown_rx, local_rx).await;
            });
        });
        crate::sync::lock(&self.rpc_handles).push(RpcHandle {
            thread: Some(thread),
            shutdown_tx,
        });
    }

    pub fn listen_uds(&self, path: impl AsRef<Path> + Send + 'static) {
        let trees = self.build_tree_map();
        let shutdown_rx = self.shutdown.subscribe_broadcast();
        let (local_tx, local_rx) = async_broadcast::broadcast::<()>(1);
        let shutdown_tx = local_tx.clone();

        let thread = std::thread::spawn(move || {
            let runtime = match RuntimeBuilder::new().build() {
                Ok(rt) => rt,
                Err(e) => {
                    tracing::error!("failed to build compio runtime for UDS RPC listener: {e}");
                    return;
                }
            };
            runtime.block_on(async move {
                let _ = std::fs::remove_file(path.as_ref());
                let listener = match UnixListener::bind(path).await {
                    Ok(l) => l,
                    Err(e) => {
                        tracing::error!("failed to bind UDS RPC listener: {e}");
                        return;
                    }
                };
                accept_uds(listener, trees, shutdown_rx, local_rx).await;
            });
        });
        crate::sync::lock(&self.rpc_handles).push(RpcHandle {
            thread: Some(thread),
            shutdown_tx,
        });
    }
}

#[cfg(test)]
mod err_frame_tests {
    use super::encode_err_frame;

    // AI-11: the length prefix is a `u16`; a message ≥ 64 KiB must be truncated so
    // the prefix equals the bytes actually written. Pre-fix the prefix wrapped mod
    // 65536 while all bytes were emitted, desyncing the client frame parser.
    #[test]
    fn err_frame_len_prefix_matches_written_bytes_when_over_u16() {
        let big = "x".repeat(70_000);

        // The pre-fix formula (`len as u16`) disagreed with the byte count.
        let buggy_prefix = big.len() as u16; // 70000 as u16 == 4464
        assert_ne!(
            buggy_prefix as usize,
            big.len(),
            "sanity: pre-fix u16 prefix disagreed with the byte count"
        );

        let mut buf = Vec::new();
        encode_err_frame(&mut buf, 500, &big);

        // Frame layout: [tag:1][code:2][len:2][msg…]
        let len = u16::from_be_bytes([buf[3], buf[4]]) as usize;
        let msg_written = buf.len() - 5;
        assert_eq!(
            len, msg_written,
            "len prefix must equal bytes actually written"
        );
        assert_eq!(len, u16::MAX as usize, "message capped at u16::MAX");
    }

    #[test]
    fn err_frame_small_message_roundtrips() {
        let mut buf = Vec::new();
        encode_err_frame(&mut buf, 404, "not found");
        assert_eq!(buf[0], 0x01);
        assert_eq!(u16::from_be_bytes([buf[1], buf[2]]), 404);
        let len = u16::from_be_bytes([buf[3], buf[4]]) as usize;
        assert_eq!(len, "not found".len());
        assert_eq!(&buf[5..], b"not found");
    }

    #[test]
    fn err_frame_truncates_at_utf8_boundary() {
        let message = format!("{}€", "a".repeat(u16::MAX as usize - 1));
        let mut buf = Vec::new();
        encode_err_frame(&mut buf, 500, &message);
        let len = u16::from_be_bytes([buf[3], buf[4]]) as usize;
        let encoded = &buf[5..5 + len];
        assert_eq!(len, u16::MAX as usize - 1);
        assert_eq!(std::str::from_utf8(encoded).unwrap(), "a".repeat(len));
        assert_eq!(buf.len(), 5 + len);
    }
}

#[cfg(test)]
mod rpc_handle_tests {
    use super::Db;
    use std::time::Duration;

    // AI-05: `RpcHandle::stop` must be able to terminate its own listener without a
    // Db-global shutdown. Pre-fix the handle's shutdown sender had no live
    // receiver, so `stop` broadcast into the void and `join` hung.
    #[test]
    fn rpc_handle_stop_terminates_listener_without_global_shutdown() {
        let dir = tempfile::tempdir().unwrap();
        let db = Db::open_test(dir.path()).unwrap();
        let sock = dir.path().join("rpc.sock");
        db.listen_uds(sock);

        // Let the listener thread bind and enter its accept loop.
        std::thread::sleep(Duration::from_millis(150));

        // Take the handle out so stopping it does not depend on Db-global shutdown.
        let handle = crate::sync::lock(&db.rpc_handles)
            .pop()
            .expect("one rpc handle");

        // Run stop() with a timeout: a regression (accept loop ignores the local
        // signal) then fails fast instead of hanging the whole test run.
        let (tx, rx) = std::sync::mpsc::channel();
        std::thread::spawn(move || {
            let mut handle = handle;
            handle.stop();
            let _ = tx.send(());
        });

        let stopped = rx.recv_timeout(Duration::from_secs(5)).is_ok();
        assert!(
            stopped,
            "RpcHandle::stop did not terminate the listener (accept loop ignored the local signal)"
        );
    }
}

#[cfg(test)]
mod dispatch_status_code_tests {
    use crate::error::{DbError, SchemaMismatchKind};

    #[test]
    fn key_not_found_maps_to_404() {
        assert_eq!(DbError::KeyNotFound.status_code(), 404);
    }
    #[test]
    fn key_exists_maps_to_409() {
        assert_eq!(DbError::KeyExists.status_code(), 409);
    }
    #[test]
    fn client_error_maps_to_400() {
        assert_eq!(DbError::Client("bad").status_code(), 400);
    }
    #[test]
    fn not_implemented_maps_to_501() {
        assert_eq!(DbError::NotImplemented.status_code(), 501);
    }
    #[test]
    fn schema_mismatch_maps_to_422() {
        let e = DbError::SchemaMismatch {
            name: "x".into(),
            kind: SchemaMismatchKind::MissingStep { from: 0 },
        };
        assert_eq!(e.status_code(), 422);
    }
    #[test]
    fn tx_conflict_maps_to_409() {
        // D4: TxConflict is a client-actionable "retry" signal → 409, not 500.
        assert_eq!(DbError::TxConflict.status_code(), 409);
    }
    #[test]
    fn shard_mismatch_stays_500() {
        // D4: wrong-shard routing is a server bug → stays opaque 500 (not 4xx).
        assert_eq!(DbError::ShardMismatch.status_code(), 500);
    }
}

#[cfg(test)]
mod err_response_policy_tests {
    use super::{Response, err_response};

    fn unpack(r: Response) -> (u16, String) {
        match r {
            Response::Err { code, message } => (code, message),
            Response::Ok(_) => panic!("expected Err, got Ok"),
        }
    }

    #[test]
    fn client_error_4xx_forwards_message_verbatim() {
        // 4xx is the client's fault: it must see *what* it got wrong.
        let (code, msg) = unpack(err_response(404, "tree not found: 0xDEAD"));
        assert_eq!(code, 404);
        assert_eq!(msg, "tree not found: 0xDEAD");
    }

    #[test]
    fn server_error_5xx_sends_opaque_message() {
        // 5xx leaks nothing: the original (paths, corruption state) must NOT
        // appear on the wire — only a generic "internal error".
        let original = "corrupted entry at offset 12345 in /var/data/shard_007/000012.data";
        let (code, msg) = unpack(err_response(500, original));
        assert_eq!(code, 500);
        assert_eq!(msg, "internal error");
        assert!(
            !msg.contains("12345") && !msg.contains("/var/data"),
            "5xx response must not echo internal details"
        );
    }

    #[test]
    fn not_implemented_501_forwards_message() {
        // 501 is exempt from the opaque policy: `DbError::NotImplemented` is
        // raised only by unsupported *client* input (`soft = true`), so the
        // client must see that the request is permanently unsupported rather
        // than a transient server fault.
        let (code, msg) = unpack(err_response(501, "not implemented"));
        assert_eq!(code, 501);
        assert_eq!(msg, "not implemented");
    }
}