armdb 0.14.2

use std::io::Read;
use std::marker::PhantomData;
use std::net::{TcpStream, ToSocketAddrs};
use std::ops::Bound;
use std::os::unix::net::UnixStream;
use std::path::Path;

use armour_rpc::{ClientConfig, DeadlineWriter, Result, RpcClient, RpcError, WriteKey};

use crate::{Codec, CollectionMeta, Key};

pub struct TypedRpcClient<Item, C, S>
where
    Item: CollectionMeta,
    C: Codec<Item>,
{
    inner: RpcClient<S>,
    codec: C,
    _phantom: PhantomData<Item>,
}

fn encode_bound<K: Key>(bound: Bound<&K>) -> Bound<Vec<u8>> {
    match bound {
        Bound::Unbounded => Bound::Unbounded,
        Bound::Included(k) => Bound::Included(k.as_bytes().to_vec()),
        Bound::Excluded(k) => Bound::Excluded(k.as_bytes().to_vec()),
    }
}

/// Decode a key received from the server into a typed `K`.
///
/// `Key::from_bytes` transmutes the slice with only a `debug_assert` on length —
/// a server that returns a wrong-length key (schema drift, a hostile or buggy
/// peer) would trigger an out-of-bounds read in release builds. Validate the
/// length here and return a protocol error instead of reading past the buffer.
fn decode_id<K: Key>(key: &[u8]) -> Result<K> {
    if key.len() != size_of::<K>() {
        return Err(RpcError::Protocol(format!(
            "key length mismatch: got {} bytes, expected {}",
            key.len(),
            size_of::<K>()
        )));
    }
    Ok(K::from_bytes(key))
}

impl<Item, C> TypedRpcClient<Item, C, TcpStream>
where
    Item: CollectionMeta,
    C: Codec<Item> + Default,
{
    pub fn connect_tcp(addr: impl ToSocketAddrs, cfg: ClientConfig) -> Result<Self> {
        let inner = RpcClient::connect_tcp_by_name(addr, Item::NAME, cfg)?;
        Ok(Self {
            inner,
            codec: C::default(),
            _phantom: PhantomData,
        })
    }
}

impl<Item, C> TypedRpcClient<Item, C, UnixStream>
where
    Item: CollectionMeta,
    C: Codec<Item> + Default,
{
    pub fn connect_uds(path: impl AsRef<Path>, cfg: ClientConfig) -> Result<Self> {
        let inner = RpcClient::connect_uds_by_name(path, Item::NAME, cfg)?;
        Ok(Self {
            inner,
            codec: C::default(),
            _phantom: PhantomData,
        })
    }
}

