Skip to main content

midnight_serialize/
deserializable.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
14use crate::VecExt;
15use crate::serializable::GLOBAL_TAG;
16use crate::tagged::Tagged;
17use std::borrow::Cow;
18use std::io::{self, BufRead, Read, Seek};
19use std::marker::PhantomData;
20use std::sync::Arc;
21use std::{collections::HashMap, collections::HashSet, hash::Hash};
22
23#[cfg(debug_assertions)]
24pub const RECURSION_LIMIT: u32 = 50;
25#[cfg(not(debug_assertions))]
26pub const RECURSION_LIMIT: u32 = 250;
27
28// Top-level deserialization function
29pub fn tagged_deserialize<T: Deserializable + Tagged>(reader: impl Read) -> std::io::Result<T> {
30    tagged_deserialize_inner(reader, true)
31}
32
33pub fn tagged_deserialize_sequence<T: Deserializable + Tagged>(
34    mut reader: impl BufRead,
35) -> std::io::Result<Vec<T>> {
36    let mut res = vec![];
37    while !reader.fill_buf()?.is_empty() {
38        res.push(tagged_deserialize_inner(&mut reader, false)?);
39    }
40    Ok(res)
41}
42
43/// Attempts to identify the tag a stream starts with without consuming it, allowing determining a
44/// stream's type *before* deserializing it.
45pub fn peek_tag(reader: &mut (impl Read + Seek)) -> std::io::Result<String> {
46    let position = reader.stream_position()?;
47    // Note that colons are special-cased -- we should expect two, one for `GLOBAL_TAG`, and one
48    // for the end of the read tag. We read up to a limit of 512 bytes, and then take up to the
49    // second b':', converting to string, returning an error if not possible.
50    const READ_LIMIT: usize = 512;
51    let mut buf = [0u8; READ_LIMIT];
52    let mut offset = 0;
53    while offset < READ_LIMIT {
54        let read = reader.read(&mut buf[offset..])?;
55        if read == 0 {
56            break;
57        }
58        offset += read;
59    }
60    reader.seek(std::io::SeekFrom::Start(position))?;
61    let err = |msg| io::Error::new(io::ErrorKind::InvalidData, msg);
62    if !buf.starts_with(GLOBAL_TAG.as_bytes()) {
63        return Err(err(format!(
64            "tagged data does not begin with '{GLOBAL_TAG}'"
65        )));
66    }
67    let second_colon = buf
68        .iter()
69        .enumerate()
70        .filter(|(_, b)| **b == b':')
71        .nth(1)
72        .ok_or_else(|| err("tagged data does not begin with a colon-separated tag".to_string()))?
73        .0;
74    let raw_tag = &buf[GLOBAL_TAG.len()..second_colon];
75    String::from_utf8(raw_tag.to_owned())
76        .map_err(|e| err(format!("tag not utf-8: {e}")))
77        .map(|s| {
78            s.replace(
79                |c: char| -> bool { !c.is_ascii_alphanumeric() && !":_-()[],".contains(c) },
80                "�",
81            )
82        })
83}
84
85fn tagged_deserialize_inner<T: Deserializable + Tagged>(
86    mut reader: impl Read,
87    ensure_consumed: bool,
88) -> std::io::Result<T> {
89    let tag_expected = format!("{GLOBAL_TAG}{}:", T::tag());
90    let mut read_tag = vec![0u8; tag_expected.len()];
91    let mut remaining_tag_buf = &mut read_tag[..];
92    while !remaining_tag_buf.is_empty() {
93        let read = reader.read(remaining_tag_buf)?;
94        if read == 0 {
95            let rem = remaining_tag_buf.len();
96            let len = read_tag.len() - rem;
97            read_tag.truncate(len);
98            break;
99        }
100        remaining_tag_buf = &mut remaining_tag_buf[read..];
101    }
102    if read_tag != tag_expected.as_bytes() {
103        let sanitised = String::from_utf8_lossy(&read_tag).replace(
104            |c: char| -> bool { !c.is_ascii_alphanumeric() && !":_-()[],".contains(c) },
105            "�",
106        );
107        return Err(std::io::Error::new(
108            std::io::ErrorKind::InvalidData,
109            format!("expected header tag '{tag_expected}', got '{sanitised}'"),
110        ));
111    }
112    let value = <T as Deserializable>::deserialize(&mut reader, 0)?;
113
114    if !ensure_consumed {
115        return Ok(value);
116    }
117
118    #[allow(clippy::unbuffered_bytes)] // we can permit a potentally inefficient count here, as in
119    let count = reader.bytes().count(); // the happy path it should be 0
120
121    if count == 0 {
122        return Ok(value);
123    }
124
125    Err(std::io::Error::new(
126        std::io::ErrorKind::InvalidData,
127        format!(
128            "Not all bytes read deserializing '{}'; {} bytes remaining",
129            tag_expected, count
130        ),
131    ))
132}
133
134pub trait Deserializable
135where
136    Self: Sized,
137{
138    const LIMIT_RECURSION: bool = true;
139
140    fn deserialize(reader: &mut impl Read, recursion_depth: u32) -> std::io::Result<Self>;
141
142    fn check_rec(depth: &mut u32) -> std::io::Result<()> {
143        if Self::LIMIT_RECURSION {
144            *depth += 1;
145            if *depth > RECURSION_LIMIT {
146                return Err(std::io::Error::new(
147                    std::io::ErrorKind::InvalidData,
148                    "exceeded recursion depth deserializing",
149                ));
150            }
151        }
152        Ok(())
153    }
154}
155
156impl<T: Deserializable> Deserializable for Vec<T> {
157    fn deserialize(reader: &mut impl Read, mut recursion_depth: u32) -> std::io::Result<Self> {
158        Self::check_rec(&mut recursion_depth)?;
159        let len = <u32 as Deserializable>::deserialize(reader, recursion_depth)?;
160        let mut result = Vec::with_bounded_capacity(len as usize);
161        for _ in 0..len {
162            result.push(<T as Deserializable>::deserialize(reader, recursion_depth)?);
163        }
164        Ok(result)
165    }
166}
167
168impl<K: Deserializable + PartialOrd + Hash + Eq, V: Deserializable> Deserializable
169    for HashMap<K, V>
170{
171    fn deserialize(reader: &mut impl Read, mut recursion_depth: u32) -> std::io::Result<Self> {
172        Self::check_rec(&mut recursion_depth)?;
173        let len = <u32 as Deserializable>::deserialize(reader, recursion_depth)?;
174        let mut result = Vec::with_bounded_capacity(len as usize);
175        for _ in 0..len {
176            let k = <K as Deserializable>::deserialize(reader, recursion_depth)?;
177            let v = <V as Deserializable>::deserialize(reader, recursion_depth)?;
178            result.push((k, v));
179        }
180        if result.iter().is_sorted_by_key(|(k, _)| k) && result.windows(2).all(|window| window[0].0 != window[1].0) {
181            Ok(result.into_iter().collect())
182        } else {
183            Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "HashMap not normalized (not sorted or duplicate keys)"))
184        }
185    }
186}
187
188impl<T: Deserializable + PartialOrd + Hash + Eq> Deserializable for HashSet<T> {
189    fn deserialize(reader: &mut impl Read, mut recursion_depth: u32) -> std::io::Result<Self> {
190        Self::check_rec(&mut recursion_depth)?;
191        let len = <u32 as Deserializable>::deserialize(reader, recursion_depth)?;
192        let mut result = Vec::with_bounded_capacity(len as usize);
193        for _ in 0..len {
194            result.push(<T as Deserializable>::deserialize(reader, recursion_depth)?);
195        }
196        if result.iter().is_sorted() && result.windows(2).all(|window| window[0] != window[1]) {
197            Ok(result.into_iter().collect())
198        } else {
199            Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "HashSet not normalized (not sorted or duplicate entries)"))
200        }
201    }
202}
203
204impl<T: Deserializable> Deserializable for Option<T> {
205    fn deserialize(reader: &mut impl Read, mut recursion_depth: u32) -> std::io::Result<Self> {
206        Self::check_rec(&mut recursion_depth)?;
207        let some = <u8 as Deserializable>::deserialize(reader, recursion_depth)?;
208        match some {
209            0 => Ok(None),
210            1 => Ok(Some(<T as Deserializable>::deserialize(
211                reader,
212                recursion_depth,
213            )?)),
214            _ => Err(std::io::Error::new(
215                std::io::ErrorKind::InvalidData,
216                format!("Invalid discriminant: {}.", some),
217            )),
218        }
219    }
220}
221
222impl<T: Deserializable> Deserializable for Arc<T> {
223    fn deserialize(
224        reader: &mut impl Read,
225        mut recursion_depth: u32,
226    ) -> Result<Self, std::io::Error> {
227        Self::check_rec(&mut recursion_depth)?;
228        Ok(Arc::new(T::deserialize(reader, recursion_depth)?))
229    }
230}
231
232impl<const N: usize> Deserializable for [u8; N] {
233    fn deserialize(reader: &mut impl Read, _recursion_depth: u32) -> std::io::Result<Self> {
234        let mut res = [0u8; N];
235        reader.read_exact(&mut res[..])?;
236        Ok(res)
237    }
238}
239
240impl<T: ?Sized> Deserializable for PhantomData<T> {
241    fn deserialize(_reader: &mut impl Read, _recursion_depth: u32) -> std::io::Result<Self> {
242        Ok(PhantomData)
243    }
244}
245
246impl<T: Deserializable> Deserializable for Box<T> {
247    fn deserialize(reader: &mut impl Read, mut recursion_depth: u32) -> std::io::Result<Self> {
248        Self::check_rec(&mut recursion_depth)?;
249        T::deserialize(reader, recursion_depth).map(Box::new)
250    }
251}
252
253impl<'a, T: ToOwned + ?Sized> Deserializable for Cow<'a, T>
254where
255    T::Owned: Deserializable,
256{
257    fn deserialize(reader: &mut impl Read, recursion_depth: u32) -> std::io::Result<Self> {
258        <T::Owned>::deserialize(reader, recursion_depth).map(Cow::Owned)
259    }
260}