Skip to main content

ytsaurus_skiff/
wire.rs

1//! Bounded encoding and decoding of schema-described Skiff values.
2//!
3//! [`Encoder`] deliberately follows the Go SDK's `NewEncoder`: it writes a
4//! `Variant16` tag of zero followed by the one table schema it was given.
5//! [`Decoder`] is the complementary job-input form: it receives a [`Format`]
6//! and uses the leading `Variant16` tag to select an input table schema.
7//!
8//! The dynamic [`Value`] form is the codec's compatibility layer. Typed
9//! `SkiffRow` support will build on it without making the framing or limit
10//! checks depend on Serde internals.
11
12use std::io::{ErrorKind, Read, Write};
13
14use thiserror::Error;
15
16use crate::{Format, Schema, WireType};
17
18/// The largest `string32` or `yson32` payload accepted by default.
19///
20/// This is YTsaurus's documented maximum row size and the same bound used by
21/// the Go SDK reader. It protects a decoder from turning an untrusted `u32`
22/// length prefix into an unbounded allocation.
23pub const DEFAULT_MAX_BLOB_BYTES: usize = 128 * 1024 * 1024;
24
25/// The largest decoded footprint accepted for one row by default.
26///
27/// Bounding blobs does not bound a row. Skiff is compact and positional: one
28/// `repeated_variant8` item can cost a single tag byte on the wire and tens of
29/// bytes of decoded [`Value`], and the item loop runs until the end tag or the
30/// end of the stream. Without this limit a few hundred megabytes of hostile or
31/// corrupt stream decode into a multi-gigabyte value and the process is
32/// OOM-killed rather than given an error.
33///
34/// This is the same ceiling `ytsaurus-job`'s YSON `JobReader` puts on one
35/// record.
36pub const DEFAULT_MAX_ROW_BYTES: usize = 256 * 1024 * 1024;
37
38/// A dynamically decoded Skiff value.
39///
40/// Every variant maps directly to one [`WireType`]. Compound values retain
41/// tags and wire order so they can be re-encoded without a lossy conversion.
42#[derive(Debug, Clone, PartialEq)]
43pub enum Value {
44    /// `nothing`, an empty payload.
45    Nothing,
46    /// `boolean`.
47    Boolean(bool),
48    /// `int8`.
49    Int8(i8),
50    /// `int16`.
51    Int16(i16),
52    /// `int32`.
53    Int32(i32),
54    /// `int64`.
55    Int64(i64),
56    /// `int128`.
57    Int128(i128),
58    /// `int256`, kept as its little-endian two's-complement bytes.
59    Int256([u8; 32]),
60    /// `uint8`.
61    Uint8(u8),
62    /// `uint16`.
63    Uint16(u16),
64    /// `uint32`.
65    Uint32(u32),
66    /// `uint64`.
67    Uint64(u64),
68    /// `double`.
69    Double(f64),
70    /// `string32`, which may contain non-UTF-8 bytes.
71    Bytes(Vec<u8>),
72    /// `yson32`, kept as one binary-YSON payload.
73    Yson(Vec<u8>),
74    /// A `variant8` or `variant16` child selected by `tag`.
75    Variant {
76        /// The selected child index.
77        tag: u16,
78        /// The selected child's value.
79        value: Box<Value>,
80    },
81    /// A `repeated_variant8` or `repeated_variant16` sequence.
82    RepeatedVariants(Vec<Variant>),
83    /// A tuple, in schema-child order.
84    Tuple(Vec<Value>),
85}
86
87/// One non-terminal element of [`Value::RepeatedVariants`].
88#[derive(Debug, Clone, PartialEq)]
89pub struct Variant {
90    /// The selected child index.
91    pub tag: u16,
92    /// The selected child's value.
93    pub value: Value,
94}
95
96impl Value {
97    /// A stable description used when a value does not fit its schema node.
98    #[must_use]
99    pub const fn kind(&self) -> &'static str {
100        match self {
101            Self::Nothing => "nothing",
102            Self::Boolean(_) => "boolean",
103            Self::Int8(_) => "int8",
104            Self::Int16(_) => "int16",
105            Self::Int32(_) => "int32",
106            Self::Int64(_) => "int64",
107            Self::Int128(_) => "int128",
108            Self::Int256(_) => "int256",
109            Self::Uint8(_) => "uint8",
110            Self::Uint16(_) => "uint16",
111            Self::Uint32(_) => "uint32",
112            Self::Uint64(_) => "uint64",
113            Self::Double(_) => "double",
114            Self::Bytes(_) => "string32",
115            Self::Yson(_) => "yson32",
116            Self::Variant { .. } => "variant",
117            Self::RepeatedVariants(_) => "repeated variant",
118            Self::Tuple(_) => "tuple",
119        }
120    }
121}
122
123/// An encoder for one Skiff table stream.
124///
125/// Each row is prefixed with `0_u16`, exactly as Go's `skiff.NewEncoder`
126/// does. A job with several output descriptors creates one encoder per output
127/// table; it does not multiplex table indexes into one descriptor.
128#[derive(Debug)]
129pub struct Encoder<W> {
130    output: W,
131    schema: Schema,
132    max_blob_bytes: usize,
133}
134
135impl<W: Write> Encoder<W> {
136    /// Creates an encoder for a single table schema.
137    ///
138    /// The schema must satisfy the YTsaurus table-format contract in full: a
139    /// tuple root whose every child is named. These are the rules
140    /// [`Format::new`] applies, and they are applied here so that a schema
141    /// which can be encoded is always one that can also be declared to the
142    /// cluster and read back by the matching [`Decoder`].
143    pub fn new(output: W, schema: Schema) -> Result<Self, CodecError> {
144        schema.validate().map_err(CodecError::InvalidSchema)?;
145        if schema.wire_type != WireType::Tuple {
146            return Err(CodecError::TableSchemaMustBeTuple {
147                found: schema.wire_type,
148            });
149        }
150        crate::schema::validate_table_schema(&schema).map_err(CodecError::InvalidSchema)?;
151        Ok(Self {
152            output,
153            schema,
154            max_blob_bytes: DEFAULT_MAX_BLOB_BYTES,
155        })
156    }
157
158    /// Changes the maximum `string32` or `yson32` payload this encoder accepts.
159    #[must_use]
160    pub fn with_max_blob_bytes(mut self, bytes: usize) -> Self {
161        self.max_blob_bytes = bytes;
162        self
163    }
164
165    /// Writes one row with the single-table `Variant16` tag.
166    pub fn write(&mut self, row: &Value) -> Result<(), CodecError> {
167        write_all(&mut self.output, &0_u16.to_le_bytes())?;
168        encode_value(&mut self.output, &self.schema, row, self.max_blob_bytes)
169    }
170
171    /// Flushes bytes held by the caller's writer.
172    pub fn flush(&mut self) -> Result<(), CodecError> {
173        self.output.flush().map_err(CodecError::Write)
174    }
175
176    /// Returns the output writer after flushing it.
177    pub fn into_inner(mut self) -> Result<W, CodecError> {
178        self.flush()?;
179        Ok(self.output)
180    }
181}
182
183/// A decoder for a multiplexed YTsaurus Skiff input stream.
184#[derive(Debug)]
185pub struct Decoder<R> {
186    input: R,
187    format: Format,
188    max_blob_bytes: usize,
189    max_row_bytes: usize,
190}
191
192impl<R: Read> Decoder<R> {
193    /// Creates a decoder for rows described by `format`.
194    #[must_use]
195    pub fn new(input: R, format: Format) -> Self {
196        Self {
197            input,
198            format,
199            max_blob_bytes: DEFAULT_MAX_BLOB_BYTES,
200            max_row_bytes: DEFAULT_MAX_ROW_BYTES,
201        }
202    }
203
204    /// Changes the maximum `string32` or `yson32` payload this decoder accepts.
205    #[must_use]
206    pub fn with_max_blob_bytes(mut self, bytes: usize) -> Self {
207        self.max_blob_bytes = bytes;
208        self
209    }
210
211    /// Changes the maximum decoded footprint this decoder accepts per row.
212    ///
213    /// See [`DEFAULT_MAX_ROW_BYTES`] for why a blob limit alone is not enough.
214    #[must_use]
215    pub fn with_max_row_bytes(mut self, bytes: usize) -> Self {
216        self.max_row_bytes = bytes;
217        self
218    }
219
220    /// Decodes the next table-indexed row, or `None` at a clean end of stream.
221    ///
222    /// An end of stream after any part of a row is [`CodecError::Truncated`],
223    /// never a successful short read.
224    pub fn next_row(&mut self) -> Result<Option<(usize, Value)>, CodecError> {
225        let Some(first) = read_first_byte(&mut self.input)? else {
226            return Ok(None);
227        };
228        let mut table_tag = [first, 0];
229        read_exact(&mut self.input, &mut table_tag[1..], "table Variant16 tag")?;
230        let index = usize::from(u16::from_le_bytes(table_tag));
231        let schema = self
232            .format
233            .table_schema(index)
234            .map_err(CodecError::InvalidSchema)?;
235        let mut budget = RowBudget::new(self.max_blob_bytes, self.max_row_bytes);
236        let row = decode_value(&mut self.input, schema, &mut budget)?;
237        Ok(Some((index, row)))
238    }
239
240    /// Advances past the next row without building it, returning its table
241    /// index, or `None` at a clean end of stream.
242    ///
243    /// Framing, schema and limit checks are the ones [`Self::next_row`]
244    /// applies, so a stream this accepts is exactly a stream that decodes.
245    /// What it does not do is allocate: no `Vec` per blob, no `Box` per
246    /// variant, no `Value` tree to throw away. That is the difference between
247    /// asking whether a stream is a whole number of rows and decoding it to
248    /// find out — the first question is the one a caller validating a table
249    /// write is asking.
250    pub fn skip_row(&mut self) -> Result<Option<usize>, CodecError> {
251        let Some(first) = read_first_byte(&mut self.input)? else {
252            return Ok(None);
253        };
254        let mut table_tag = [first, 0];
255        read_exact(&mut self.input, &mut table_tag[1..], "table Variant16 tag")?;
256        let index = usize::from(u16::from_le_bytes(table_tag));
257        let schema = self
258            .format
259            .table_schema(index)
260            .map_err(CodecError::InvalidSchema)?;
261        let mut budget = RowBudget::new(self.max_blob_bytes, self.max_row_bytes);
262        skip_value(&mut self.input, schema, &mut budget)?;
263        Ok(Some(index))
264    }
265
266    /// Returns the wrapped input reader.
267    #[must_use]
268    pub fn into_inner(self) -> R {
269        self.input
270    }
271}
272
273/// A wire or I/O failure while encoding or decoding Skiff.
274#[derive(Debug, Error)]
275pub enum CodecError {
276    /// A schema was invalid before a stream was touched.
277    #[error("invalid Skiff schema: {0}")]
278    InvalidSchema(#[source] crate::SchemaError),
279    /// A table format schema did not have the required tuple root.
280    #[error("Skiff table schema root must be tuple, got {found}")]
281    TableSchemaMustBeTuple {
282        /// The root wire type that was supplied.
283        found: WireType,
284    },
285    /// The stream ended after a row had started.
286    #[error("Skiff stream ended while reading {context}")]
287    Truncated {
288        /// The incomplete portion of the stream.
289        context: &'static str,
290    },
291    /// One row would decode into more memory than the configured limit.
292    #[error("Skiff row does not fit in the {limit}-byte decode limit")]
293    RowTooLarge {
294        /// The configured maximum decoded footprint of one row.
295        limit: usize,
296    },
297    /// A declared blob size would exceed the configured resource limit.
298    #[error("Skiff {wire_type} payload is {length} bytes, exceeding the {limit}-byte limit")]
299    BlobTooLarge {
300        /// The schema type that carried the length.
301        wire_type: WireType,
302        /// The decoded or supplied payload size.
303        length: usize,
304        /// The configured maximum.
305        limit: usize,
306    },
307    /// An encoded value did not match its schema node.
308    #[error("Skiff {expected} node cannot encode {actual}")]
309    ValueDoesNotMatchSchema {
310        /// The schema's requested wire type.
311        expected: WireType,
312        /// The supplied value's stable kind.
313        actual: &'static str,
314    },
315    /// A tuple carried a different number of values than its schema children.
316    #[error("Skiff tuple has {actual} values, but its schema has {expected}")]
317    TupleLength {
318        /// The schema child count.
319        expected: usize,
320        /// The supplied value count.
321        actual: usize,
322    },
323    /// A variant tag did not select a child in its schema.
324    #[error("Skiff {wire_type} tag {tag} has no matching child in a {children}-child schema")]
325    InvalidVariantTag {
326        /// The variant wire type.
327        wire_type: WireType,
328        /// The invalid tag.
329        tag: u16,
330        /// The schema child count.
331        children: usize,
332    },
333    /// A repeated-variant value carried a tag wider than the variant supports.
334    #[error("Skiff {wire_type} tag {tag} cannot fit in its tag width")]
335    VariantTagTooWide {
336        /// The repeated-variant wire type.
337        wire_type: WireType,
338        /// The tag that did not fit.
339        tag: u16,
340    },
341    /// A caller supplied a payload that cannot fit into its `u32` length prefix.
342    #[error("Skiff {wire_type} payload is {length} bytes, which cannot fit in u32")]
343    BlobLengthOverflowsU32 {
344        /// The affected wire type.
345        wire_type: WireType,
346        /// The supplied payload length.
347        length: usize,
348    },
349    /// The underlying writer failed.
350    #[error("writing Skiff stream: {0}")]
351    Write(#[source] std::io::Error),
352    /// The underlying reader failed for a reason other than end of stream.
353    #[error("reading Skiff stream: {0}")]
354    Read(#[source] std::io::Error),
355}
356
357fn encode_value<W: Write>(
358    output: &mut W,
359    schema: &Schema,
360    value: &Value,
361    max_blob_bytes: usize,
362) -> Result<(), CodecError> {
363    match (schema.wire_type, value) {
364        (WireType::Nothing, Value::Nothing) => Ok(()),
365        (WireType::Boolean, Value::Boolean(value)) => write_all(output, &[u8::from(*value)]),
366        (WireType::Int8, Value::Int8(value)) => write_all(output, &value.to_le_bytes()),
367        (WireType::Int16, Value::Int16(value)) => write_all(output, &value.to_le_bytes()),
368        (WireType::Int32, Value::Int32(value)) => write_all(output, &value.to_le_bytes()),
369        (WireType::Int64, Value::Int64(value)) => write_all(output, &value.to_le_bytes()),
370        (WireType::Int128, Value::Int128(value)) => write_all(output, &value.to_le_bytes()),
371        (WireType::Int256, Value::Int256(value)) => write_all(output, value),
372        (WireType::Uint8, Value::Uint8(value)) => write_all(output, &value.to_le_bytes()),
373        (WireType::Uint16, Value::Uint16(value)) => write_all(output, &value.to_le_bytes()),
374        (WireType::Uint32, Value::Uint32(value)) => write_all(output, &value.to_le_bytes()),
375        (WireType::Uint64, Value::Uint64(value)) => write_all(output, &value.to_le_bytes()),
376        (WireType::Double, Value::Double(value)) => write_all(output, &value.to_le_bytes()),
377        (WireType::String32, Value::Bytes(value)) | (WireType::Yson32, Value::Yson(value)) => {
378            write_blob(output, schema.wire_type, value, max_blob_bytes)
379        }
380        (WireType::Variant8 | WireType::Variant16, Value::Variant { tag, value }) => {
381            let child = variant_child(schema, *tag)?;
382            write_variant_tag(output, schema.wire_type, *tag)?;
383            encode_value(output, child, value, max_blob_bytes)
384        }
385        (
386            WireType::RepeatedVariant8 | WireType::RepeatedVariant16,
387            Value::RepeatedVariants(items),
388        ) => {
389            for item in items {
390                let child = variant_child(schema, item.tag)?;
391                write_variant_tag(output, schema.wire_type, item.tag)?;
392                encode_value(output, child, &item.value, max_blob_bytes)?;
393            }
394            write_repeated_variant_end(output, schema.wire_type)
395        }
396        (WireType::Tuple, Value::Tuple(values)) => {
397            if values.len() != schema.children.len() {
398                return Err(CodecError::TupleLength {
399                    expected: schema.children.len(),
400                    actual: values.len(),
401                });
402            }
403            for (child, value) in schema.children.iter().zip(values) {
404                encode_value(output, child, value, max_blob_bytes)?;
405            }
406            Ok(())
407        }
408        (expected, value) => Err(CodecError::ValueDoesNotMatchSchema {
409            expected,
410            actual: value.kind(),
411        }),
412    }
413}
414
415/// How much decoded row one call to [`Decoder::next_row`] may still produce.
416///
417/// The charge is an estimate of the decoded footprint — one [`Value`] per
418/// decoded node, plus each blob payload — rather than an exact allocation
419/// count. That is deliberate: the bound has to be proportional to what a row
420/// costs in memory, and for repeated variants that differs from what it costs
421/// on the wire by more than an order of magnitude. The blob limit rides along
422/// because both bounds are consumed at the same points.
423struct RowBudget {
424    max_blob_bytes: usize,
425    limit: usize,
426    remaining: usize,
427}
428
429impl RowBudget {
430    fn new(max_blob_bytes: usize, max_row_bytes: usize) -> Self {
431        Self {
432            max_blob_bytes,
433            limit: max_row_bytes,
434            remaining: max_row_bytes,
435        }
436    }
437
438    /// Charges `bytes` against the row, before the memory is committed.
439    fn charge(&mut self, bytes: usize) -> Result<(), CodecError> {
440        self.remaining = self
441            .remaining
442            .checked_sub(bytes)
443            .ok_or(CodecError::RowTooLarge { limit: self.limit })?;
444        Ok(())
445    }
446}
447
448fn decode_value<R: Read>(
449    input: &mut R,
450    schema: &Schema,
451    budget: &mut RowBudget,
452) -> Result<Value, CodecError> {
453    // Every decoded node lands in a Box, a Vec or the row itself, so it costs
454    // at least one Value. Charging here also bounds the repeated-variant loop
455    // below: each item costs a Value, so the item count cannot outrun the
456    // budget however few wire bytes an item takes.
457    budget.charge(size_of::<Value>())?;
458    match schema.wire_type {
459        WireType::Nothing => Ok(Value::Nothing),
460        WireType::Boolean => Ok(Value::Boolean(read_byte(input, "boolean")? != 0)),
461        WireType::Int8 => Ok(Value::Int8(i8::from_le_bytes(read_array(input, "int8")?))),
462        WireType::Int16 => Ok(Value::Int16(i16::from_le_bytes(read_array(
463            input, "int16",
464        )?))),
465        WireType::Int32 => Ok(Value::Int32(i32::from_le_bytes(read_array(
466            input, "int32",
467        )?))),
468        WireType::Int64 => Ok(Value::Int64(i64::from_le_bytes(read_array(
469            input, "int64",
470        )?))),
471        WireType::Int128 => Ok(Value::Int128(i128::from_le_bytes(read_array(
472            input, "int128",
473        )?))),
474        WireType::Int256 => Ok(Value::Int256(read_array(input, "int256")?)),
475        WireType::Uint8 => Ok(Value::Uint8(read_byte(input, "uint8")?)),
476        WireType::Uint16 => Ok(Value::Uint16(u16::from_le_bytes(read_array(
477            input, "uint16",
478        )?))),
479        WireType::Uint32 => Ok(Value::Uint32(u32::from_le_bytes(read_array(
480            input, "uint32",
481        )?))),
482        WireType::Uint64 => Ok(Value::Uint64(u64::from_le_bytes(read_array(
483            input, "uint64",
484        )?))),
485        WireType::Double => Ok(Value::Double(f64::from_le_bytes(read_array(
486            input, "double",
487        )?))),
488        WireType::String32 => Ok(Value::Bytes(read_blob(input, WireType::String32, budget)?)),
489        WireType::Yson32 => Ok(Value::Yson(read_blob(input, WireType::Yson32, budget)?)),
490        WireType::Variant8 | WireType::Variant16 => {
491            let tag = read_variant_tag(input, schema.wire_type)?;
492            let child = variant_child(schema, tag)?;
493            let value = decode_value(input, child, budget)?;
494            Ok(Value::Variant {
495                tag,
496                value: Box::new(value),
497            })
498        }
499        WireType::RepeatedVariant8 | WireType::RepeatedVariant16 => {
500            let mut items = Vec::new();
501            loop {
502                let tag = read_variant_tag(input, schema.wire_type)?;
503                if is_repeated_variant_end(schema.wire_type, tag) {
504                    break;
505                }
506                let child = variant_child(schema, tag)?;
507                items.push(Variant {
508                    tag,
509                    value: decode_value(input, child, budget)?,
510                });
511            }
512            Ok(Value::RepeatedVariants(items))
513        }
514        WireType::Tuple => {
515            let values = schema
516                .children
517                .iter()
518                .map(|child| decode_value(input, child, budget))
519                .collect::<Result<_, _>>()?;
520            Ok(Value::Tuple(values))
521        }
522    }
523}
524
525/// The wire width and truncation context of a fixed-size type.
526///
527/// The contexts match [`decode_value`]'s, so a truncated stream reports the
528/// same field whichever way it was read.
529const fn fixed_width(wire_type: WireType) -> Option<(u64, &'static str)> {
530    match wire_type {
531        WireType::Boolean => Some((1, "boolean")),
532        WireType::Int8 => Some((1, "int8")),
533        WireType::Int16 => Some((2, "int16")),
534        WireType::Int32 => Some((4, "int32")),
535        WireType::Int64 => Some((8, "int64")),
536        WireType::Int128 => Some((16, "int128")),
537        WireType::Int256 => Some((32, "int256")),
538        WireType::Uint8 => Some((1, "uint8")),
539        WireType::Uint16 => Some((2, "uint16")),
540        WireType::Uint32 => Some((4, "uint32")),
541        WireType::Uint64 => Some((8, "uint64")),
542        WireType::Double => Some((8, "double")),
543        _ => None,
544    }
545}
546
547fn skip_value<R: Read>(
548    input: &mut R,
549    schema: &Schema,
550    budget: &mut RowBudget,
551) -> Result<(), CodecError> {
552    // Charged as decode_value charges, so the two accept the same rows. A
553    // stream that skips is a stream the decoder can read, which is the whole
554    // value of asking the cheap question.
555    budget.charge(size_of::<Value>())?;
556    if let Some((width, context)) = fixed_width(schema.wire_type) {
557        return skip_exact(input, width, context);
558    }
559    match schema.wire_type {
560        WireType::Nothing => Ok(()),
561        WireType::String32 | WireType::Yson32 => {
562            skip_blob(input, schema.wire_type, budget)?;
563            Ok(())
564        }
565        WireType::Variant8 | WireType::Variant16 => {
566            let tag = read_variant_tag(input, schema.wire_type)?;
567            skip_value(input, variant_child(schema, tag)?, budget)
568        }
569        WireType::RepeatedVariant8 | WireType::RepeatedVariant16 => loop {
570            let tag = read_variant_tag(input, schema.wire_type)?;
571            if is_repeated_variant_end(schema.wire_type, tag) {
572                return Ok(());
573            }
574            skip_value(input, variant_child(schema, tag)?, budget)?;
575        },
576        WireType::Tuple => {
577            for child in &schema.children {
578                skip_value(input, child, budget)?;
579            }
580            Ok(())
581        }
582        // Every fixed-width type is answered above.
583        _ => unreachable!("fixed_width covers the remaining wire types"),
584    }
585}
586
587fn skip_blob<R: Read>(
588    input: &mut R,
589    wire_type: WireType,
590    budget: &mut RowBudget,
591) -> Result<(), CodecError> {
592    let length = usize::try_from(u32::from_le_bytes(read_array(input, "blob length")?))
593        .expect("u32 always fits usize on supported Rust targets");
594    check_blob_length(wire_type, length, budget.max_blob_bytes)?;
595    budget.charge(length)?;
596    skip_exact(input, length as u64, "blob payload")
597}
598
599fn skip_exact<R: Read>(input: &mut R, count: u64, context: &'static str) -> Result<(), CodecError> {
600    let skipped = std::io::copy(&mut input.by_ref().take(count), &mut std::io::sink())
601        .map_err(CodecError::Read)?;
602    if skipped != count {
603        return Err(CodecError::Truncated { context });
604    }
605    Ok(())
606}
607
608fn variant_child(schema: &Schema, tag: u16) -> Result<&Schema, CodecError> {
609    schema
610        .children
611        .get(usize::from(tag))
612        .ok_or(CodecError::InvalidVariantTag {
613            wire_type: schema.wire_type,
614            tag,
615            children: schema.children.len(),
616        })
617}
618
619fn write_variant_tag<W: Write>(
620    output: &mut W,
621    wire_type: WireType,
622    tag: u16,
623) -> Result<(), CodecError> {
624    match wire_type {
625        WireType::Variant8 | WireType::RepeatedVariant8 => {
626            let tag =
627                u8::try_from(tag).map_err(|_| CodecError::VariantTagTooWide { wire_type, tag })?;
628            write_all(output, &[tag])
629        }
630        WireType::Variant16 | WireType::RepeatedVariant16 => write_all(output, &tag.to_le_bytes()),
631        _ => unreachable!("only variant schema nodes request a variant tag"),
632    }
633}
634
635fn read_variant_tag<R: Read>(input: &mut R, wire_type: WireType) -> Result<u16, CodecError> {
636    match wire_type {
637        WireType::Variant8 | WireType::RepeatedVariant8 => {
638            Ok(u16::from(read_byte(input, "variant8 tag")?))
639        }
640        WireType::Variant16 | WireType::RepeatedVariant16 => {
641            Ok(u16::from_le_bytes(read_array(input, "variant16 tag")?))
642        }
643        _ => unreachable!("only variant schema nodes request a variant tag"),
644    }
645}
646
647fn write_repeated_variant_end<W: Write>(
648    output: &mut W,
649    wire_type: WireType,
650) -> Result<(), CodecError> {
651    match wire_type {
652        WireType::RepeatedVariant8 => write_all(output, &[u8::MAX]),
653        WireType::RepeatedVariant16 => write_all(output, &u16::MAX.to_le_bytes()),
654        _ => unreachable!("only repeated-variant schema nodes have an end tag"),
655    }
656}
657
658fn is_repeated_variant_end(wire_type: WireType, tag: u16) -> bool {
659    match wire_type {
660        WireType::RepeatedVariant8 => tag == u16::from(u8::MAX),
661        WireType::RepeatedVariant16 => tag == u16::MAX,
662        _ => unreachable!("only repeated-variant schema nodes have an end tag"),
663    }
664}
665
666fn write_blob<W: Write>(
667    output: &mut W,
668    wire_type: WireType,
669    value: &[u8],
670    max_blob_bytes: usize,
671) -> Result<(), CodecError> {
672    check_blob_length(wire_type, value.len(), max_blob_bytes)?;
673    let length = u32::try_from(value.len()).map_err(|_| CodecError::BlobLengthOverflowsU32 {
674        wire_type,
675        length: value.len(),
676    })?;
677    write_all(output, &length.to_le_bytes())?;
678    write_all(output, value)
679}
680
681fn read_blob<R: Read>(
682    input: &mut R,
683    wire_type: WireType,
684    budget: &mut RowBudget,
685) -> Result<Vec<u8>, CodecError> {
686    let length = usize::try_from(u32::from_le_bytes(read_array(input, "blob length")?))
687        .expect("u32 always fits usize on supported Rust targets");
688    check_blob_length(wire_type, length, budget.max_blob_bytes)?;
689    budget.charge(length)?;
690    let mut value = vec![0; length];
691    read_exact(input, &mut value, "blob payload")?;
692    Ok(value)
693}
694
695fn check_blob_length(
696    wire_type: WireType,
697    length: usize,
698    max_blob_bytes: usize,
699) -> Result<(), CodecError> {
700    if length > max_blob_bytes {
701        return Err(CodecError::BlobTooLarge {
702            wire_type,
703            length,
704            limit: max_blob_bytes,
705        });
706    }
707    Ok(())
708}
709
710fn write_all<W: Write>(output: &mut W, bytes: &[u8]) -> Result<(), CodecError> {
711    output.write_all(bytes).map_err(CodecError::Write)
712}
713
714fn read_first_byte<R: Read>(input: &mut R) -> Result<Option<u8>, CodecError> {
715    let mut byte = [0; 1];
716    loop {
717        match input.read(&mut byte) {
718            Ok(0) => return Ok(None),
719            Ok(_) => return Ok(Some(byte[0])),
720            Err(error) if error.kind() == ErrorKind::Interrupted => {}
721            Err(error) => return Err(CodecError::Read(error)),
722        }
723    }
724}
725
726fn read_byte<R: Read>(input: &mut R, context: &'static str) -> Result<u8, CodecError> {
727    let mut byte = [0; 1];
728    read_exact(input, &mut byte, context)?;
729    Ok(byte[0])
730}
731
732fn read_array<R: Read, const N: usize>(
733    input: &mut R,
734    context: &'static str,
735) -> Result<[u8; N], CodecError> {
736    let mut bytes = [0; N];
737    read_exact(input, &mut bytes, context)?;
738    Ok(bytes)
739}
740
741fn read_exact<R: Read>(
742    input: &mut R,
743    bytes: &mut [u8],
744    context: &'static str,
745) -> Result<(), CodecError> {
746    input.read_exact(bytes).map_err(|error| {
747        if error.kind() == ErrorKind::UnexpectedEof {
748            CodecError::Truncated { context }
749        } else {
750            CodecError::Read(error)
751        }
752    })
753}