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    Upsert = 0x10,
16    Remove = 0x11,
17    Take = 0x12,
18    ApplyBatch = 0x13,
19    ListCollections = 0x20,
20}
21
22#[derive(Debug, Clone)]
23pub enum UpsertKey {
24    Sequence,
25    Provided(Vec<u8>),
26}
27
28#[derive(Debug)]
29pub struct Request {
30    pub op: OpCode,
31    pub hashname: u64,
32    pub payload: RequestPayload,
33}
34
35#[derive(Debug)]
36pub enum RequestPayload {
37    Key(Vec<u8>),
38    Empty,
39    Range {
40        start: Bound<Vec<u8>>,
41        end: Bound<Vec<u8>>,
42    },
43    Upsert {
44        key: UpsertKey,
45        flag: Option<bool>,
46        value: Vec<u8>,
47    },
48    Remove {
49        key: Vec<u8>,
50        soft: bool,
51    },
52    Take {
53        key: Vec<u8>,
54        soft: bool,
55    },
56    Count {
57        exact: bool,
58    },
59    Batch(Vec<(Vec<u8>, Option<Vec<u8>>)>),
60    ListCollections,
61}
62
63/// Response decoded from the wire.
64/// `Ok` contains the raw payload bytes; each client method decodes them.
65#[derive(Debug)]
66pub enum Response {
67    Ok(Vec<u8>),
68    Err { code: u16, message: String },
69}
70
71// --- Encoding helpers ---
72
73pub fn write_bound(buf: &mut Vec<u8>, bound: &Bound<Vec<u8>>) {
74    match bound {
75        Bound::Unbounded => buf.push(0x00),
76        Bound::Included(bytes) => {
77            buf.push(0x01);
78            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
79            buf.extend_from_slice(bytes);
80        }
81        Bound::Excluded(bytes) => {
82            buf.push(0x02);
83            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
84            buf.extend_from_slice(bytes);
85        }
86    }
87}
88
89pub fn read_bound(buf: &[u8], pos: &mut usize) -> Result<Bound<Vec<u8>>, RpcError> {
90    let tag = read_u8(buf, pos)?;
91    match tag {
92        0x00 => Ok(Bound::Unbounded),
93        0x01 => {
94            let bytes = read_bytes(buf, pos)?;
95            Ok(Bound::Included(bytes))
96        }
97        0x02 => {
98            let bytes = read_bytes(buf, pos)?;
99            Ok(Bound::Excluded(bytes))
100        }
101        _ => Err(RpcError::Protocol("invalid bound tag".to_string())),
102    }
103}
104
105pub fn read_upsert_key(buf: &[u8], pos: &mut usize) -> Result<UpsertKey, RpcError> {
106    let tag = read_u8(buf, pos)?;
107    match tag {
108        0x00 => Ok(UpsertKey::Sequence),
109        0x01 => {
110            let key = read_bytes(buf, pos)?;
111            Ok(UpsertKey::Provided(key))
112        }
113        _ => Err(RpcError::Protocol("invalid upsert key tag".to_string())),
114    }
115}
116
117pub fn write_upsert_key(buf: &mut Vec<u8>, key: &UpsertKey) {
118    match key {
119        UpsertKey::Sequence => buf.push(0x00),
120        UpsertKey::Provided(bytes) => {
121            buf.push(0x01);
122            buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
123            buf.extend_from_slice(bytes);
124        }
125    }
126}
127
128// --- Primitive read/write helpers ---
129
130pub fn read_u8(buf: &[u8], pos: &mut usize) -> Result<u8, RpcError> {
131    if *pos >= buf.len() {
132        return Err(RpcError::UnexpectedEof);
133    }
134    let v = buf[*pos];
135    *pos += 1;
136    Ok(v)
137}
138
139pub fn read_u16_be(buf: &[u8], pos: &mut usize) -> Result<u16, RpcError> {
140    if *pos + 2 > buf.len() {
141        return Err(RpcError::UnexpectedEof);
142    }
143    let v = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]);
144    *pos += 2;
145    Ok(v)
146}
147
148pub fn read_u32_be(buf: &[u8], pos: &mut usize) -> Result<u32, RpcError> {
149    if *pos + 4 > buf.len() {
150        return Err(RpcError::UnexpectedEof);
151    }
152    let v = u32::from_be_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
153    *pos += 4;
154    Ok(v)
155}
156
157pub fn read_u64_be(buf: &[u8], pos: &mut usize) -> Result<u64, RpcError> {
158    if *pos + 8 > buf.len() {
159        return Err(RpcError::UnexpectedEof);
160    }
161    let v = u64::from_be_bytes([
162        buf[*pos],
163        buf[*pos + 1],
164        buf[*pos + 2],
165        buf[*pos + 3],
166        buf[*pos + 4],
167        buf[*pos + 5],
168        buf[*pos + 6],
169        buf[*pos + 7],
170    ]);
171    *pos += 8;
172    Ok(v)
173}
174
175pub fn read_bytes(buf: &[u8], pos: &mut usize) -> Result<Vec<u8>, RpcError> {
176    let len = read_u32_be(buf, pos)? as usize;
177    if *pos + len > buf.len() {
178        return Err(RpcError::UnexpectedEof);
179    }
180    let v = buf[*pos..*pos + len].to_vec();
181    *pos += len;
182    Ok(v)
183}
184
185pub fn write_bytes(buf: &mut Vec<u8>, bytes: &[u8]) {
186    buf.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
187    buf.extend_from_slice(bytes);
188}