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            let bytes_len = values
207                .len()
208                .checked_mul(std::mem::size_of::<f32>())
209                .ok_or(ParamCodecError::LengthOverflow("vector"))?;
210            let start = out.len();
211            let end = start
212                .checked_add(bytes_len)
213                .ok_or(ParamCodecError::LengthOverflow("vector"))?;
214            out.resize(end, 0);
215            for (bytes, value) in out[start..].chunks_exact_mut(4).zip(values) {
216                bytes.copy_from_slice(&value.to_le_bytes());
217            }
218        }
219        ParamValue::Json(value) => {
220            out.push(TAG_JSON);
221            write_len_prefixed(value, out, "json")?;
222        }
223        ParamValue::Timestamp(value) => {
224            out.push(TAG_TIMESTAMP);
225            out.extend_from_slice(&value.to_le_bytes());
226        }
227        ParamValue::Uuid(value) => {
228            out.push(TAG_UUID);
229            out.extend_from_slice(value);
230        }
231    }
232    Ok(())
233}
234
235pub fn decode_value(payload: &[u8], pos: &mut usize) -> Result<ParamValue, ParamCodecError> {
236    let tag = *read_bytes(payload, pos, 1, "value tag")?
237        .first()
238        .expect("read one byte");
239    match tag {
240        TAG_NULL => Ok(ParamValue::Null),
241        TAG_BOOL => {
242            let value = read_bytes(payload, pos, 1, "bool")?[0];
243            match value {
244                0 => Ok(ParamValue::Bool(false)),
245                1 => Ok(ParamValue::Bool(true)),
246                other => Err(ParamCodecError::InvalidBool(other)),
247            }
248        }
249        TAG_INT => Ok(ParamValue::Int(read_i64(payload, pos, "int")?)),
250        TAG_FLOAT => Ok(ParamValue::Float(f64::from_le_bytes(read_array(
251            payload, pos, "float",
252        )?))),
253        TAG_TEXT => {
254            let len = read_u32(payload, pos, "text_len")? as usize;
255            let bytes = read_bytes(payload, pos, len, "text")?;
256            let text = std::str::from_utf8(bytes)
257                .map_err(|_| ParamCodecError::InvalidUtf8("text"))?
258                .to_string();
259            Ok(ParamValue::Text(text))
260        }
261        TAG_BYTES => {
262            let len = read_u32(payload, pos, "bytes_len")? as usize;
263            Ok(ParamValue::Bytes(
264                read_bytes(payload, pos, len, "bytes")?.to_vec(),
265            ))
266        }
267        TAG_VECTOR => {
268            let len = read_u32(payload, pos, "vector_len")? as usize;
269            let byte_len = len
270                .checked_mul(4)
271                .ok_or(ParamCodecError::LengthOverflow("vector"))?;
272            ensure_remaining(payload, *pos, byte_len, "vector")?;
273            let mut values = Vec::with_capacity(len);
274            for _ in 0..len {
275                values.push(f32::from_le_bytes(read_array(payload, pos, "vector")?));
276            }
277            Ok(ParamValue::Vector(values))
278        }
279        TAG_JSON => {
280            let len = read_u32(payload, pos, "json_len")? as usize;
281            Ok(ParamValue::Json(
282                read_bytes(payload, pos, len, "json")?.to_vec(),
283            ))
284        }
285        TAG_TIMESTAMP => Ok(ParamValue::Timestamp(read_i64(payload, pos, "timestamp")?)),
286        TAG_UUID => Ok(ParamValue::Uuid(read_array(payload, pos, "uuid")?)),
287        other => Err(ParamCodecError::UnknownTag(other)),
288    }
289}
290
291fn write_len_prefixed(
292    value: &[u8],
293    out: &mut Vec<u8>,
294    field: &'static str,
295) -> Result<(), ParamCodecError> {
296    if value.len() > u32::MAX as usize {
297        return Err(ParamCodecError::LengthOverflow(field));
298    }
299    out.extend_from_slice(&(value.len() as u32).to_le_bytes());
300    out.extend_from_slice(value);
301    Ok(())
302}
303
304fn read_u32(payload: &[u8], pos: &mut usize, field: &'static str) -> Result<u32, ParamCodecError> {
305    Ok(u32::from_le_bytes(read_array(payload, pos, field)?))
306}
307
308fn read_i64(payload: &[u8], pos: &mut usize, field: &'static str) -> Result<i64, ParamCodecError> {
309    Ok(i64::from_le_bytes(read_array(payload, pos, field)?))
310}
311
312fn read_array<const N: usize>(
313    payload: &[u8],
314    pos: &mut usize,
315    field: &'static str,
316) -> Result<[u8; N], ParamCodecError> {
317    let bytes = read_bytes(payload, pos, N, field)?;
318    let mut out = [0u8; N];
319    out.copy_from_slice(bytes);
320    Ok(out)
321}
322
323fn read_bytes<'a>(
324    payload: &'a [u8],
325    pos: &mut usize,
326    len: usize,
327    field: &'static str,
328) -> Result<&'a [u8], ParamCodecError> {
329    let end = pos
330        .checked_add(len)
331        .ok_or(ParamCodecError::Truncated(field))?;
332    if end > payload.len() {
333        return Err(ParamCodecError::Truncated(field));
334    }
335    let bytes = &payload[*pos..end];
336    *pos = end;
337    Ok(bytes)
338}
339
340fn ensure_remaining(
341    payload: &[u8],
342    pos: usize,
343    len: usize,
344    field: &'static str,
345) -> Result<(), ParamCodecError> {
346    let end = pos
347        .checked_add(len)
348        .ok_or(ParamCodecError::Truncated(field))?;
349    if end > payload.len() {
350        return Err(ParamCodecError::Truncated(field));
351    }
352    Ok(())
353}