Skip to main content

reddb_wire/redwire/
bulk_stream.rs

1//! RedWire legacy bulk-stream payload codec.
2//!
3//! Start layout:
4//! `[collection_len:u16][collection][column_count:u16]([column_len:u16][column_name])*`.
5//! Rows layout:
6//! `[row_count:u32]([legacy WireValue])*(column_count * row_count)`.
7
8use crate::legacy::{encode_value, try_decode_value, WireValue};
9use std::fmt;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub struct BulkStreamStartPayload {
13    pub collection: String,
14    pub columns: Vec<String>,
15}
16
17#[derive(Debug, Clone, PartialEq)]
18pub struct BulkStreamRowsPayload {
19    pub rows: Vec<Vec<WireValue>>,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum BulkStreamError {
24    MissingCollectionLength,
25    TruncatedCollectionName,
26    InvalidCollectionName,
27    MissingColumnCount,
28    MissingColumnNameLength,
29    TruncatedColumnName,
30    InvalidColumnName,
31    MissingRowCount,
32    Value(&'static str),
33    LengthOverflow(&'static str),
34    RowWidthMismatch { got: usize, expected: usize },
35    RowCountOverflow,
36}
37
38impl fmt::Display for BulkStreamError {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        match self {
41            Self::MissingCollectionLength => write!(f, "stream start: missing collection length"),
42            Self::TruncatedCollectionName => write!(f, "stream start: truncated collection name"),
43            Self::InvalidCollectionName => write!(f, "stream start: invalid collection name"),
44            Self::MissingColumnCount => write!(f, "stream start: missing column count"),
45            Self::MissingColumnNameLength => {
46                write!(f, "stream start: missing column name length")
47            }
48            Self::TruncatedColumnName => write!(f, "stream start: truncated column name"),
49            Self::InvalidColumnName => write!(f, "stream start: invalid column name"),
50            Self::MissingRowCount => write!(f, "stream rows: missing row count"),
51            Self::Value(err) => write!(f, "stream rows: {err}"),
52            Self::LengthOverflow(field) => {
53                write!(f, "{field} is too large for RedWire bulk stream")
54            }
55            Self::RowWidthMismatch { got, expected } => {
56                write!(f, "row had {got} values for {expected} columns")
57            }
58            Self::RowCountOverflow => write!(f, "row count is too large for RedWire bulk stream"),
59        }
60    }
61}
62
63impl std::error::Error for BulkStreamError {}
64
65pub fn encode_bulk_stream_start_payload(
66    collection: &str,
67    columns: &[&str],
68) -> Result<Vec<u8>, BulkStreamError> {
69    write_len_u16(collection.len(), "collection")?;
70    write_len_u16(columns.len(), "columns")?;
71    let mut out = Vec::with_capacity(4 + collection.len() + columns.len() * 16);
72    write_string_u16(&mut out, collection, "collection")?;
73    out.extend_from_slice(&(columns.len() as u16).to_le_bytes());
74    for column in columns {
75        write_string_u16(&mut out, column, "column")?;
76    }
77    Ok(out)
78}
79
80pub fn decode_bulk_stream_start_payload(
81    payload: &[u8],
82) -> Result<BulkStreamStartPayload, BulkStreamError> {
83    let mut pos = 0;
84    let coll_len = read_u16(payload, &mut pos, BulkStreamError::MissingCollectionLength)? as usize;
85    let collection = read_string(
86        payload,
87        &mut pos,
88        coll_len,
89        BulkStreamError::TruncatedCollectionName,
90        BulkStreamError::InvalidCollectionName,
91    )?;
92    let ncols = read_u16(payload, &mut pos, BulkStreamError::MissingColumnCount)? as usize;
93    let mut columns = Vec::with_capacity(ncols);
94    for _ in 0..ncols {
95        let len = read_u16(payload, &mut pos, BulkStreamError::MissingColumnNameLength)? as usize;
96        columns.push(read_string(
97            payload,
98            &mut pos,
99            len,
100            BulkStreamError::TruncatedColumnName,
101            BulkStreamError::InvalidColumnName,
102        )?);
103    }
104    Ok(BulkStreamStartPayload {
105        collection,
106        columns,
107    })
108}
109
110pub fn encode_bulk_stream_rows_payload(
111    rows: &[Vec<WireValue>],
112    column_count: usize,
113) -> Result<Vec<u8>, BulkStreamError> {
114    if rows.len() > u32::MAX as usize {
115        return Err(BulkStreamError::RowCountOverflow);
116    }
117    let mut out = Vec::with_capacity(4 + rows.len() * column_count * 16);
118    out.extend_from_slice(&(rows.len() as u32).to_le_bytes());
119    for row in rows {
120        if row.len() != column_count {
121            return Err(BulkStreamError::RowWidthMismatch {
122                got: row.len(),
123                expected: column_count,
124            });
125        }
126        for value in row {
127            encode_value(&mut out, value);
128        }
129    }
130    Ok(out)
131}
132
133pub fn decode_bulk_stream_rows_payload(
134    payload: &[u8],
135    column_count: usize,
136) -> Result<BulkStreamRowsPayload, BulkStreamError> {
137    let mut pos = 0;
138    let nrows = read_u32(payload, &mut pos, BulkStreamError::MissingRowCount)? as usize;
139    // A rows frame with no columns carries no data, and with `column_count`
140    // zero the value loop below consumes nothing — so nothing would ever hit
141    // the truncation error that normally terminates decoding, and the row
142    // loop would run up to `u32::MAX` times pushing empty vectors. Refuse
143    // the frame instead.
144    if column_count == 0 {
145        return Err(BulkStreamError::RowWidthMismatch {
146            got: 0,
147            expected: 1,
148        });
149    }
150    // Reject, rather than merely under-allocate for, a row count the payload
151    // cannot possibly hold: every value costs at least one byte, so a count
152    // above `remaining` is a malformed frame. Clamping only the capacity hint
153    // still let the loop itself run for hours on a few bytes of input.
154    let remaining = payload.len().saturating_sub(pos);
155    if nrows > remaining / column_count {
156        return Err(BulkStreamError::RowCountOverflow);
157    }
158    let mut rows = Vec::with_capacity(nrows);
159    for _ in 0..nrows {
160        let mut values = Vec::with_capacity(column_count);
161        for _ in 0..column_count {
162            values.push(try_decode_value(payload, &mut pos).map_err(BulkStreamError::Value)?);
163        }
164        rows.push(values);
165    }
166    Ok(BulkStreamRowsPayload { rows })
167}
168
169fn write_string_u16(
170    out: &mut Vec<u8>,
171    value: &str,
172    field: &'static str,
173) -> Result<(), BulkStreamError> {
174    write_len_u16(value.len(), field)?;
175    out.extend_from_slice(&(value.len() as u16).to_le_bytes());
176    out.extend_from_slice(value.as_bytes());
177    Ok(())
178}
179
180fn write_len_u16(len: usize, field: &'static str) -> Result<(), BulkStreamError> {
181    if len > u16::MAX as usize {
182        return Err(BulkStreamError::LengthOverflow(field));
183    }
184    Ok(())
185}
186
187fn read_u16(payload: &[u8], pos: &mut usize, err: BulkStreamError) -> Result<u16, BulkStreamError> {
188    let bytes = read_bytes(payload, pos, 2, err)?;
189    Ok(u16::from_le_bytes([bytes[0], bytes[1]]))
190}
191
192fn read_u32(payload: &[u8], pos: &mut usize, err: BulkStreamError) -> Result<u32, BulkStreamError> {
193    let bytes = read_bytes(payload, pos, 4, err)?;
194    Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
195}
196
197fn read_string(
198    payload: &[u8],
199    pos: &mut usize,
200    len: usize,
201    truncated_err: BulkStreamError,
202    utf8_err: BulkStreamError,
203) -> Result<String, BulkStreamError> {
204    let bytes = read_bytes(payload, pos, len, truncated_err)?;
205    std::str::from_utf8(bytes)
206        .map(str::to_owned)
207        .map_err(|_| utf8_err)
208}
209
210fn read_bytes<'a>(
211    payload: &'a [u8],
212    pos: &mut usize,
213    len: usize,
214    err: BulkStreamError,
215) -> Result<&'a [u8], BulkStreamError> {
216    let end = pos.saturating_add(len);
217    if end > payload.len() {
218        return Err(err);
219    }
220    let bytes = &payload[*pos..end];
221    *pos = end;
222    Ok(bytes)
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn stream_start_payload_round_trips() {
231        let bytes = encode_bulk_stream_start_payload("events", &["id", "name"]).unwrap();
232        let decoded = decode_bulk_stream_start_payload(&bytes).unwrap();
233        assert_eq!(decoded.collection, "events");
234        assert_eq!(decoded.columns, vec!["id", "name"]);
235    }
236
237    #[test]
238    fn stream_rows_payload_round_trips_values() {
239        let rows = vec![vec![WireValue::I64(7), WireValue::Text("Ada".into())]];
240        let bytes = encode_bulk_stream_rows_payload(&rows, 2).unwrap();
241        assert_eq!(
242            decode_bulk_stream_rows_payload(&bytes, 2).unwrap().rows,
243            rows
244        );
245    }
246
247    #[test]
248    fn stream_rows_payload_preserves_error_prefix() {
249        let payload = vec![1, 0, 0, 0, 1];
250        assert_eq!(
251            decode_bulk_stream_rows_payload(&payload, 1)
252                .unwrap_err()
253                .to_string(),
254            "stream rows: truncated i64 value"
255        );
256    }
257
258    #[test]
259    fn hostile_row_count_is_rejected_not_merely_under_allocated() {
260        // A 4-byte payload claiming u32::MAX rows used to request a ~96 GiB
261        // `Vec`, and an allocation failure aborts the process rather than
262        // erroring the connection. Clamping only the capacity hint was not
263        // enough: the row loop still ran, so the count itself must be
264        // rejected against what the payload can hold.
265        let mut payload = Vec::new();
266        payload.extend_from_slice(&u32::MAX.to_le_bytes());
267        let err = decode_bulk_stream_rows_payload(&payload, 4)
268            .expect_err("a row count with no rows behind it must be refused");
269        assert!(
270            matches!(err, BulkStreamError::RowCountOverflow),
271            "unexpected error: {err:?}"
272        );
273    }
274
275    #[test]
276    fn zero_column_rows_frame_is_refused() {
277        // Found by the `redwire_bulk` fuzz target: with no columns the value
278        // loop consumes nothing, so the truncation error that normally ends
279        // decoding never fires and the row loop runs up to u32::MAX times
280        // pushing empty vectors — gigabytes of allocation from four bytes.
281        let mut payload = Vec::new();
282        payload.extend_from_slice(&0x0a7f_007fu32.to_le_bytes());
283        let err = decode_bulk_stream_rows_payload(&payload, 0)
284            .expect_err("a rows frame with no columns carries no data");
285        assert!(
286            matches!(err, BulkStreamError::RowWidthMismatch { .. }),
287            "unexpected error: {err:?}"
288        );
289    }
290}