Skip to main content

reddb_wire/
query_with_params.rs

1//! RedWire `QueryWithParams` payload codec.
2//!
3//! Payload layout v1:
4//! `u32 sql_len` + UTF-8 SQL + `u32 param_count` + encoded values.
5
6use std::fmt;
7
8pub const FEATURE_PARAMS: u32 = 0x0000_0001;
9pub const MAX_PARAM_COUNT: usize = 65_536;
10
11const OPTIONS_MAGIC: &[u8; 4] = b"RQP1";
12
13const TAG_NULL: u8 = 0x00;
14const TAG_BOOL: u8 = 0x01;
15const TAG_INT: u8 = 0x02;
16const TAG_FLOAT: u8 = 0x03;
17const TAG_TEXT: u8 = 0x04;
18const TAG_BYTES: u8 = 0x05;
19const TAG_VECTOR: u8 = 0x06;
20const TAG_JSON: u8 = 0x07;
21const TAG_TIMESTAMP: u8 = 0x08;
22const TAG_UUID: u8 = 0x09;
23
24#[derive(Debug, Clone, PartialEq)]
25pub enum ParamValue {
26    Null,
27    Bool(bool),
28    Int(i64),
29    Float(f64),
30    Text(String),
31    Bytes(Vec<u8>),
32    Vector(Vec<f32>),
33    Json(Vec<u8>),
34    Timestamp(i64),
35    Uuid([u8; 16]),
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub enum ParamCodecError {
40    LengthOverflow(&'static str),
41    ParamCountOverLimit(u32),
42    Truncated(&'static str),
43    InvalidUtf8(&'static str),
44    InvalidBool(u8),
45    UnknownTag(u8),
46    InvalidOptionsMagic,
47    TrailingBytes(usize),
48}
49
50impl fmt::Display for ParamCodecError {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        match self {
53            Self::LengthOverflow(field) => write!(f, "{field} is too large for RedWire v1"),
54            Self::ParamCountOverLimit(count) => {
55                write!(f, "param_count {count} exceeds RedWire v1 limit")
56            }
57            Self::Truncated(field) => write!(f, "truncated {field}"),
58            Self::InvalidUtf8(field) => write!(f, "{field} must be valid UTF-8"),
59            Self::InvalidBool(byte) => write!(f, "invalid bool payload byte {byte}"),
60            Self::UnknownTag(tag) => write!(f, "unknown parameter value tag 0x{tag:02x}"),
61            Self::InvalidOptionsMagic => write!(f, "invalid QueryWithParams options extension"),
62            Self::TrailingBytes(count) => write!(f, "{count} trailing bytes after payload"),
63        }
64    }
65}
66
67impl std::error::Error for ParamCodecError {}
68
69#[derive(Debug, Clone, PartialEq, Eq, Default)]
70pub struct QueryWithParamsOptions {
71    pub commit_policy: Option<String>,
72}
73
74#[derive(Debug, Clone, PartialEq)]
75pub struct QueryWithParamsRequest {
76    pub sql: String,
77    pub params: Vec<ParamValue>,
78    pub options: QueryWithParamsOptions,
79}
80
81pub fn encode_query_with_params(
82    sql: &str,
83    params: &[ParamValue],
84) -> Result<Vec<u8>, ParamCodecError> {
85    encode_query_with_params_request(sql, params, &QueryWithParamsOptions::default())
86}
87
88pub fn encode_query_with_params_request(
89    sql: &str,
90    params: &[ParamValue],
91    options: &QueryWithParamsOptions,
92) -> Result<Vec<u8>, ParamCodecError> {
93    if sql.len() > u32::MAX as usize {
94        return Err(ParamCodecError::LengthOverflow("sql"));
95    }
96    if params.len() > u32::MAX as usize {
97        return Err(ParamCodecError::LengthOverflow("params"));
98    }
99    if params.len() > MAX_PARAM_COUNT {
100        return Err(ParamCodecError::ParamCountOverLimit(params.len() as u32));
101    }
102
103    let mut out = Vec::new();
104    out.extend_from_slice(&(sql.len() as u32).to_le_bytes());
105    out.extend_from_slice(sql.as_bytes());
106    out.extend_from_slice(&(params.len() as u32).to_le_bytes());
107    for value in params {
108        encode_value(value, &mut out)?;
109    }
110    if let Some(policy) = &options.commit_policy {
111        if policy.len() > u32::MAX as usize {
112            return Err(ParamCodecError::LengthOverflow("commit_policy"));
113        }
114        out.extend_from_slice(OPTIONS_MAGIC);
115        write_len_prefixed(policy.as_bytes(), &mut out, "commit_policy")?;
116    }
117    Ok(out)
118}
119
120pub fn decode_query_with_params(
121    payload: &[u8],
122) -> Result<(String, Vec<ParamValue>), ParamCodecError> {
123    let request = decode_query_with_params_request(payload)?;
124    Ok((request.sql, request.params))
125}
126
127pub fn decode_query_with_params_request(
128    payload: &[u8],
129) -> Result<QueryWithParamsRequest, ParamCodecError> {
130    let mut pos = 0;
131    let sql_len = read_u32(payload, &mut pos, "sql_len")? as usize;
132    let sql_bytes = read_bytes(payload, &mut pos, sql_len, "sql")?;
133    let sql = std::str::from_utf8(sql_bytes)
134        .map_err(|_| ParamCodecError::InvalidUtf8("sql"))?
135        .to_string();
136    let param_count = read_u32(payload, &mut pos, "param_count")?;
137    if param_count as usize > MAX_PARAM_COUNT {
138        return Err(ParamCodecError::ParamCountOverLimit(param_count));
139    }
140    let mut params = Vec::with_capacity(param_count as usize);
141    for _ in 0..param_count {
142        params.push(decode_value(payload, &mut pos)?);
143    }
144    let options = if pos == payload.len() {
145        QueryWithParamsOptions::default()
146    } else {
147        decode_options(payload, &mut pos)?
148    };
149    if pos != payload.len() {
150        return Err(ParamCodecError::TrailingBytes(payload.len() - pos));
151    }
152    Ok(QueryWithParamsRequest {
153        sql,
154        params,
155        options,
156    })
157}
158
159fn decode_options(
160    payload: &[u8],
161    pos: &mut usize,
162) -> Result<QueryWithParamsOptions, ParamCodecError> {
163    let magic = read_bytes(payload, pos, OPTIONS_MAGIC.len(), "options_magic")?;
164    if magic != OPTIONS_MAGIC {
165        return Err(ParamCodecError::InvalidOptionsMagic);
166    }
167    let len = read_u32(payload, pos, "commit_policy_len")? as usize;
168    let bytes = read_bytes(payload, pos, len, "commit_policy")?;
169    let commit_policy = std::str::from_utf8(bytes)
170        .map_err(|_| ParamCodecError::InvalidUtf8("commit_policy"))?
171        .to_string();
172    Ok(QueryWithParamsOptions {
173        commit_policy: Some(commit_policy),
174    })
175}
176
177pub fn encode_value(value: &ParamValue, out: &mut Vec<u8>) -> Result<(), ParamCodecError> {
178    match value {
179        ParamValue::Null => out.push(TAG_NULL),
180        ParamValue::Bool(value) => {
181            out.push(TAG_BOOL);
182            out.push(u8::from(*value));
183        }
184        ParamValue::Int(value) => {
185            out.push(TAG_INT);
186            out.extend_from_slice(&value.to_le_bytes());
187        }
188        ParamValue::Float(value) => {
189            out.push(TAG_FLOAT);
190            out.extend_from_slice(&value.to_le_bytes());
191        }
192        ParamValue::Text(value) => {
193            out.push(TAG_TEXT);
194            write_len_prefixed(value.as_bytes(), out, "text")?;
195        }
196        ParamValue::Bytes(value) => {
197            out.push(TAG_BYTES);
198            write_len_prefixed(value, out, "bytes")?;
199        }
200        ParamValue::Vector(values) => {
201            out.push(TAG_VECTOR);
202            if values.len() > u32::MAX as usize {
203                return Err(ParamCodecError::LengthOverflow("vector"));
204            }
205            out.extend_from_slice(&(values.len() as u32).to_le_bytes());
206            for value in values {
207                out.extend_from_slice(&value.to_le_bytes());
208            }
209        }
210        ParamValue::Json(value) => {
211            out.push(TAG_JSON);
212            write_len_prefixed(value, out, "json")?;
213        }
214        ParamValue::Timestamp(value) => {
215            out.push(TAG_TIMESTAMP);
216            out.extend_from_slice(&value.to_le_bytes());
217        }
218        ParamValue::Uuid(value) => {
219            out.push(TAG_UUID);
220            out.extend_from_slice(value);
221        }
222    }
223    Ok(())
224}
225
226pub fn decode_value(payload: &[u8], pos: &mut usize) -> Result<ParamValue, ParamCodecError> {
227    let tag = *read_bytes(payload, pos, 1, "value tag")?
228        .first()
229        .expect("read one byte");
230    match tag {
231        TAG_NULL => Ok(ParamValue::Null),
232        TAG_BOOL => {
233            let value = read_bytes(payload, pos, 1, "bool")?[0];
234            match value {
235                0 => Ok(ParamValue::Bool(false)),
236                1 => Ok(ParamValue::Bool(true)),
237                other => Err(ParamCodecError::InvalidBool(other)),
238            }
239        }
240        TAG_INT => Ok(ParamValue::Int(read_i64(payload, pos, "int")?)),
241        TAG_FLOAT => Ok(ParamValue::Float(f64::from_le_bytes(read_array(
242            payload, pos, "float",
243        )?))),
244        TAG_TEXT => {
245            let len = read_u32(payload, pos, "text_len")? as usize;
246            let bytes = read_bytes(payload, pos, len, "text")?;
247            let text = std::str::from_utf8(bytes)
248                .map_err(|_| ParamCodecError::InvalidUtf8("text"))?
249                .to_string();
250            Ok(ParamValue::Text(text))
251        }
252        TAG_BYTES => {
253            let len = read_u32(payload, pos, "bytes_len")? as usize;
254            Ok(ParamValue::Bytes(
255                read_bytes(payload, pos, len, "bytes")?.to_vec(),
256            ))
257        }
258        TAG_VECTOR => {
259            let len = read_u32(payload, pos, "vector_len")? as usize;
260            let byte_len = len
261                .checked_mul(4)
262                .ok_or(ParamCodecError::LengthOverflow("vector"))?;
263            ensure_remaining(payload, *pos, byte_len, "vector")?;
264            let mut values = Vec::with_capacity(len);
265            for _ in 0..len {
266                values.push(f32::from_le_bytes(read_array(payload, pos, "vector")?));
267            }
268            Ok(ParamValue::Vector(values))
269        }
270        TAG_JSON => {
271            let len = read_u32(payload, pos, "json_len")? as usize;
272            Ok(ParamValue::Json(
273                read_bytes(payload, pos, len, "json")?.to_vec(),
274            ))
275        }
276        TAG_TIMESTAMP => Ok(ParamValue::Timestamp(read_i64(payload, pos, "timestamp")?)),
277        TAG_UUID => Ok(ParamValue::Uuid(read_array(payload, pos, "uuid")?)),
278        other => Err(ParamCodecError::UnknownTag(other)),
279    }
280}
281
282fn write_len_prefixed(
283    value: &[u8],
284    out: &mut Vec<u8>,
285    field: &'static str,
286) -> Result<(), ParamCodecError> {
287    if value.len() > u32::MAX as usize {
288        return Err(ParamCodecError::LengthOverflow(field));
289    }
290    out.extend_from_slice(&(value.len() as u32).to_le_bytes());
291    out.extend_from_slice(value);
292    Ok(())
293}
294
295fn read_u32(payload: &[u8], pos: &mut usize, field: &'static str) -> Result<u32, ParamCodecError> {
296    Ok(u32::from_le_bytes(read_array(payload, pos, field)?))
297}
298
299fn read_i64(payload: &[u8], pos: &mut usize, field: &'static str) -> Result<i64, ParamCodecError> {
300    Ok(i64::from_le_bytes(read_array(payload, pos, field)?))
301}
302
303fn read_array<const N: usize>(
304    payload: &[u8],
305    pos: &mut usize,
306    field: &'static str,
307) -> Result<[u8; N], ParamCodecError> {
308    let bytes = read_bytes(payload, pos, N, field)?;
309    let mut out = [0u8; N];
310    out.copy_from_slice(bytes);
311    Ok(out)
312}
313
314fn read_bytes<'a>(
315    payload: &'a [u8],
316    pos: &mut usize,
317    len: usize,
318    field: &'static str,
319) -> Result<&'a [u8], ParamCodecError> {
320    let end = pos
321        .checked_add(len)
322        .ok_or(ParamCodecError::Truncated(field))?;
323    if end > payload.len() {
324        return Err(ParamCodecError::Truncated(field));
325    }
326    let bytes = &payload[*pos..end];
327    *pos = end;
328    Ok(bytes)
329}
330
331fn ensure_remaining(
332    payload: &[u8],
333    pos: usize,
334    len: usize,
335    field: &'static str,
336) -> Result<(), ParamCodecError> {
337    let end = pos
338        .checked_add(len)
339        .ok_or(ParamCodecError::Truncated(field))?;
340    if end > payload.len() {
341        return Err(ParamCodecError::Truncated(field));
342    }
343    Ok(())
344}