armour-rpc 0.2.1

use compio::buf::{IoBuf, IoBufMut};
use compio::io::framed::codec::{Decoder, Encoder};

use crate::error::RpcError;
use crate::protocol::*;

pub struct ClientCodec;

/// Encodes a Request into the wire format expected by the server:
/// `[op: u8][hashname: u64 BE][payload...]`
impl<B: IoBufMut> Encoder<Request, B> for ClientCodec {
    type Error = RpcError;

    fn encode(&mut self, item: Request, buf: &mut B) -> Result<(), Self::Error> {
        let mut tmp = Vec::new();
        tmp.push(item.op as u8);
        tmp.extend_from_slice(&item.hashname.to_be_bytes());
        encode_request_payload(&mut tmp, item.payload);
        buf.extend_from_slice(&tmp)
            .map_err(|e| RpcError::Other(e.to_string()))?;
        Ok(())
    }
}

/// Decodes a Response from the wire format sent by the server:
/// `[status: u8][payload...]`
///
/// The `Ok` variant contains the raw payload bytes; the caller decodes
/// them according to which operation was sent.
impl<B: IoBuf> Decoder<Response, B> for ClientCodec {
    type Error = RpcError;

    fn decode(&mut self, buf: &compio::buf::Slice<B>) -> Result<Response, Self::Error> {
        let bytes: &[u8] = buf;
        let mut pos = 0;
        let status = read_u8(bytes, &mut pos)?;
        match status {
            0x00 => Ok(Response::Ok(bytes[pos..].to_vec())),
            0x01 => {
                let code = read_u16_be(bytes, &mut pos)?;
                let msg_len = read_u16_be(bytes, &mut pos)? as usize;
                if pos + msg_len > bytes.len() {
                    return Err(RpcError::UnexpectedEof);
                }
                let message = String::from_utf8_lossy(&bytes[pos..pos + msg_len]).into_owned();
                Ok(Response::Err { code, message })
            }
            _ => Err(RpcError::Protocol(
                "invalid response status byte".to_string(),
            )),
        }
    }
}

pub(crate) fn encode_request_payload(buf: &mut Vec<u8>, payload: RequestPayload) {
    match payload {
        RequestPayload::Key(key) => {
            write_bytes(buf, &key);
        }
        RequestPayload::Empty => {}
        RequestPayload::Range { start, end, limit } => {
            write_bound(buf, &start);
            write_bound(buf, &end);
            buf.extend_from_slice(&limit.to_be_bytes());
        }
        RequestPayload::Upsert { key, flag, value } => {
            write_upsert_key(buf, &key);
            let flag_byte = match flag {
                None => 0u8,
                Some(true) => 1u8,
                Some(false) => 2u8,
            };
            buf.push(flag_byte);
            write_bytes(buf, &value);
        }
        RequestPayload::Remove { key, soft } => {
            write_bytes(buf, &key);
            buf.push(soft as u8);
        }
        RequestPayload::Take { key, soft } => {
            write_bytes(buf, &key);
            buf.push(soft as u8);
        }
        RequestPayload::Count { exact } => {
            buf.push(exact as u8);
        }
        RequestPayload::EntryLen { key } => {
            write_bytes(buf, &key);
        }
        RequestPayload::ListCollections => {}
        RequestPayload::Batch(items) => {
            buf.extend_from_slice(&(items.len() as u32).to_be_bytes());
            for (key, val) in items {
                write_bytes(buf, &key);
                match val {
                    None => buf.push(0),
                    Some(v) => {
                        buf.push(1);
                        write_bytes(buf, &v);
                    }
                }
            }
        }
    }
}