impl<Item, C, S> TypedRpcClient<Item, C, S>
where
    Item: CollectionMeta,
    C: Codec<Item>,
    S: Read + DeadlineWriter,
{
    pub fn get(&mut self, id: &Item::SelfId) -> Result<Option<Item>> {
        let data = self.inner.get(id.as_bytes())?;
        match data {
            None => Ok(None),
            Some(bytes) => {
                let item = self
                    .codec
                    .decode_from(&bytes)
                    .map_err(|e| RpcError::Encoding(e.to_string()))?;
                Ok(Some(item))
            }
        }
    }

    pub fn contains(&mut self, id: &Item::SelfId) -> Result<bool> {
        self.inner.contains(id.as_bytes())
    }

    pub fn entry_len(&mut self, id: &Item::SelfId) -> Result<Option<u32>> {
        self.inner.entry_len(id.as_bytes())
    }

    pub fn first(&mut self) -> Result<Option<(Item::SelfId, Item)>> {
        let data = self.inner.first()?;
        self.decode_kv(data)
    }

    pub fn last(&mut self) -> Result<Option<(Item::SelfId, Item)>> {
        let data = self.inner.last()?;
        self.decode_kv(data)
    }

    pub fn range(
        &mut self,
        start: Bound<&Item::SelfId>,
        end: Bound<&Item::SelfId>,
    ) -> Result<Vec<(Item::SelfId, Item)>> {
        let start = encode_bound::<Item::SelfId>(start);
        let end = encode_bound::<Item::SelfId>(end);
        let data = self.inner.range(start, end, 0)?;
        data.into_iter()
            .map(|(key, val)| {
                let id = decode_id::<Item::SelfId>(&key)?;
                let item = self
                    .codec
                    .decode_from(&val)
                    .map_err(|e| RpcError::Encoding(e.to_string()))?;
                Ok((id, item))
            })
            .collect()
    }

    pub fn range_keys(
        &mut self,
        start: Bound<&Item::SelfId>,
        end: Bound<&Item::SelfId>,
    ) -> Result<Vec<Item::SelfId>> {
        let start = encode_bound::<Item::SelfId>(start);
        let end = encode_bound::<Item::SelfId>(end);
        let data = self.inner.range_keys(start, end, 0)?;
        data.iter()
            .map(|key| decode_id::<Item::SelfId>(key))
            .collect()
    }

    pub fn count(&mut self, exact: bool) -> Result<u64> {
        self.inner.count(exact)
    }

    pub fn write(&mut self, id: &Item::SelfId, item: &Item) -> Result<()> {
        let value = self.encode_value(item)?;
        self.inner
            .write(WriteKey::Provided(id.as_bytes().to_vec()), None, value)?;
        Ok(())
    }

    pub fn insert(&mut self, id: &Item::SelfId, item: &Item) -> Result<()> {
        let value = self.encode_value(item)?;
        self.inner.write(
            WriteKey::Provided(id.as_bytes().to_vec()),
            Some(false),
            value,
        )?;
        Ok(())
    }

    pub fn update(&mut self, id: &Item::SelfId, item: &Item) -> Result<()> {
        let value = self.encode_value(item)?;
        self.inner.write(
            WriteKey::Provided(id.as_bytes().to_vec()),
            Some(true),
            value,
        )?;
        Ok(())
    }

    pub fn remove(&mut self, id: &Item::SelfId) -> Result<()> {
        self.inner.remove(id.as_bytes(), false)
    }

    pub fn take(&mut self, id: &Item::SelfId) -> Result<Option<Item>> {
        let data = self.inner.take(id.as_bytes(), false)?;
        match data {
            None => Ok(None),
            Some(bytes) => {
                let item = self
                    .codec
                    .decode_from(&bytes)
                    .map_err(|e| RpcError::Encoding(e.to_string()))?;
                Ok(Some(item))
            }
        }
    }

    pub fn apply_batch(
        &mut self,
        items: impl Iterator<Item = (Item::SelfId, Option<Item>)>,
    ) -> Result<()> {
        let raw: Vec<(Vec<u8>, Option<Vec<u8>>)> = items
            .map(|(id, val)| -> Result<_> {
                let key = id.as_bytes().to_vec();
                let value = val.map(|v| self.encode_value(&v)).transpose()?;
                Ok((key, value))
            })
            .collect::<Result<_>>()?;
        self.inner.apply_batch(raw)
    }

    fn decode_kv(&self, data: Option<(Vec<u8>, Vec<u8>)>) -> Result<Option<(Item::SelfId, Item)>> {
        match data {
            None => Ok(None),
            Some((key, val)) => {
                let id = decode_id::<Item::SelfId>(&key)?;
                let item = self
                    .codec
                    .decode_from(&val)
                    .map_err(|e| RpcError::Encoding(e.to_string()))?;
                Ok(Some((id, item)))
            }
        }
    }

    fn encode_value(&self, item: &Item) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
        self.codec
            .encode_to(item, &mut buf)
            .map_err(|e| RpcError::Encoding(e.to_string()))?;
        Ok(buf)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(armour_core::GetType)]
    struct TestItem;

    impl CollectionMeta for TestItem {
        type SelfId = [u8; 8];
        const NAME: &'static str = "typed_rpc_client_test_items";
    }

    #[derive(Default)]
    struct TestCodec;

    impl Codec<TestItem> for TestCodec {
        fn encode_to(&self, _: &TestItem, _: &mut Vec<u8>) -> crate::DbResult<()> {
            Ok(())
        }

        fn decode_from(&self, _: &[u8]) -> crate::DbResult<TestItem> {
            Ok(TestItem)
        }
    }

    // AI-08: a wrong-length key from the server must produce a protocol error,
    // not the out-of-bounds read that `Key::from_bytes` (transmute under a bare
    // `debug_assert`) would perform in a release build.
    #[test]
    fn decode_id_rejects_wrong_length_key() {
        // `[u8; 8]` is an 8-byte key type.
        let too_short = [0u8; 4];
        let err = decode_id::<[u8; 8]>(&too_short).unwrap_err();
        assert!(
            matches!(err, RpcError::Protocol(_)),
            "expected Protocol error, got {err:?}"
        );

        let too_long = [0u8; 12];
        assert!(matches!(
            decode_id::<[u8; 8]>(&too_long),
            Err(RpcError::Protocol(_))
        ));
    }

    #[test]
    fn decode_id_accepts_exact_length_key() {
        let exact = [1u8, 2, 3, 4, 5, 6, 7, 8];
        let id = decode_id::<[u8; 8]>(&exact).expect("exact length must decode");
        assert_eq!(id, exact);
    }

    #[test]
    fn typed_client_connect_tcp_accepts_client_config() {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind listener");
        let addr = listener.local_addr().expect("listener address");
        let config = armour_rpc::ClientConfig::default();

        let client: TypedRpcClient<TestItem, TestCodec, std::net::TcpStream> =
            TypedRpcClient::connect_tcp(addr, config).expect("connect typed client");
        drop(client);
    }

    #[test]
    fn typed_client_connect_uds_accepts_client_config() {
        let dir = tempfile::tempdir().expect("create temporary directory");
        let path = dir.path().join("typed-client.sock");
        let listener = std::os::unix::net::UnixListener::bind(&path).expect("bind listener");
        let config = armour_rpc::ClientConfig::default();

        let client: TypedRpcClient<TestItem, TestCodec, std::os::unix::net::UnixStream> =
            TypedRpcClient::connect_uds(&path, config).expect("connect typed client");
        drop(client);
        drop(listener);
    }
}