OXYGENHUB

Coverage / crates/jet/src/images.rs 95.8%

182 of 190 lines run, 51 of 66 functions · run #181 on main · run missed

1
//! Images by hash (oxygenflow#39, docs/jetstream.md §8).
2
//!
3
//! A rootfs image is named by the BLAKE3 of its bytes, as pipelines pin it
4
//! (`cas://blob/HASH`). To move it, it is cut into chunks with Oxygen's
5
//! FastCDC (32 KiB to 256 KiB, `oxy_core::chunk`), each named by its own
6
//! BLAKE3, and described by a [`Manifest`]: the image's hash, size and
7
//! architecture, and its chunks in order. A sender sends only the chunks a
8
//! receiver lacks; a receiver writes the image back, leaving runs of zeros
9
//! as holes without fetching them, and checks the whole hash before using
10
//! it. Chunk boundaries depend only on content, so an image rebuilt with a
11
//! small change shares nearly all its chunks with the one before.
12
13
use oxy_core::chunk::{cut_point, MAX_SIZE};
14
use serde::{Deserialize, Serialize};
15
use std::collections::HashSet;
16
use std::io::{Read, Seek, SeekFrom, Write};
17
use std::path::Path;
18
19
/// The manifest format this crate writes and reads.
20
pub const VERSION: u32 = 1;
21
/// Largest image accepted: 64 GiB.
22
pub const MAX_IMAGE: u64 = 64 << 30;
23
24
#[derive(Debug, thiserror::Error)]
25
pub enum Error {
26
    #[error("{0}")]
27
    Io(#[from] std::io::Error),
28
    #[error("{0}")]
29
    Bad(String),
30
}
31
32
pub type Result<T> = std::result::Result<T, Error>;
33
34
/// One chunk: its BLAKE3 (hex) and length.
35
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
36
pub struct Chunk(pub String, pub u32);
37
38
impl Chunk {
392002×
    pub fn hash(&self) -> &str {
402002×
        &self.0
412002×
    }
42
431666×
    pub fn len(&self) -> usize {
441666×
        self.1 as usize
451666×
    }
46
47246×
    pub fn is_empty(&self) -> bool {
48246×
        self.1 == 0
49246×
    }
50
51
    /// Whether this is a chunk of zeros, which a receiver writes as a hole
52
    /// without fetching it.
53885×
    pub fn is_zeros(&self) -> bool {
54885×
        self.0 == zeros_hash(self.len())
55885×
    }
56
}
57
58
/// An image: its hash, size, architecture and chunks.
59
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
60
pub struct Manifest {
61
    pub version: u32,
62
    /// BLAKE3 of the whole image, hex: its name.
63
    pub hash: String,
64
    pub size: u64,
65
    /// `x86_64`, `riscv64`, …, as its import recorded (OxygenFlow spec §11.2).
66
    #[serde(default, skip_serializing_if = "Option::is_none")]
67
    pub arch: Option<String>,
68
    pub chunks: Vec<Chunk>,
69
}
70
71
impl Manifest {
72
    /// The distinct chunks a receiver must fetch: all but zeros.
7311×
    pub fn needed(&self) -> Vec<&str> {
7411×
        let mut seen = HashSet::new();
75370×
        self.chunks.iter().filter(|c| !c.is_zeros() && seen.insert(c.hash())).map(|c| c.hash()).collect()
7611×
    }
77
78
    /// Check its shape: a known version, hex names, chunks in the
79
    /// chunker's bounds that add up to the size.
8010×
    pub fn check(&self) -> Result<()> {
8110×
        if self.version != VERSION {
820×
            return Err(Error::Bad(format!("image manifest version {} is not known; this reads {VERSION}", self.version)));
8310×
        }
8410×
        if !is_hash(&self.hash) {
851×
            return Err(Error::Bad(format!("{:?} is not an image hash", self.hash)));
869×
        }
879×
        if self.size > MAX_IMAGE {
880×
            return Err(Error::Bad(format!("the image is {} GiB; at most {} GiB", self.size >> 30, MAX_IMAGE >> 30)));
899×
        }
909×
        let mut total = 0u64;
91246×
        for c in &self.chunks {
92246×
            if !is_hash(c.hash()) || c.is_empty() || c.len() > MAX_SIZE {
930×
                return Err(Error::Bad(format!("chunk {:?} of {} bytes is not a chunk", c.hash(), c.len())));
94246×
            }
95246×
            total += c.1 as u64;
96
        }
979×
        if total != self.size {
981×
            return Err(Error::Bad(format!("the chunks hold {total} bytes; the image is {}", self.size)));
998×
        }
1008×
        Ok(())
10110×
    }
102
}
103
104
/// Whether `s` is 64 lower-case hex digits.
105325×
pub fn is_hash(s: &str) -> bool {
10620736×
    s.len() == 64 && s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
107325×
}
108
109
/// The BLAKE3 of `len` zero bytes, hex; remembered per length.
110885×
pub fn zeros_hash(len: usize) -> String {
111
    use std::sync::Mutex;
112
    static MEMO: Mutex<Vec<(usize, String)>> = Mutex::new(Vec::new());
113885×
    let mut memo = MEMO.lock().unwrap_or_else(|p| p.into_inner());
11416200×
    if let Some((_, h)) = memo.iter().find(|(l, _)| *l == len) {
115730×
        return h.clone();
116155×
    }
117155×
    let h = blake3::hash(&vec![0u8; len]).to_hex().to_string();
118155×
    if memo.len() < 64 {
119155×
        memo.push((len, h.clone()));
120155×
    }
121155×
    h
122885×
}
123
124
/// Cut `r` into chunks as Oxygen's chunker would cut it whole, calling
125
/// `each` with every chunk's bytes, and answer the manifest. Reads in
126
/// bounded memory: the chunker never looks past [`MAX_SIZE`].
1279×
pub fn scan<R: Read>(mut r: R, arch: Option<String>, mut each: impl FnMut(&Chunk, &[u8]) -> Result<()>) -> Result<Manifest> {
128
    const FILL: usize = 8 << 20;
1299×
    let mut buf: Vec<u8> = Vec::with_capacity(FILL + MAX_SIZE);
1309×
    let (mut pos, mut eof, mut size) = (0usize, false, 0u64);
1319×
    let mut whole = blake3::Hasher::new();
1329×
    let mut chunks = Vec::new();
133
    loop {
134279×
        if !eof && buf.len() - pos < MAX_SIZE {
1359×
            buf.drain(..pos);
1369×
            pos = 0;
1379×
            let start = buf.len();
1389×
            buf.resize(start + FILL, 0);
1399×
            let mut got = 0;
14018×
            while got < FILL {
14118×
                match r.read(&mut buf[start + got..])? {
142
                    0 => {
1439×
                        eof = true;
1449×
                        break;
145
                    }
1469×
                    n => got += n,
147
                }
148
            }
1499×
            buf.truncate(start + got);
150270×
        }
151279×
        if pos == buf.len() {
1529×
            break;
153270×
        }
154270×
        let n = cut_point(&buf[pos..]);
155270×
        let bytes = &buf[pos..pos + n];
156270×
        whole.update(bytes);
157270×
        let c = Chunk(blake3::hash(bytes).to_hex().to_string(), n as u32);
158270×
        each(&c, bytes)?;
159270×
        chunks.push(c);
160270×
        size += n as u64;
161270×
        pos += n;
162
    }
1639×
    Ok(Manifest { version: VERSION, hash: whole.finalize().to_hex().to_string(), size, arch, chunks })
1649×
}
165
166
/// [`scan`] of a file.
1672×
pub fn scan_file(path: &Path, arch: Option<String>, each: impl FnMut(&Chunk, &[u8]) -> Result<()>) -> Result<Manifest> {
1682×
    scan(std::io::BufReader::with_capacity(1 << 20, std::fs::File::open(path)?), arch, each)
1692×
}
170
171
/// Write image `m` to `out`, getting each chunk's bytes from `fetch`
172
/// (zeros are holes, never fetched; a chunk seen before is fetched once
173
/// per run of repeats), checking every chunk and then the whole hash. The
174
/// image is written beside `out` and renamed into place only when whole.
1757×
pub fn assemble(m: &Manifest, out: &Path, mut fetch: impl FnMut(&Chunk) -> Result<Vec<u8>>) -> Result<()> {
1767×
    m.check()?;
1777×
    let tmp = out.with_extension("part");
1787×
    let mut f = std::fs::File::create(&tmp)?;
1797×
    let mut whole = blake3::Hasher::new();
1807×
    let mut last: Option<(String, Vec<u8>)> = None;
181205×
    for c in &m.chunks {
182205×
        if c.is_zeros() {
18311×
            whole.update(&vec![0u8; c.len()]);
18411×
            f.seek(SeekFrom::Current(c.len() as i64))?;
18511×
            continue;
186194×
        }
187193×
        let bytes = match &last {
188187×
            Some((h, b)) if h == c.hash() => b.clone(),
189
            _ => {
190194×
                let b = fetch(c)?;
191194×
                if b.len() != c.len() || blake3::hash(&b).to_hex().as_str() != c.hash() {
1921×
                    let _ = std::fs::remove_file(&tmp);
1931×
                    return Err(Error::Bad(format!("chunk {} came back with other bytes", &c.hash()[..12])));
194193×
                }
195193×
                last = Some((c.0.clone(), b.clone()));
196193×
                b
197
            }
198
        };
199193×
        whole.update(&bytes);
200193×
        f.write_all(&bytes)?;
201
    }
2026×
    f.set_len(m.size)?;
2036×
    f.sync_all()?;
2046×
    drop(f);
2056×
    let got = whole.finalize().to_hex().to_string();
2066×
    if got != m.hash {
2070×
        let _ = std::fs::remove_file(&tmp);
2080×
        return Err(Error::Bad(format!("the image came out as {}, not {}", &got[..12], &m.hash[..12])));
2096×
    }
2106×
    std::fs::rename(&tmp, out)?;
2116×
    Ok(())
2127×
}
213
214
#[cfg(test)]
215
mod tests {
216
    use super::*;
217
2182×
    fn image(seed: u8) -> Vec<u8> {
219
        // Some text-like bytes, a long run of zeros (as an ext4 image has),
220
        // then more bytes.
221600000×
        let mut v: Vec<u8> = (0..300_000u32).map(|i| ((i * 7 + seed as u32) % 251) as u8).collect();
2222×
        v.extend(std::iter::repeat_n(0u8, 3 << 20));
223400000×
        v.extend((0..200_000u32).map(|i| ((i * 13 + seed as u32) % 241) as u8));
2242×
        v
2252×
    }
226
227
    #[test]
2281×
    fn a_scanned_image_comes_back_whole_and_zeros_are_never_fetched() {
2291×
        let data = image(1);
2301×
        let mut stored = std::collections::HashMap::new();
23114×
        let m = scan(&data[..], Some("x86_64".into()), |c, b| {
23214×
            stored.insert(c.hash().to_string(), b.to_vec());
23314×
            Ok(())
23414×
        })
2351×
        .unwrap();
2361×
        m.check().unwrap();
2371×
        assert_eq!(m.hash, blake3::hash(&data).to_hex().to_string(), "named as pipelines pin it");
2381×
        assert_eq!(m.size, data.len() as u64);
2391×
        assert_eq!(
24014×
            m.chunks.iter().map(|c| c.len()).collect::<Vec<_>>(),
24114×
            oxy_core::chunk::chunks(&data).iter().map(|c| c.len()).collect::<Vec<_>>(),
242
            "the same cuts as Oxygen's chunker over the whole"
243
        );
2443×
        assert!(m.chunks.iter().any(|c| c.is_zeros()));
2451×
        let dir = tempfile::tempdir().unwrap();
2461×
        let out = dir.path().join("img.ext4");
2471×
        let mut fetched = 0;
2483×
        assemble(&m, &out, |c| {
2493×
            fetched += 1;
2503×
            assert!(!c.is_zeros());
2513×
            Ok(stored[c.hash()].clone())
2523×
        })
2531×
        .unwrap();
2541×
        assert_eq!(std::fs::read(&out).unwrap(), data);
2551×
        assert!(fetched <= m.needed().len() + 1);
256
257
        // A rebuild with a small change shares nearly every chunk.
2581×
        let mut changed = data.clone();
2591×
        changed[100] ^= 0xff;
26014×
        let m2 = scan(&changed[..], None, |_, _| Ok(())).unwrap();
26114×
        let have: HashSet<&str> = m.chunks.iter().map(|c| c.hash()).collect();
2623×
        let new = m2.needed().into_iter().filter(|h| !have.contains(h)).count();
2631×
        assert!(new <= 2, "{new} new chunks for a one-byte change");
2641×
    }
265
266
    #[test]
2671×
    fn bad_bytes_and_bad_manifests_are_refused() {
2681×
        let data = image(2);
26914×
        let m = scan(&data[..], None, |_, _| Ok(())).unwrap();
2701×
        let dir = tempfile::tempdir().unwrap();
2711×
        let out = dir.path().join("img.ext4");
2721×
        let e = assemble(&m, &out, |c| Ok(vec![1u8; c.len()])).unwrap_err();
2731×
        assert!(e.to_string().contains("other bytes"), "{e}");
2741×
        assert!(!out.exists() && !out.with_extension("part").exists());
2751×
        let mut short = m.clone();
2761×
        short.size += 1;
2771×
        assert!(short.check().unwrap_err().to_string().contains("chunks hold"));
2781×
        let mut odd = m.clone();
2791×
        odd.hash = "nope".into();
2801×
        assert!(odd.check().is_err());
2811×
    }
282
}