Skip to main content

iota_storage/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5#![allow(dead_code)]
6
7use std::{
8    fs,
9    fs::File,
10    io,
11    io::{BufReader, Read, Write},
12    path::PathBuf,
13    sync::Arc,
14};
15
16use anyhow::{Result, anyhow};
17use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
18use bytes::{Buf, Bytes};
19use fastcrypto::hash::{HashFunction, Sha3_256};
20use iota_types::{
21    committee::Committee,
22    messages_checkpoint::{CertifiedCheckpointSummary, VerifiedCheckpoint},
23    storage::WriteStore,
24};
25use num_enum::{IntoPrimitive, TryFromPrimitive};
26use serde::{Deserialize, Serialize, de::DeserializeOwned};
27use tracing::debug;
28
29use crate::blob::BlobIter;
30
31pub mod blob;
32pub mod http_key_value_store;
33pub mod key_value_store;
34pub mod key_value_store_metrics;
35pub mod mutex_table;
36pub mod object_store;
37pub mod package_object_cache;
38pub mod sharded_lru;
39pub mod write_path_pending_tx_log;
40
41pub const SHA3_BYTES: usize = 32;
42
43#[derive(
44    Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, TryFromPrimitive, IntoPrimitive,
45)]
46#[repr(u8)]
47pub enum StorageFormat {
48    Blob = 0,
49}
50
51#[derive(
52    Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize, TryFromPrimitive, IntoPrimitive,
53)]
54#[repr(u8)]
55pub enum FileCompression {
56    None = 0,
57    Zstd,
58}
59
60impl FileCompression {
61    pub fn zstd_compress<R: Read, W: Write>(reader: &mut R, writer: &mut W) -> io::Result<()> {
62        // TODO: Add zstd compression level as function argument
63        let mut encoder = zstd::Encoder::new(writer, 1)?;
64        io::copy(reader, &mut encoder)?;
65        encoder.finish()?;
66        Ok(())
67    }
68    pub fn compress(&self, source: &std::path::Path) -> io::Result<()> {
69        match self {
70            FileCompression::Zstd => {
71                let mut input = File::open(source)?;
72                let tmp_file_name = source.with_extension("tmp");
73                let mut output = File::create(&tmp_file_name)?;
74                Self::zstd_compress(&mut input, &mut output)?;
75                fs::rename(tmp_file_name, source)?;
76            }
77            FileCompression::None => {}
78        }
79        Ok(())
80    }
81    pub fn decompress(&self, source: &PathBuf) -> Result<Box<dyn Read>> {
82        let file = File::open(source)?;
83        let res: Box<dyn Read> = match self {
84            FileCompression::Zstd => Box::new(zstd::stream::Decoder::new(file)?),
85            FileCompression::None => Box::new(BufReader::new(file)),
86        };
87        Ok(res)
88    }
89    pub fn bytes_decompress(&self, bytes: Bytes) -> Result<Box<dyn Read>> {
90        let res: Box<dyn Read> = match self {
91            FileCompression::Zstd => Box::new(zstd::stream::Decoder::new(bytes.reader())?),
92            FileCompression::None => Box::new(BufReader::new(bytes.reader())),
93        };
94        Ok(res)
95    }
96}
97
98pub fn compute_sha3_checksum_for_bytes(bytes: Bytes) -> Result<[u8; 32]> {
99    let mut hasher = Sha3_256::default();
100    io::copy(&mut bytes.reader(), &mut hasher)?;
101    Ok(hasher.finalize().digest)
102}
103
104pub fn compute_sha3_checksum_for_file(file: &mut File) -> Result<[u8; 32]> {
105    let mut hasher = Sha3_256::default();
106    io::copy(file, &mut hasher)?;
107    Ok(hasher.finalize().digest)
108}
109
110pub fn compute_sha3_checksum(source: &std::path::Path) -> Result<[u8; 32]> {
111    let mut file = fs::File::open(source)?;
112    compute_sha3_checksum_for_file(&mut file)
113}
114
115pub fn compress<R: Read, W: Write>(reader: &mut R, writer: &mut W) -> Result<()> {
116    let magic = reader.read_u32::<BigEndian>()?;
117    writer.write_u32::<BigEndian>(magic)?;
118    let storage_format = reader.read_u8()?;
119    writer.write_u8(storage_format)?;
120    let file_compression = FileCompression::try_from(reader.read_u8()?)?;
121    writer.write_u8(file_compression.into())?;
122    match file_compression {
123        FileCompression::Zstd => {
124            FileCompression::zstd_compress(reader, writer)?;
125        }
126        FileCompression::None => {}
127    }
128    Ok(())
129}
130
131pub fn read<R: Read + 'static>(
132    expected_magic: u32,
133    mut reader: R,
134) -> Result<(Box<dyn Read>, StorageFormat)> {
135    let magic = reader.read_u32::<BigEndian>()?;
136    if magic != expected_magic {
137        Err(anyhow!(
138            "Unexpected magic string in file: {magic:?}, expected: {expected_magic:?}"
139        ))
140    } else {
141        let storage_format = StorageFormat::try_from(reader.read_u8()?)?;
142        let file_compression = FileCompression::try_from(reader.read_u8()?)?;
143        let reader: Box<dyn Read> = match file_compression {
144            FileCompression::Zstd => Box::new(zstd::stream::Decoder::new(reader)?),
145            FileCompression::None => Box::new(BufReader::new(reader)),
146        };
147        Ok((reader, storage_format))
148    }
149}
150
151pub fn make_iterator<T: DeserializeOwned, R: Read + 'static>(
152    expected_magic: u32,
153    reader: R,
154) -> Result<impl Iterator<Item = T>> {
155    let (reader, storage_format) = read(expected_magic, reader)?;
156    match storage_format {
157        StorageFormat::Blob => Ok(BlobIter::new(reader)),
158    }
159}
160
161#[expect(clippy::result_large_err)]
162pub fn verify_checkpoint_with_committee(
163    committee: Arc<Committee>,
164    current: &VerifiedCheckpoint,
165    checkpoint: CertifiedCheckpointSummary,
166) -> Result<VerifiedCheckpoint, CertifiedCheckpointSummary> {
167    assert_eq!(
168        checkpoint.sequence_number(),
169        current.sequence_number().checked_add(1).unwrap()
170    );
171
172    if Some(*current.digest()) != checkpoint.previous_digest {
173        debug!(
174            current_checkpoint_seq = current.sequence_number(),
175            current_digest =% current.digest(),
176            checkpoint_seq = checkpoint.sequence_number(),
177            checkpoint_digest =% checkpoint.digest(),
178            checkpoint_previous_digest =? checkpoint.previous_digest,
179            "checkpoint not on same chain"
180        );
181        return Err(checkpoint);
182    }
183
184    let current_epoch = current.epoch();
185    if checkpoint.epoch() != current_epoch
186        && checkpoint.epoch() != current_epoch.checked_add(1).unwrap()
187    {
188        debug!(
189            checkpoint_seq = checkpoint.sequence_number(),
190            checkpoint_epoch = checkpoint.epoch(),
191            current_checkpoint_seq = current.sequence_number(),
192            current_epoch = current_epoch,
193            "cannot verify checkpoint with too high of an epoch",
194        );
195        return Err(checkpoint);
196    }
197
198    if checkpoint.epoch() == current_epoch.checked_add(1).unwrap()
199        && current.next_epoch_committee().is_none()
200    {
201        debug!(
202            checkpoint_seq = checkpoint.sequence_number(),
203            checkpoint_epoch = checkpoint.epoch(),
204            current_checkpoint_seq = current.sequence_number(),
205            current_epoch = current_epoch,
206            "next checkpoint claims to be from the next epoch but the latest verified \
207            checkpoint does not indicate that it is the last checkpoint of an epoch"
208        );
209        return Err(checkpoint);
210    }
211
212    checkpoint
213        .verify_authority_signatures(&committee)
214        .map_err(|e| {
215            debug!("error verifying checkpoint: {e}");
216            checkpoint.clone()
217        })?;
218    Ok(VerifiedCheckpoint::new_unchecked(checkpoint))
219}
220
221#[expect(clippy::result_large_err)]
222pub fn verify_checkpoint<S>(
223    current: &VerifiedCheckpoint,
224    store: S,
225    checkpoint: CertifiedCheckpointSummary,
226) -> Result<VerifiedCheckpoint, CertifiedCheckpointSummary>
227where
228    S: WriteStore,
229{
230    let committee = store.get_committee(checkpoint.epoch()).unwrap_or_else(|| {
231        panic!(
232            "BUG: should have committee for epoch {} before we try to verify checkpoint {}",
233            checkpoint.epoch(),
234            checkpoint.sequence_number()
235        )
236    });
237
238    verify_checkpoint_with_committee(committee, current, checkpoint)
239}