1use std::mem::size_of;
13
14use nodedb_types::decode_bounds::checked_decode_capacity;
15use serde::{Deserialize, Serialize};
16use sonic_rs;
17use zerompk::{FromMessagePack, ToMessagePack};
18
19#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
21#[serde(rename_all = "snake_case")]
22#[non_exhaustive]
23pub enum ColumnarWalRecord {
24 #[serde(rename = "insert_row")]
30 InsertRow {
31 collection: String,
32 row_data: Vec<u8>,
36 },
37
38 #[serde(rename = "delete_rows")]
43 DeleteRows {
44 collection: String,
45 segment_id: u64,
46 row_indices: Vec<u32>,
47 },
48
49 #[serde(rename = "memtable_flushed")]
55 MemtableFlushed {
56 collection: String,
57 segment_id: u64,
58 row_count: u64,
59 },
60}
61
62impl ColumnarWalRecord {
63 pub fn collection(&self) -> &str {
65 match self {
66 Self::InsertRow { collection, .. }
67 | Self::DeleteRows { collection, .. }
68 | Self::MemtableFlushed { collection, .. } => collection,
69 }
70 }
71
72 pub fn to_bytes(&self) -> Result<Vec<u8>, crate::error::ColumnarError> {
74 zerompk::to_msgpack_vec(self)
75 .map_err(|e| crate::error::ColumnarError::Serialization(e.to_string()))
76 }
77
78 pub fn from_bytes(data: &[u8]) -> Result<Self, crate::error::ColumnarError> {
80 zerompk::from_msgpack(data)
81 .map_err(|e| crate::error::ColumnarError::Serialization(e.to_string()))
82 }
83}
84
85pub fn encode_row_for_wal(
91 values: &[nodedb_types::value::Value],
92) -> Result<Vec<u8>, crate::error::ColumnarError> {
93 use nodedb_types::value::Value;
94
95 let mut buf = Vec::with_capacity(values.len() * 10); for value in values {
98 match value {
99 Value::Null => buf.push(0),
100 Value::Integer(v) => {
101 buf.push(1);
102 buf.extend_from_slice(&v.to_le_bytes());
103 }
104 Value::Float(v) => {
105 buf.push(2);
106 buf.extend_from_slice(&v.to_le_bytes());
107 }
108 Value::Bool(v) => {
109 buf.push(3);
110 buf.push(*v as u8);
111 }
112 Value::String(s) => {
113 buf.push(4);
114 let bytes = s.as_bytes();
115 buf.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
116 buf.extend_from_slice(bytes);
117 }
118 Value::Bytes(b) => {
119 buf.push(5);
120 buf.extend_from_slice(&(b.len() as u32).to_le_bytes());
121 buf.extend_from_slice(b);
122 }
123 Value::DateTime(dt) => {
124 buf.push(6);
125 buf.extend_from_slice(&dt.micros.to_le_bytes());
126 }
127 Value::Decimal(d) => {
128 buf.push(7);
129 buf.extend_from_slice(&d.serialize());
130 }
131 Value::Uuid(s) => {
132 buf.push(8);
133 let bytes = s.as_bytes();
134 buf.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
135 buf.extend_from_slice(bytes);
136 }
137 Value::Array(arr) if arr.iter().all(|v| matches!(v, Value::Float(_))) => {
138 buf.push(9);
141 buf.extend_from_slice(&(arr.len() as u32).to_le_bytes());
142 for v in arr {
143 if let Value::Float(f) = v {
144 buf.extend_from_slice(&f.to_le_bytes());
145 }
146 }
147 }
148 _ => {
149 buf.push(10);
152 let json = sonic_rs::to_vec(value).map_err(|e| {
153 crate::error::ColumnarError::Serialization(format!(
154 "failed to serialize value as JSON: {e}"
155 ))
156 })?;
157 buf.extend_from_slice(&(json.len() as u32).to_le_bytes());
158 buf.extend_from_slice(&json);
159 }
160 }
161 }
162
163 Ok(buf)
164}
165
166const MAX_FIELD_LEN: usize = 256 * 1024 * 1024;
169
170fn corrupt(offset: usize, reason: impl Into<String>) -> crate::error::ColumnarError {
172 crate::error::ColumnarError::WalRowCorrupt {
173 offset,
174 reason: reason.into(),
175 }
176}
177
178fn read_array<const N: usize>(
181 data: &[u8],
182 cursor: &mut usize,
183 context: &str,
184) -> Result<[u8; N], crate::error::ColumnarError> {
185 let at = *cursor;
186 let slice = read_slice(data, cursor, N, context)?;
187 slice
188 .try_into()
189 .map_err(|_| corrupt(at, format!("truncated {context}")))
190}
191
192fn read_slice<'a>(
195 data: &'a [u8],
196 cursor: &mut usize,
197 n: usize,
198 context: &str,
199) -> Result<&'a [u8], crate::error::ColumnarError> {
200 let end = cursor
201 .checked_add(n)
202 .ok_or_else(|| corrupt(*cursor, format!("overflow in {context}")))?;
203 if end > data.len() {
204 return Err(corrupt(
205 *cursor,
206 format!(
207 "truncated {context}: need {n} bytes, have {}",
208 data.len().saturating_sub(*cursor)
209 ),
210 ));
211 }
212 let slice = &data[*cursor..end];
213 *cursor = end;
214 Ok(slice)
215}
216
217fn read_length_prefixed<'a>(
220 data: &'a [u8],
221 cursor: &mut usize,
222 context: &str,
223) -> Result<&'a [u8], crate::error::ColumnarError> {
224 let at = *cursor;
225 let len = u32::from_le_bytes(read_array::<4>(data, cursor, context)?) as usize;
226 if len > MAX_FIELD_LEN {
227 return Err(corrupt(
228 at,
229 format!("{context} length {len} exceeds maximum {MAX_FIELD_LEN}"),
230 ));
231 }
232 read_slice(data, cursor, len, context)
233}
234
235fn read_utf8(
238 data: &[u8],
239 cursor: &mut usize,
240 context: &str,
241) -> Result<String, crate::error::ColumnarError> {
242 let at = *cursor;
243 let bytes = read_length_prefixed(data, cursor, context)?;
244 String::from_utf8(bytes.to_vec())
245 .map_err(|e| corrupt(at, format!("{context} is not UTF-8: {e}")))
246}
247
248pub fn decode_row_from_wal(
253 data: &[u8],
254) -> Result<Vec<nodedb_types::value::Value>, crate::error::ColumnarError> {
255 decode_row(data).inspect_err(crate::diag::wal_row_corrupt)
256}
257
258fn decode_row(data: &[u8]) -> Result<Vec<nodedb_types::value::Value>, crate::error::ColumnarError> {
260 use nodedb_types::value::Value;
261
262 let mut values = Vec::new();
263 let mut cursor = 0;
264
265 while cursor < data.len() {
266 let tag_at = cursor;
267 let [tag] = read_array::<1>(data, &mut cursor, "tag")?;
268
269 let value = match tag {
270 0 => Value::Null,
271 1 => Value::Integer(i64::from_le_bytes(read_array(data, &mut cursor, "i64")?)),
272 2 => Value::Float(f64::from_le_bytes(read_array(data, &mut cursor, "f64")?)),
273 3 => {
274 let [b] = read_array::<1>(data, &mut cursor, "bool")?;
275 Value::Bool(b != 0)
276 }
277 4 => Value::String(read_utf8(data, &mut cursor, "string")?),
278 5 => Value::Bytes(read_length_prefixed(data, &mut cursor, "bytes")?.to_vec()),
279 8 => Value::Uuid(read_utf8(data, &mut cursor, "uuid")?),
280 6 => {
281 let micros = i64::from_le_bytes(read_array(data, &mut cursor, "timestamp")?);
282 Value::DateTime(nodedb_types::datetime::NdbDateTime::from_micros(micros))
283 }
284 7 => Value::Decimal(rust_decimal::Decimal::deserialize(read_array(
285 data,
286 &mut cursor,
287 "decimal",
288 )?)),
289 9 => {
290 let count_at = cursor;
291 let count =
292 u32::from_le_bytes(read_array(data, &mut cursor, "vector count")?) as usize;
293 let remaining_values = data.len().saturating_sub(cursor) / 8;
294 let max_count = (MAX_FIELD_LEN / 8).min(remaining_values);
295 if count > max_count {
296 return Err(corrupt(
297 count_at,
298 format!("vector count {count} exceeds maximum {max_count}"),
299 ));
300 }
301 let capacity = checked_decode_capacity(
302 count,
303 size_of::<nodedb_types::value::Value>(),
304 data.len().saturating_sub(cursor),
305 8,
306 max_count,
307 usize::MAX,
308 )
309 .ok_or_else(|| {
310 corrupt(count_at, "vector count exceeds decode allocation bounds")
311 })?;
312 let mut arr = Vec::with_capacity(capacity);
313 for _ in 0..count {
314 let f = f64::from_le_bytes(read_array(data, &mut cursor, "vector f64")?);
315 arr.push(Value::Float(f));
316 }
317 Value::Array(arr)
318 }
319 10 => {
320 let json_at = cursor;
321 let json_bytes = read_length_prefixed(data, &mut cursor, "json")?;
322 sonic_rs::from_slice(json_bytes)
323 .map_err(|e| corrupt(json_at, format!("json value does not decode: {e}")))?
324 }
325 _ => return Err(corrupt(tag_at, format!("unknown WAL value tag: {tag}"))),
326 };
327
328 values.push(value);
329 }
330
331 Ok(values)
332}
333
334#[cfg(test)]
335mod tests {
336 use nodedb_types::datetime::NdbDateTime;
337 use nodedb_types::value::Value;
338
339 use super::*;
340
341 #[test]
342 fn rejects_huge_vector_count_with_tiny_payload_before_allocation() {
343 let mut bytes = vec![9];
345 bytes.extend_from_slice(&u32::MAX.to_le_bytes());
346 assert!(matches!(
347 decode_row_from_wal(&bytes),
348 Err(crate::error::ColumnarError::WalRowCorrupt { .. })
349 ));
350 }
351
352 #[test]
353 fn a_string_that_is_not_utf8_is_refused() {
354 let mut bytes = vec![4];
355 bytes.extend_from_slice(&2u32.to_le_bytes());
356 bytes.extend_from_slice(&[0xC3, 0x28]);
357 assert!(matches!(
358 decode_row_from_wal(&bytes),
359 Err(crate::error::ColumnarError::WalRowCorrupt { offset: 1, .. })
360 ));
361 }
362
363 #[test]
364 fn a_uuid_that_is_not_utf8_is_refused() {
365 let mut bytes = vec![8];
366 bytes.extend_from_slice(&1u32.to_le_bytes());
367 bytes.push(0xFF);
368 assert!(matches!(
369 decode_row_from_wal(&bytes),
370 Err(crate::error::ColumnarError::WalRowCorrupt { .. })
371 ));
372 }
373
374 #[test]
375 fn a_json_value_that_does_not_decode_is_refused() {
376 let json = b"{not json";
377 let mut bytes = vec![10];
378 bytes.extend_from_slice(&(json.len() as u32).to_le_bytes());
379 bytes.extend_from_slice(json);
380 assert!(matches!(
381 decode_row_from_wal(&bytes),
382 Err(crate::error::ColumnarError::WalRowCorrupt { .. })
383 ));
384 }
385
386 #[test]
387 fn arrays_decode_to_the_values_they_were() {
388 let values = vec![
389 Value::Array(vec![Value::Float(0.1), Value::Float(-2.5)]),
390 Value::Array(vec![Value::String("a".into()), Value::String("b".into())]),
391 ];
392 let encoded = encode_row_for_wal(&values).expect("encode");
393 assert_eq!(decode_row_from_wal(&encoded).expect("decode"), values);
394 }
395
396 #[test]
397 fn wal_record_roundtrip() {
398 let records = vec![
399 ColumnarWalRecord::InsertRow {
400 collection: "test".into(),
401 row_data: vec![1, 2, 3],
402 },
403 ColumnarWalRecord::DeleteRows {
404 collection: "test".into(),
405 segment_id: 0,
406 row_indices: vec![5, 10, 15],
407 },
408 ColumnarWalRecord::MemtableFlushed {
409 collection: "test".into(),
410 segment_id: 3,
411 row_count: 1024,
412 },
413 ];
414
415 for record in &records {
416 let bytes = record.to_bytes().expect("serialize");
417 let restored = ColumnarWalRecord::from_bytes(&bytes).expect("deserialize");
418 assert_eq!(restored.collection(), record.collection());
419 }
420 }
421
422 #[test]
423 fn row_wire_format_roundtrip() {
424 let values = vec![
425 Value::Integer(42),
426 Value::Float(0.75),
427 Value::Bool(true),
428 Value::String("hello".into()),
429 Value::Bytes(vec![0xDE, 0xAD]),
430 Value::DateTime(NdbDateTime::from_micros(1_700_000_000)),
431 Value::Decimal(rust_decimal::Decimal::new(314, 2)),
432 Value::Uuid("550e8400-e29b-41d4-a716-446655440000".into()),
433 Value::Null,
434 Value::Array(vec![Value::Float(1.0), Value::Float(2.0)]),
435 ];
436
437 let encoded = encode_row_for_wal(&values).expect("encode");
438 let decoded = decode_row_from_wal(&encoded).expect("decode");
439
440 assert_eq!(decoded.len(), values.len());
441 assert_eq!(decoded[0], Value::Integer(42));
442 assert_eq!(decoded[1], Value::Float(0.75));
443 assert_eq!(decoded[2], Value::Bool(true));
444 assert_eq!(decoded[3], Value::String("hello".into()));
445 assert_eq!(decoded[4], Value::Bytes(vec![0xDE, 0xAD]));
446 assert_eq!(
447 decoded[5],
448 Value::DateTime(NdbDateTime::from_micros(1_700_000_000))
449 );
450 assert_eq!(
451 decoded[7],
452 Value::Uuid("550e8400-e29b-41d4-a716-446655440000".into())
453 );
454 assert_eq!(decoded[8], Value::Null);
455 }
456
457 #[test]
458 fn decode_truncated_i64_returns_error() {
459 let result = decode_row_from_wal(&[1]);
464 assert!(
465 result.is_err(),
466 "truncated i64 payload must return Err, not panic"
467 );
468 }
469
470 #[test]
471 fn decode_truncated_string_returns_error() {
472 let input = {
476 let mut v = vec![4u8]; v.extend_from_slice(&255u32.to_le_bytes()); v
480 };
481 let result = decode_row_from_wal(&input);
482 assert!(
483 result.is_err(),
484 "truncated string payload must return Err, not panic"
485 );
486 }
487
488 #[test]
489 fn decode_huge_vector_count_returns_error() {
490 let input = {
495 let mut v = vec![9u8]; v.extend_from_slice(&0x7FFF_FFFFu32.to_le_bytes()); v
499 };
500 let result = decode_row_from_wal(&input);
501 assert!(
502 result.is_err(),
503 "huge vector count with no payload must return Err, not panic or OOM"
504 );
505 }
506
507 #[test]
508 fn decode_truncated_decimal_returns_error() {
509 let input = {
512 let mut v = vec![7u8]; v.extend_from_slice(&[0u8; 4]); v
515 };
516 let result = decode_row_from_wal(&input);
517 assert!(
518 result.is_err(),
519 "truncated decimal payload must return Err, not panic"
520 );
521 }
522}