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#[derive(Debug)]
72pub enum Response {
73 Ok(Vec<u8>),
74 Err { code: u16, message: String },
75}
76
77pub 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
134pub 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 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 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#[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}