midnight_serialize/
shape.rs

1// This file is part of midnight-ledger.
2// Copyright (C) Midnight Foundation
3// SPDX-License-Identifier: Apache-2.0
4// Licensed under the Apache License, Version 2.0 (the "License");
5// You may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7// http://www.apache.org/licenses/LICENSE-2.0
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14//! Descriptions of exactly which byte strings decode as a type, as finite
15//! automata, so that many encodings sharing structure can be checked without
16//! enumerating them.
17
18use std::fmt::Debug;
19use std::hash::Hash;
20
21/// The byte strings that [`crate::Deserializable::deserialize`] reads in full, as a
22/// finite automaton over bytes.
23///
24/// Encodings must be prefix-free, as self-delimiting encodings are: no
25/// accepted string continues into another. Checking a set of encodings costs
26/// time proportional to the number of distinct states reached, so `State`
27/// should take few values.
28pub trait EncodingShape {
29    /// A position in the automaton.
30    type State: Clone + Eq + Hash + Debug + 'static;
31    /// The state before any byte is read.
32    fn start() -> Self::State;
33    /// The state after reading `byte`, or `None` if no encoding continues so.
34    fn step(state: &Self::State, byte: u8) -> Option<Self::State>;
35    /// Whether the bytes read so far are a complete encoding.
36    fn accepts(state: &Self::State) -> bool;
37}
38
39/// Whether `T`'s shape accepts `bytes`.
40#[cfg(any(test, feature = "proptest"))]
41fn shape_accepts<T: EncodingShape>(bytes: &[u8]) -> bool {
42    bytes
43        .iter()
44        .try_fold(T::start(), |state, b| T::step(&state, *b))
45        .is_some_and(|state| T::accepts(&state))
46}
47
48/// Implements [`EncodingShape`] for a type by delegating to a shape type.
49#[macro_export]
50macro_rules! encoding_shape_as {
51    ($ty:ty, $shape:ty) => {
52        impl $crate::EncodingShape for $ty {
53            type State = <$shape as $crate::EncodingShape>::State;
54            fn start() -> Self::State {
55                <$shape as $crate::EncodingShape>::start()
56            }
57            fn step(state: &Self::State, byte: u8) -> Option<Self::State> {
58                <$shape as $crate::EncodingShape>::step(state, byte)
59            }
60            fn accepts(state: &Self::State) -> bool {
61                <$shape as $crate::EncodingShape>::accepts(state)
62            }
63        }
64    };
65}
66
67impl<const N: usize> EncodingShape for [u8; N] {
68    type State = usize;
69    fn start() -> usize {
70        0
71    }
72    fn step(read: &usize, _: u8) -> Option<usize> {
73        (*read < N).then_some(read + 1)
74    }
75    fn accepts(read: &usize) -> bool {
76        *read == N
77    }
78}
79
80impl EncodingShape for bool {
81    type State = bool;
82    fn start() -> bool {
83        false
84    }
85    fn step(done: &bool, byte: u8) -> Option<bool> {
86        (!done && byte <= 1).then_some(true)
87    }
88    fn accepts(done: &bool) -> bool {
89        *done
90    }
91}
92
93/// A position in a pair's shape.
94#[derive(Clone, PartialEq, Eq, Hash, Debug)]
95pub enum SeqState<A, B> {
96    /// Inside the first encoding.
97    First(A),
98    /// Inside the second encoding.
99    Second(B),
100}
101
102// Prefix-freedom makes the switch to `B` unambiguous.
103fn seq_normalize<A: EncodingShape, B: EncodingShape>(
104    state: SeqState<A::State, B::State>,
105) -> SeqState<A::State, B::State> {
106    match state {
107        SeqState::First(a) if A::accepts(&a) => SeqState::Second(B::start()),
108        other => other,
109    }
110}
111
112impl<A: EncodingShape, B: EncodingShape> EncodingShape for (A, B) {
113    type State = SeqState<A::State, B::State>;
114    fn start() -> Self::State {
115        seq_normalize::<A, B>(SeqState::First(A::start()))
116    }
117    fn step(state: &Self::State, byte: u8) -> Option<Self::State> {
118        Some(seq_normalize::<A, B>(match state {
119            SeqState::First(a) => SeqState::First(A::step(a, byte)?),
120            SeqState::Second(b) => SeqState::Second(B::step(b, byte)?),
121        }))
122    }
123    fn accepts(state: &Self::State) -> bool {
124        matches!(state, SeqState::Second(b) if B::accepts(b))
125    }
126}
127
128/// A position in the shape of a canonical SCALE compact integer.
129#[derive(Clone, PartialEq, Eq, Hash, Debug)]
130pub enum ScaleState {
131    /// Before the first byte.
132    Start,
133    /// A complete encoding.
134    Done,
135    /// Two-byte mode, before its nonzero second byte.
136    TwoByte,
137    /// Four-byte mode, with this many bytes left, and whether the third byte
138    /// was nonzero.
139    FourByte(u8, bool),
140    /// Big-integer mode, with this many of its `n` bytes left.
141    NByte { left: u8, n: u8 },
142}
143
144/// One step of a canonical SCALE compact integer whose value fits in `M`
145/// bytes, as `u32`, `u64` and `u128` encode.
146fn scale_step<const M: usize>(state: &ScaleState, byte: u8) -> Option<ScaleState> {
147    use ScaleState::*;
148    match *state {
149        Start => match byte & 0b11 {
150            0b00 => Some(Done),
151            0b01 => Some(TwoByte),
152            0b10 if M >= 4 => Some(FourByte(3, false)),
153            0b11 => {
154                // Big-integer mode is n = 4 + (byte >> 2) bytes, which must fit in `M`.
155                let n = (byte >> 2) + 4;
156                (n as usize <= M).then_some(NByte { left: n, n })
157            }
158            _ => None,
159        },
160        Done => None,
161        TwoByte => (byte != 0).then_some(Done),
162        FourByte(3, _) => Some(FourByte(2, false)),
163        FourByte(2, _) => Some(FourByte(1, byte != 0)),
164        FourByte(_, third_nonzero) => (third_nonzero || byte != 0).then_some(Done),
165        // Canonical: no zero top byte, and 4-byte big integers only above
166        // 2^30, which the shorter modes cover.
167        NByte { left: 1, n } => (byte != 0 && (n != 4 || byte >= 64)).then_some(Done),
168        NByte { left, n } => Some(NByte { left: left - 1, n }),
169    }
170}
171
172macro_rules! scale_shape {
173    ($ty:ty, $m:literal) => {
174        impl EncodingShape for $ty {
175            type State = ScaleState;
176            fn start() -> ScaleState {
177                ScaleState::Start
178            }
179            fn step(state: &ScaleState, byte: u8) -> Option<ScaleState> {
180                scale_step::<$m>(state, byte)
181            }
182            fn accepts(state: &ScaleState) -> bool {
183                *state == ScaleState::Done
184            }
185        }
186    };
187}
188
189scale_shape!(u32, 4);
190scale_shape!(u64, 8);
191scale_shape!(u128, 16);
192
193encoding_shape_as!((), [u8; 0]);
194encoding_shape_as!(u8, [u8; 1]);
195encoding_shape_as!(u16, [u8; 2]);
196encoding_shape_as!(i8, [u8; 1]);
197encoding_shape_as!(i16, [u8; 2]);
198encoding_shape_as!(i32, [u8; 4]);
199encoding_shape_as!(i64, [u8; 8]);
200encoding_shape_as!(i128, [u8; 16]);
201
202/// Checks that `T`'s shape accepts exactly the byte strings `T` decodes from
203/// in full, over encodings of arbitrary values and corruptions of them.
204#[cfg(feature = "proptest")]
205pub fn check_encoding_shape<T>()
206where
207    T: EncodingShape + crate::Serializable + crate::Deserializable + proptest::arbitrary::Arbitrary,
208{
209    use proptest::prelude::*;
210    use proptest::test_runner::{Config, TestRunner};
211
212    let decodes = |bytes: &[u8]| {
213        let mut reader = bytes;
214        T::deserialize(&mut reader, 0).is_ok() && reader.is_empty()
215    };
216    let encoding = || {
217        any::<T>().prop_map(|v| {
218            let mut bytes = Vec::new();
219            v.serialize(&mut bytes)
220                .expect("in-memory serialization succeeds");
221            bytes
222        })
223    };
224    let corrupted = (
225        encoding(),
226        any::<prop::sample::Index>(),
227        any::<u8>(),
228        0u8..4,
229    )
230        .prop_map(|(mut bytes, i, b, how): (Vec<u8>, _, u8, u8)| {
231            match how {
232                0 if !bytes.is_empty() => bytes.truncate(i.index(bytes.len())),
233                1 => bytes.push(b),
234                2 if !bytes.is_empty() => {
235                    let i = i.index(bytes.len());
236                    bytes[i] = b;
237                }
238                _ => bytes.insert(0, b),
239            }
240            bytes
241        });
242    let short = prop::collection::vec(any::<u8>(), 0..12);
243    let strategy = prop_oneof![encoding(), corrupted, short];
244    TestRunner::new(Config::with_cases(4096))
245        .run(&strategy, |bytes: Vec<u8>| {
246            prop_assert_eq!(
247                shape_accepts::<T>(&bytes),
248                decodes(&bytes),
249                "{:02x?}",
250                bytes
251            );
252            Ok(())
253        })
254        .unwrap_or_else(|e| {
255            panic!(
256                "{} shape disagrees with decoding: {e}",
257                std::any::type_name::<T>()
258            )
259        });
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use crate::{Deserializable, Serializable};
266
267    #[test]
268    fn primitive_shapes_match_decoding() {
269        check_encoding_shape::<()>();
270        check_encoding_shape::<bool>();
271        check_encoding_shape::<u8>();
272        check_encoding_shape::<u16>();
273        check_encoding_shape::<i32>();
274        check_encoding_shape::<i128>();
275        check_encoding_shape::<u32>();
276        check_encoding_shape::<u64>();
277        check_encoding_shape::<u128>();
278        check_encoding_shape::<[u8; 3]>();
279        check_encoding_shape::<(u64, bool)>();
280        check_encoding_shape::<(u8, u32)>();
281    }
282
283    // Arbitrary integers are mostly large, so cover each SCALE mode's
284    // boundaries explicitly.
285    #[test]
286    fn scale_shapes_match_decoding_at_mode_boundaries() {
287        fn check<T: EncodingShape + Serializable + Deserializable + TryFrom<u128>>(m: u32) {
288            for bits in 0..(8 * m) {
289                for v in [
290                    (1u128 << bits).wrapping_sub(1),
291                    1u128 << bits,
292                    (1u128 << bits) + 1,
293                ] {
294                    let Ok(v) = T::try_from(v) else { continue };
295                    let mut bytes = Vec::new();
296                    v.serialize(&mut bytes).unwrap();
297                    assert!(shape_accepts::<T>(&bytes), "{bytes:02x?}");
298                    for i in 0..bytes.len() {
299                        assert!(!shape_accepts::<T>(&bytes[..i]), "{bytes:02x?}");
300                    }
301                }
302            }
303        }
304        check::<u32>(4);
305        check::<u64>(8);
306        check::<u128>(16);
307    }
308}
midnight_serialize/shape.rs

midnight_serialize/
shape.rs
