Skip to main content

armour_rpc/
protocol.rs

1use crate::error::RpcError;
2use std::ops::Bound;
3use strum::FromRepr;
4
5#[repr(u8)]
6#[derive(Debug, Clone, Copy, PartialEq, Eq, FromRepr)]
7pub enum OpCode {
8    Get = 0x01,
9    Contains = 0x02,
10    First = 0x03,
11    Last = 0x04,
12    Range = 0x05,
13    RangeKeys = 0x06,
14    Count = 0x07,
15    EntryLen = 0x08,
16    Upsert = 0x10,
17    Remove = 0x11,
18    Take = 0x12,
19    ApplyBatch = 0x13,
20    ListCollections = 0x20,
21    GetSchema = 0x21,
22}
23
24#[derive(Debug, Clone)]
25pub enum UpsertKey {
26    Sequence,
27    Provided(Vec<u8>),
28}
29
30#[derive(Debug)]
31pub struct Request {
32    pub op: OpCode,
33    pub hashname: u64,
34    pub payload: RequestPayload,
35}
36
37#[derive(Debug)]
38pub enum RequestPayload {
39    Key(Vec<u8>),
40    Empty,
41    Range {
42        start: Bound<Vec<u8>>,
43        end: Bound<Vec<u8>>,
44        limit: u32,
45    },
46    Upsert {
47        key: UpsertKey,
48        flag: Option<bool>,
49        value: Vec<u8>,
50    },
51    Remove {
52        key: Vec<u8>,
53        soft: bool,
54    },
55    Take {
56        key: Vec<u8>,
57        soft: bool,
58    },
59    Count {
60        exact: bool,
61    },
62    EntryLen {
63        key: Vec<u8>,
64    },
65    Batch(Vec<(Vec<u8>, Option<Vec<u8>>)>),
66    ListCollections,
67}
68
69/// Response decoded from the wire.
70/// `Ok` contains the raw payload bytes; each client method decodes them.
71#[derive(Debug)]
72pub enum Response {
73    Ok(Vec<u8>),
74    Err { code: u16, message: String },
75}
76
77// --- Encoding helpers ---
78
79pub fn write_bound(buf: &mut Vec<u8>, bound: &Bound<Vec<u8>>) {
80    match bound {
81        Bound::Unbounded => buf.push(0x00),
82        Bound::Included(bytes) => {
83            buf.push(0x01);
84            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
85            buf.extend_from_slice(bytes);
86        }
87        Bound::Excluded(bytes) => {
88            buf.push(0x02);
89            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
90            buf.extend_from_slice(bytes);
91        }
92    }
93}
94
95pub fn read_bound(buf: &[u8], pos: &mut usize) -> Result<Bound<Vec<u8>>, RpcError> {
96    let tag = read_u8(buf, pos)?;
97    match tag {
98        0x00 => Ok(Bound::Unbounded),
99        0x01 => {
100            let bytes = read_bytes(buf, pos)?;
101            Ok(Bound::Included(bytes))
102        }
103        0x02 => {
104            let bytes = read_bytes(buf, pos)?;
105            Ok(Bound::Excluded(bytes))
106        }
107        _ => Err(RpcError::Protocol("invalid bound tag".to_string())),
108    }
109}
110
111pub fn read_upsert_key(buf: &[u8], pos: &mut usize) -> Result<UpsertKey, RpcError> {
112    let tag = read_u8(buf, pos)?;
113    match tag {
114        0x00 => Ok(UpsertKey::Sequence),
115        0x01 => {
116            let key = read_bytes(buf, pos)?;
117            Ok(UpsertKey::Provided(key))
118        }
119        _ => Err(RpcError::Protocol("invalid upsert key tag".to_string())),
120    }
121}
122
123pub fn write_upsert_key(buf: &mut Vec<u8>, key: &UpsertKey) {
124    match key {
125        UpsertKey::Sequence => buf.push(0x00),
126        UpsertKey::Provided(bytes) => {
127            buf.push(0x01);
128            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
129            buf.extend_from_slice(bytes);
130        }
131    }
132}
133
134// --- Primitive read/write helpers ---
135
136pub fn read_u8(buf: &[u8], pos: &mut usize) -> Result<u8, RpcError> {
137    if *pos >= buf.len() {
138        return Err(RpcError::UnexpectedEof);
139    }
140    let v = buf[*pos];
141    *pos += 1;
142    Ok(v)
143}
144
145pub fn read_u16_be(buf: &[u8], pos: &mut usize) -> Result<u16, RpcError> {
146    if *pos + 2 > buf.len() {
147        return Err(RpcError::UnexpectedEof);
148    }
149    let v = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]);
150    *pos += 2;
151    Ok(v)
152}
153
154pub fn read_u32_be(buf: &[u8], pos: &mut usize) -> Result<u32, RpcError> {
155    if *pos + 4 > buf.len() {
156        return Err(RpcError::UnexpectedEof);
157    }
158    let v = u32::from_be_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
159    *pos += 4;
160    Ok(v)
161}
162
163pub fn read_u64_be(buf: &[u8], pos: &mut usize) -> Result<u64, RpcError> {
164    if *pos + 8 > buf.len() {
165        return Err(RpcError::UnexpectedEof);
166    }
167    let v = u64::from_be_bytes([
168        buf[*pos],
169        buf[*pos + 1],
170        buf[*pos + 2],
171        buf[*pos + 3],
172        buf[*pos + 4],
173        buf[*pos + 5],
174        buf[*pos + 6],
175        buf[*pos + 7],
176    ]);
177    *pos += 8;
178    Ok(v)
179}
180
181pub fn read_bytes(buf: &[u8], pos: &mut usize) -> Result<Vec<u8>, RpcError> {
182    let len = read_u32_be(buf, pos)? as usize;
183    if *pos + len > buf.len() {
184        return Err(RpcError::UnexpectedEof);
185    }
186    let v = buf[*pos..*pos + len].to_vec();
187    *pos += len;
188    Ok(v)
189}
190
191pub fn write_bytes(buf: &mut Vec<u8>, bytes: &[u8]) {
192    buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
193    buf.extend_from_slice(bytes);
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use crate::client::decode_bool;
200    use crate::codec::encode_request_payload;
201
202    // Helpers that mirror the server-side encoding used in response decoders.
203
204    fn encode_optional_value(v: Option<&[u8]>) -> Vec<u8> {
205        let mut buf = Vec::new();
206        match v {
207            None => buf.push(0u8),
208            Some(bytes) => {
209                buf.push(1u8);
210                write_bytes(&mut buf, bytes);
211            }
212        }
213        buf
214    }
215
216    fn decode_optional_value(buf: &[u8]) -> Result<Option<Vec<u8>>, RpcError> {
217        let mut pos = 0;
218        let flag = read_u8(buf, &mut pos)?;
219        if flag == 0 {
220            return Ok(None);
221        }
222        Ok(Some(read_bytes(buf, &mut pos)?))
223    }
224
225    #[test]
226    fn entry_len_request_round_trip() {
227        let key = vec![1u8, 2, 3];
228        let mut buf = Vec::new();
229        encode_request_payload(&mut buf, RequestPayload::EntryLen { key: key.clone() });
230        // The payload for EntryLen is: [u32 BE len][key bytes]
231        let mut pos = 0;
232        let decoded_key = read_bytes(&buf, &mut pos).expect("decode key");
233        assert_eq!(decoded_key, key);
234    }
235
236    #[test]
237    fn response_bool_round_trip() {
238        for v in [true, false] {
239            let buf = vec![v as u8];
240            let got = decode_bool(&buf).expect("decode bool");
241            assert_eq!(got, v);
242        }
243    }
244
245    #[test]
246    fn response_optional_value_round_trip() {
247        for v in [None, Some(vec![]), Some(vec![10u8, 20, 30])] {
248            let buf = encode_optional_value(v.as_deref());
249            let got = decode_optional_value(&buf).expect("decode optional value");
250            assert_eq!(got, v);
251        }
252    }
253
254    #[test]
255    fn range_payload_with_limit_round_trip() {
256        let start = Bound::Included(vec![1, 2, 3]);
257        let end = Bound::Excluded(vec![4, 5]);
258        let limit = 42u32;
259
260        let mut buf = Vec::new();
261        encode_request_payload(
262            &mut buf,
263            RequestPayload::Range {
264                start: start.clone(),
265                end: end.clone(),
266                limit,
267            },
268        );
269
270        let mut pos = 0;
271        let decoded_start = read_bound(&buf, &mut pos).expect("decode start");
272        let decoded_end = read_bound(&buf, &mut pos).expect("decode end");
273        let decoded_limit = read_u32_be(&buf, &mut pos).expect("decode limit");
274
275        assert_eq!(decoded_start, start);
276        assert_eq!(decoded_end, end);
277        assert_eq!(decoded_limit, limit);
278        assert_eq!(pos, buf.len());
279    }
280}
281
282/// JSON-serialised schema response for [`OpCode::GetSchema`].
283///
284/// Available under the `schema` feature. Clients typically parse it out of
285/// [`Response::Ok`] bytes via `serde_json::from_slice`.
286#[cfg(feature = "schema")]
287#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
288pub struct SchemaResponse {
289    pub name: String,
290    pub version: u16,
291    pub typ: armour_core::Typ,
292    pub key_scheme: armour_core::KeyScheme,
293}