armour-rpc 0.3.0

use crate::error::RpcError;
use std::ops::Bound;
use strum::FromRepr;

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, FromRepr)]
pub enum OpCode {
    Get = 0x01,
    Contains = 0x02,
    First = 0x03,
    Last = 0x04,
    Range = 0x05,
    RangeKeys = 0x06,
    Count = 0x07,
    EntryLen = 0x08,
    Upsert = 0x10,
    Remove = 0x11,
    Take = 0x12,
    ApplyBatch = 0x13,
    ListCollections = 0x20,
    GetSchema = 0x21,
}

#[derive(Debug, Clone)]
pub enum UpsertKey {
    Sequence,
    Provided(Vec<u8>),
}

#[derive(Debug)]
pub struct Request {
    pub op: OpCode,
    pub hashname: u64,
    pub payload: RequestPayload,
}

#[derive(Debug)]
pub enum RequestPayload {
    Key(Vec<u8>),
    Empty,
    Range {
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
        limit: u32,
    },
    Upsert {
        key: UpsertKey,
        flag: Option<bool>,
        value: Vec<u8>,
    },
    Remove {
        key: Vec<u8>,
        soft: bool,
    },
    Take {
        key: Vec<u8>,
        soft: bool,
    },
    Count {
        exact: bool,
    },
    EntryLen {
        key: Vec<u8>,
    },
    Batch(Vec<(Vec<u8>, Option<Vec<u8>>)>),
    ListCollections,
}

/// Response decoded from the wire.
/// `Ok` contains the raw payload bytes; each client method decodes them.
#[derive(Debug)]
pub enum Response {
    Ok(Vec<u8>),
    Err { code: u16, message: String },
}

// --- Encoding helpers ---

pub fn write_bound(buf: &mut Vec<u8>, bound: &Bound<Vec<u8>>) {
    match bound {
        Bound::Unbounded => buf.push(0x00),
        Bound::Included(bytes) => {
            buf.push(0x01);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
        Bound::Excluded(bytes) => {
            buf.push(0x02);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
    }
}

pub fn read_bound(buf: &[u8], pos: &mut usize) -> Result<Bound<Vec<u8>>, RpcError> {
    let tag = read_u8(buf, pos)?;
    match tag {
        0x00 => Ok(Bound::Unbounded),
        0x01 => {
            let bytes = read_bytes(buf, pos)?;
            Ok(Bound::Included(bytes))
        }
        0x02 => {
            let bytes = read_bytes(buf, pos)?;
            Ok(Bound::Excluded(bytes))
        }
        _ => Err(RpcError::Protocol("invalid bound tag".to_string())),
    }
}

pub fn read_upsert_key(buf: &[u8], pos: &mut usize) -> Result<UpsertKey, RpcError> {
    let tag = read_u8(buf, pos)?;
    match tag {
        0x00 => Ok(UpsertKey::Sequence),
        0x01 => {
            let key = read_bytes(buf, pos)?;
            Ok(UpsertKey::Provided(key))
        }
        _ => Err(RpcError::Protocol("invalid upsert key tag".to_string())),
    }
}

pub fn write_upsert_key(buf: &mut Vec<u8>, key: &UpsertKey) {
    match key {
        UpsertKey::Sequence => buf.push(0x00),
        UpsertKey::Provided(bytes) => {
            buf.push(0x01);
            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
            buf.extend_from_slice(bytes);
        }
    }
}

// --- Primitive read/write helpers ---

pub fn read_u8(buf: &[u8], pos: &mut usize) -> Result<u8, RpcError> {
    if *pos >= buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = buf[*pos];
    *pos += 1;
    Ok(v)
}

pub fn read_u16_be(buf: &[u8], pos: &mut usize) -> Result<u16, RpcError> {
    if *pos + 2 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]);
    *pos += 2;
    Ok(v)
}

pub fn read_u32_be(buf: &[u8], pos: &mut usize) -> Result<u32, RpcError> {
    if *pos + 4 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u32::from_be_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
    *pos += 4;
    Ok(v)
}

