iota_storage/
blob.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    io::{Read, Write},
7    marker::PhantomData,
8};
9
10use anyhow::{Result, anyhow, bail};
11use byteorder::ReadBytesExt;
12use integer_encoding::{VarInt, VarIntReader};
13use num_enum::{IntoPrimitive, TryFromPrimitive};
14use serde::{Serialize, de::DeserializeOwned};
15use tracing::error;
16
17pub const MAX_VARINT_LENGTH: usize = 10;
18pub const BLOB_ENCODING_BYTES: usize = 1;
19const MAX_PREALLOCATION: usize = 1 << 20;
20/// The largest blob [`Blob::read`] accepts. The length comes from the input,
21/// and a small zstd-compressed file can decompress to any length. Mainnet's
22/// checkpoint 0 is a single blob of about 2.15 GiB.
23const MAX_BLOB_LENGTH: u64 = 1 << 32;
24
25#[derive(Copy, Clone, Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
26#[repr(u8)]
27pub enum BlobEncoding {
28    Bcs = 1,
29}
30
31pub struct Blob {
32    pub data: Vec<u8>,
33    pub encoding: BlobEncoding,
34}
35
36impl Blob {
37    pub fn encode<T: Serialize>(value: &T, encoding: BlobEncoding) -> Result<Self> {
38        let value_buf = bcs::to_bytes(value)?;
39        let (data, encoding) = match encoding {
40            BlobEncoding::Bcs => (value_buf, encoding),
41        };
42        Ok(Blob { data, encoding })
43    }
44    pub fn decode<T: DeserializeOwned>(self) -> Result<T> {
45        let data = match &self.encoding {
46            BlobEncoding::Bcs => self.data,
47        };
48        let res = bcs::from_bytes(&data)?;
49        Ok(res)
50    }
51    /// Reads one blob from `rbuf`.
52    ///
53    /// Memory grows with the bytes actually read, not with the length the
54    /// blob declares: at most 1 MiB is reserved before reading, and at most
55    /// 4 GiB is read.
56    ///
57    /// # Errors
58    ///
59    /// Fails if reading `rbuf` fails, if the declared length is 0 or above
60    /// 4 GiB, if `rbuf` ends before that many data bytes, or if the encoding
61    /// is unknown.
62    pub fn read<R: Read>(rbuf: &mut R) -> Result<Blob> {
63        let len = rbuf.read_varint::<u64>()?;
64        if len == 0 {
65            bail!("Invalid object length of 0 in file");
66        }
67        if len > MAX_BLOB_LENGTH {
68            bail!("blob length {len} is above the maximum of {MAX_BLOB_LENGTH} bytes");
69        }
70        let encoding = rbuf.read_u8()?;
71        // The length comes from the input, so only part of it is reserved up
72        // front.
73        let mut data = Vec::with_capacity(len.min(MAX_PREALLOCATION as u64) as usize);
74        rbuf.take(len).read_to_end(&mut data)?;
75        if data.len() as u64 != len {
76            bail!(
77                "blob length {len} is beyond the remaining input of {} bytes",
78                data.len()
79            );
80        }
81        let blob = Blob {
82            data,
83            encoding: BlobEncoding::try_from(encoding)?,
84        };
85        Ok(blob)
86    }
87    pub fn write<W: Write>(&self, wbuf: &mut W) -> Result<usize> {
88        let mut buf = [0u8; MAX_VARINT_LENGTH];
89        let mut counter = 0;
90        let n = (self.data.len() as u64).encode_var(&mut buf);
91        wbuf.write_all(&buf[0..n])?;
92        counter += n;
93        buf[0] = self.encoding.into();
94        wbuf.write_all(&buf[0..BLOB_ENCODING_BYTES])?;
95        counter += 1;
96        wbuf.write_all(&self.data)?;
97        counter += self.data.len();
98        Ok(counter)
99    }
100    pub fn size(&self) -> usize {
101        let mut blob_size = self.data.len().required_space();
102        blob_size += BLOB_ENCODING_BYTES;
103        blob_size += self.data.len();
104        blob_size
105    }
106    pub fn to_bytes(&self) -> Vec<u8> {
107        [vec![self.encoding.into()], self.data.clone()].concat()
108    }
109    pub fn from_bytes<T: DeserializeOwned>(bytes: &[u8]) -> Result<T> {
110        let (encoding, data) = bytes.split_first().ok_or(anyhow!("empty bytes"))?;
111        Blob {
112            data: data.to_vec(),
113            encoding: BlobEncoding::try_from(*encoding)?,
114        }
115        .decode()
116    }
117}
118
119/// An iterator over blobs in a blob file. It ends at the end of the input, or
120/// logs an error and ends at the first blob it cannot read.
121pub struct BlobIter<T> {
122    reader: Box<dyn Read>,
123    _phantom: PhantomData<T>,
124}
125
126impl<T: DeserializeOwned> BlobIter<T> {
127    pub fn new(reader: Box<dyn Read>) -> Self {
128        Self {
129            reader,
130            _phantom: PhantomData,
131        }
132    }
133    /// Reads the next blob, or returns `None` at the end of the input.
134    fn next_blob(&mut self) -> Result<Option<T>> {
135        // Reads the first byte separately to tell the end of the input apart
136        // from a truncated blob.
137        let mut first = [0u8; 1];
138        if self.reader.read(&mut first)? == 0 {
139            return Ok(None);
140        }
141        let blob = Blob::read(&mut first.as_slice().chain(&mut self.reader))?;
142        blob.decode().map(Some)
143    }
144}
145
146impl<T: DeserializeOwned> Iterator for BlobIter<T> {
147    type Item = T;
148    fn next(&mut self) -> Option<Self::Item> {
149        self.next_blob().unwrap_or_else(|err| {
150            error!("failed to read a blob, ignoring the rest of the file: {err:#}");
151            None
152        })
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use std::io::Cursor;
159
160    use super::*;
161    use crate::{FileCompression, StorageFormat, make_iterator};
162
163    const MAGIC: u32 = 0x0000BEEF;
164
165    /// Encodes a blob header: the length varint followed by the BCS encoding
166    /// byte.
167    fn header(len: u64) -> Vec<u8> {
168        let mut bytes = len.encode_var_vec();
169        bytes.push(BlobEncoding::Bcs.into());
170        bytes
171    }
172
173    #[test]
174    fn write_read() {
175        // The second value is larger than the space reserved before reading.
176        for value in [vec![1u8, 2, 3], vec![7u8; 3 * MAX_PREALLOCATION]] {
177            let blob = Blob::encode(&value, BlobEncoding::Bcs).unwrap();
178            let mut bytes = Vec::new();
179            blob.write(&mut bytes).unwrap();
180            let read = Blob::read(&mut Cursor::new(bytes)).unwrap();
181            assert_eq!(read.data, blob.data);
182            assert_eq!(read.decode::<Vec<u8>>().unwrap(), value);
183        }
184    }
185
186    #[test]
187    fn rejects_length_beyond_input() {
188        // The middle length is that of mainnet's checkpoint 0.
189        for len in [MAX_BLOB_LENGTH, 2_313_050_709, 4] {
190            let mut bytes = header(len);
191            bytes.extend_from_slice(&[0u8; 3]);
192            let err = Blob::read(&mut Cursor::new(bytes)).err().unwrap();
193            assert!(
194                err.to_string()
195                    .contains("is beyond the remaining input of 3 bytes"),
196                "{err}"
197            );
198        }
199    }
200
201    /// A blob longer than the maximum is refused, also from a zstd-compressed
202    /// file.
203    #[test]
204    fn rejects_length_above_maximum() {
205        for compression in [FileCompression::None, FileCompression::Zstd] {
206            for len in [MAX_BLOB_LENGTH + 1, u64::MAX] {
207                let mut blob = header(len);
208                blob.extend_from_slice(&[0u8; 3]);
209                let mut file = MAGIC.to_be_bytes().to_vec();
210                file.push(StorageFormat::Blob.into());
211                file.push(compression.into());
212                match compression {
213                    FileCompression::None => file.extend(blob),
214                    FileCompression::Zstd => {
215                        FileCompression::zstd_compress(&mut Cursor::new(blob), &mut file).unwrap()
216                    }
217                }
218
219                let (mut reader, _) = crate::read(MAGIC, Cursor::new(file)).unwrap();
220                let err = Blob::read(&mut reader).err().unwrap();
221                assert!(
222                    err.to_string().contains(&format!(
223                        "blob length {len} is above the maximum of {MAX_BLOB_LENGTH} bytes"
224                    )),
225                    "{err}"
226                );
227            }
228        }
229    }
230
231    #[test]
232    fn blob_iter_reads_until_end_of_input() {
233        for count in [0u64, 1, 3] {
234            let mut blobs = Vec::new();
235            for value in 0..count {
236                Blob::encode(&value, BlobEncoding::Bcs)
237                    .unwrap()
238                    .write(&mut blobs)
239                    .unwrap();
240            }
241            let mut file = MAGIC.to_be_bytes().to_vec();
242            file.push(StorageFormat::Blob.into());
243            file.push(FileCompression::None.into());
244            file.extend(blobs);
245
246            let values: Vec<u64> = make_iterator(MAGIC, Cursor::new(file)).unwrap().collect();
247            assert_eq!(values, (0..count).collect::<Vec<_>>());
248        }
249    }
250
251    #[test]
252    fn blob_iter_stops_at_length_beyond_input() {
253        for compression in [FileCompression::None, FileCompression::Zstd] {
254            let mut blobs = Vec::new();
255            Blob::encode(&7u64, BlobEncoding::Bcs)
256                .unwrap()
257                .write(&mut blobs)
258                .unwrap();
259            blobs.extend(header(4));
260            blobs.extend_from_slice(&[0u8; 3]);
261
262            let mut file = MAGIC.to_be_bytes().to_vec();
263            file.push(StorageFormat::Blob.into());
264            file.push(compression.into());
265            match compression {
266                FileCompression::None => file.extend(blobs),
267                FileCompression::Zstd => {
268                    FileCompression::zstd_compress(&mut Cursor::new(blobs), &mut file).unwrap()
269                }
270            }
271
272            let values: Vec<u64> = make_iterator(MAGIC, Cursor::new(file)).unwrap().collect();
273            assert_eq!(values, vec![7]);
274        }
275    }
276}
iota_storage/blob.rs

iota_storage/
blob.rs
