1use 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 if column_count == 0 {
145 return Err(BulkStreamError::RowWidthMismatch {
146 got: 0,
147 expected: 1,
148 });
149 }
150 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 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 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}