pub fn read_u64_be(buf: &[u8], pos: &mut usize) -> Result<u64, RpcError> {
    if *pos + 8 > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = u64::from_be_bytes([
        buf[*pos],
        buf[*pos + 1],
        buf[*pos + 2],
        buf[*pos + 3],
        buf[*pos + 4],
        buf[*pos + 5],
        buf[*pos + 6],
        buf[*pos + 7],
    ]);
    *pos += 8;
    Ok(v)
}

pub fn read_bytes(buf: &[u8], pos: &mut usize) -> Result<Vec<u8>, RpcError> {
    let len = read_u32_be(buf, pos)? as usize;
    if *pos + len > buf.len() {
        return Err(RpcError::UnexpectedEof);
    }
    let v = buf[*pos..*pos + len].to_vec();
    *pos += len;
    Ok(v)
}

pub fn write_bytes(buf: &mut Vec<u8>, bytes: &[u8]) {
    buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
    buf.extend_from_slice(bytes);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client::decode_bool;
    use crate::codec::encode_request_payload;

    // Helpers that mirror the server-side encoding used in response decoders.

    fn encode_optional_value(v: Option<&[u8]>) -> Vec<u8> {
        let mut buf = Vec::new();
        match v {
            None => buf.push(0u8),
            Some(bytes) => {
                buf.push(1u8);
                write_bytes(&mut buf, bytes);
            }
        }
        buf
    }

    fn decode_optional_value(buf: &[u8]) -> Result<Option<Vec<u8>>, RpcError> {
        let mut pos = 0;
        let flag = read_u8(buf, &mut pos)?;
        if flag == 0 {
            return Ok(None);
        }
        Ok(Some(read_bytes(buf, &mut pos)?))
    }

    #[test]
    fn entry_len_request_round_trip() {
        let key = vec![1u8, 2, 3];
        let mut buf = Vec::new();
        encode_request_payload(&mut buf, RequestPayload::EntryLen { key: key.clone() });
        // The payload for EntryLen is: [u32 BE len][key bytes]
        let mut pos = 0;
        let decoded_key = read_bytes(&buf, &mut pos).expect("decode key");
        assert_eq!(decoded_key, key);
    }

    #[test]
    fn response_bool_round_trip() {
        for v in [true, false] {
            let buf = vec![v as u8];
            let got = decode_bool(&buf).expect("decode bool");
            assert_eq!(got, v);
        }
    }

    #[test]
    fn response_optional_value_round_trip() {
        for v in [None, Some(vec![]), Some(vec![10u8, 20, 30])] {
            let buf = encode_optional_value(v.as_deref());
            let got = decode_optional_value(&buf).expect("decode optional value");
            assert_eq!(got, v);
        }
    }

    #[test]
    fn range_payload_with_limit_round_trip() {
        let start = Bound::Included(vec![1, 2, 3]);
        let end = Bound::Excluded(vec![4, 5]);
        let limit = 42u32;

        let mut buf = Vec::new();
        encode_request_payload(
            &mut buf,
            RequestPayload::Range {
                start: start.clone(),
                end: end.clone(),
                limit,
            },
        );

        let mut pos = 0;
        let decoded_start = read_bound(&buf, &mut pos).expect("decode start");
        let decoded_end = read_bound(&buf, &mut pos).expect("decode end");
        let decoded_limit = read_u32_be(&buf, &mut pos).expect("decode limit");

        assert_eq!(decoded_start, start);
        assert_eq!(decoded_end, end);
        assert_eq!(decoded_limit, limit);
        assert_eq!(pos, buf.len());
    }
}

/// JSON-serialised schema response for [`OpCode::GetSchema`].
///
/// Available under the `schema` feature. Clients typically parse it out of
/// [`Response::Ok`] bytes via `serde_json::from_slice`.
#[cfg(feature = "schema")]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SchemaResponse {
    pub name: String,
    pub version: u16,
    pub typ: armour_core::Typ,
    pub key_scheme: armour_core::KeyScheme,
}