| 1 | |
|
| 2 | |
|
| 3 | |
|
| 4 | |
|
| 5 | |
|
| 6 | |
|
| 7 | |
|
| 8 | |
|
| 9 | |
|
| 10 | |
|
| 11 | |
|
| 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 | |
|
| 20 | | pub const VERSION: u32 = 1; |
| 21 | |
|
| 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 | |
|
| 35 | | #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] |
| 36 | | pub struct Chunk(pub String, pub u32); |
| 37 | |
|
| 38 | | impl Chunk { |
| 39 | 2002× | pub fn hash(&self) -> &str { |
| 40 | 2002× | &self.0 |
| 41 | 2002× | } |
| 42 | |
|
| 43 | 1666× | pub fn len(&self) -> usize { |
| 44 | 1666× | self.1 as usize |
| 45 | 1666× | } |
| 46 | |
|
| 47 | 246× | pub fn is_empty(&self) -> bool { |
| 48 | 246× | self.1 == 0 |
| 49 | 246× | } |
| 50 | |
|
| 51 | | |
| 52 | | |
| 53 | 885× | pub fn is_zeros(&self) -> bool { |
| 54 | 885× | self.0 == zeros_hash(self.len()) |
| 55 | 885× | } |
| 56 | | } |
| 57 | |
|
| 58 | |
|
| 59 | | #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] |
| 60 | | pub struct Manifest { |
| 61 | | pub version: u32, |
| 62 | | |
| 63 | | pub hash: String, |
| 64 | | pub size: u64, |
| 65 | | |
| 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 | | |
| 73 | 11× | pub fn needed(&self) -> Vec<&str> { |
| 74 | 11× | let mut seen = HashSet::new(); |
| 75 | 370× | self.chunks.iter().filter(|c| !c.is_zeros() && seen.insert(c.hash())).map(|c| c.hash()).collect() |
| 76 | 11× | } |
| 77 | |
|
| 78 | | |
| 79 | | |
| 80 | 10× | pub fn check(&self) -> Result<()> { |
| 81 | 10× | if self.version != VERSION { |
| 82 | 0× | return Err(Error::Bad(format!("image manifest version {} is not known; this reads {VERSION}", self.version))); |
| 83 | 10× | } |
| 84 | 10× | if !is_hash(&self.hash) { |
| 85 | 1× | return Err(Error::Bad(format!("{:?} is not an image hash", self.hash))); |
| 86 | 9× | } |
| 87 | 9× | if self.size > MAX_IMAGE { |
| 88 | 0× | return Err(Error::Bad(format!("the image is {} GiB; at most {} GiB", self.size >> 30, MAX_IMAGE >> 30))); |
| 89 | 9× | } |
| 90 | 9× | let mut total = 0u64; |
| 91 | 246× | for c in &self.chunks { |
| 92 | 246× | if !is_hash(c.hash()) || c.is_empty() || c.len() > MAX_SIZE { |
| 93 | 0× | return Err(Error::Bad(format!("chunk {:?} of {} bytes is not a chunk", c.hash(), c.len()))); |
| 94 | 246× | } |
| 95 | 246× | total += c.1 as u64; |
| 96 | | } |
| 97 | 9× | if total != self.size { |
| 98 | 1× | return Err(Error::Bad(format!("the chunks hold {total} bytes; the image is {}", self.size))); |
| 99 | 8× | } |
| 100 | 8× | Ok(()) |
| 101 | 10× | } |
| 102 | | } |
| 103 | |
|
| 104 | |
|
| 105 | 325× | pub fn is_hash(s: &str) -> bool { |
| 106 | 20736× | s.len() == 64 && s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) |
| 107 | 325× | } |
| 108 | |
|
| 109 | |
|
| 110 | 885× | pub fn zeros_hash(len: usize) -> String { |
| 111 | | use std::sync::Mutex; |
| 112 | | static MEMO: Mutex<Vec<(usize, String)>> = Mutex::new(Vec::new()); |
| 113 | 885× | let mut memo = MEMO.lock().unwrap_or_else(|p| p.into_inner()); |
| 114 | 16200× | if let Some((_, h)) = memo.iter().find(|(l, _)| *l == len) { |
| 115 | 730× | return h.clone(); |
| 116 | 155× | } |
| 117 | 155× | let h = blake3::hash(&vec![0u8; len]).to_hex().to_string(); |
| 118 | 155× | if memo.len() < 64 { |
| 119 | 155× | memo.push((len, h.clone())); |
| 120 | 155× | } |
| 121 | 155× | h |
| 122 | 885× | } |
| 123 | |
|
| 124 | |
|
| 125 | |
|
| 126 | |
|
| 127 | 9× | 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; |
| 129 | 9× | let mut buf: Vec<u8> = Vec::with_capacity(FILL + MAX_SIZE); |
| 130 | 9× | let (mut pos, mut eof, mut size) = (0usize, false, 0u64); |
| 131 | 9× | let mut whole = blake3::Hasher::new(); |
| 132 | 9× | let mut chunks = Vec::new(); |
| 133 | | loop { |
| 134 | 279× | if !eof && buf.len() - pos < MAX_SIZE { |
| 135 | 9× | buf.drain(..pos); |
| 136 | 9× | pos = 0; |
| 137 | 9× | let start = buf.len(); |
| 138 | 9× | buf.resize(start + FILL, 0); |
| 139 | 9× | let mut got = 0; |
| 140 | 18× | while got < FILL { |
| 141 | 18× | match r.read(&mut buf[start + got..])? { |
| 142 | | 0 => { |
| 143 | 9× | eof = true; |
| 144 | 9× | break; |
| 145 | | } |
| 146 | 9× | n => got += n, |
| 147 | | } |
| 148 | | } |
| 149 | 9× | buf.truncate(start + got); |
| 150 | 270× | } |
| 151 | 279× | if pos == buf.len() { |
| 152 | 9× | break; |
| 153 | 270× | } |
| 154 | 270× | let n = cut_point(&buf[pos..]); |
| 155 | 270× | let bytes = &buf[pos..pos + n]; |
| 156 | 270× | whole.update(bytes); |
| 157 | 270× | let c = Chunk(blake3::hash(bytes).to_hex().to_string(), n as u32); |
| 158 | 270× | each(&c, bytes)?; |
| 159 | 270× | chunks.push(c); |
| 160 | 270× | size += n as u64; |
| 161 | 270× | pos += n; |
| 162 | | } |
| 163 | 9× | Ok(Manifest { version: VERSION, hash: whole.finalize().to_hex().to_string(), size, arch, chunks }) |
| 164 | 9× | } |
| 165 | |
|
| 166 | |
|
| 167 | 2× | pub fn scan_file(path: &Path, arch: Option<String>, each: impl FnMut(&Chunk, &[u8]) -> Result<()>) -> Result<Manifest> { |
| 168 | 2× | scan(std::io::BufReader::with_capacity(1 << 20, std::fs::File::open(path)?), arch, each) |
| 169 | 2× | } |
| 170 | |
|
| 171 | |
|
| 172 | |
|
| 173 | |
|
| 174 | |
|
| 175 | 7× | pub fn assemble(m: &Manifest, out: &Path, mut fetch: impl FnMut(&Chunk) -> Result<Vec<u8>>) -> Result<()> { |
| 176 | 7× | m.check()?; |
| 177 | 7× | let tmp = out.with_extension("part"); |
| 178 | 7× | let mut f = std::fs::File::create(&tmp)?; |
| 179 | 7× | let mut whole = blake3::Hasher::new(); |
| 180 | 7× | let mut last: Option<(String, Vec<u8>)> = None; |
| 181 | 205× | for c in &m.chunks { |
| 182 | 205× | if c.is_zeros() { |
| 183 | 11× | whole.update(&vec![0u8; c.len()]); |
| 184 | 11× | f.seek(SeekFrom::Current(c.len() as i64))?; |
| 185 | 11× | continue; |
| 186 | 194× | } |
| 187 | 193× | let bytes = match &last { |
| 188 | 187× | Some((h, b)) if h == c.hash() => b.clone(), |
| 189 | | _ => { |
| 190 | 194× | let b = fetch(c)?; |
| 191 | 194× | if b.len() != c.len() || blake3::hash(&b).to_hex().as_str() != c.hash() { |
| 192 | 1× | let _ = std::fs::remove_file(&tmp); |
| 193 | 1× | return Err(Error::Bad(format!("chunk {} came back with other bytes", &c.hash()[..12]))); |
| 194 | 193× | } |
| 195 | 193× | last = Some((c.0.clone(), b.clone())); |
| 196 | 193× | b |
| 197 | | } |
| 198 | | }; |
| 199 | 193× | whole.update(&bytes); |
| 200 | 193× | f.write_all(&bytes)?; |
| 201 | | } |
| 202 | 6× | f.set_len(m.size)?; |
| 203 | 6× | f.sync_all()?; |
| 204 | 6× | drop(f); |
| 205 | 6× | let got = whole.finalize().to_hex().to_string(); |
| 206 | 6× | if got != m.hash { |
| 207 | 0× | let _ = std::fs::remove_file(&tmp); |
| 208 | 0× | return Err(Error::Bad(format!("the image came out as {}, not {}", &got[..12], &m.hash[..12]))); |
| 209 | 6× | } |
| 210 | 6× | std::fs::rename(&tmp, out)?; |
| 211 | 6× | Ok(()) |
| 212 | 7× | } |
| 213 | |
|
| 214 | | #[cfg(test)] |
| 215 | | mod tests { |
| 216 | | use super::*; |
| 217 | |
|
| 218 | 2× | fn image(seed: u8) -> Vec<u8> { |
| 219 | | |
| 220 | | |
| 221 | 600000× | let mut v: Vec<u8> = (0..300_000u32).map(|i| ((i * 7 + seed as u32) % 251) as u8).collect(); |
| 222 | 2× | v.extend(std::iter::repeat_n(0u8, 3 << 20)); |
| 223 | 400000× | v.extend((0..200_000u32).map(|i| ((i * 13 + seed as u32) % 241) as u8)); |
| 224 | 2× | v |
| 225 | 2× | } |
| 226 | |
|
| 227 | | #[test] |
| 228 | 1× | fn a_scanned_image_comes_back_whole_and_zeros_are_never_fetched() { |
| 229 | 1× | let data = image(1); |
| 230 | 1× | let mut stored = std::collections::HashMap::new(); |
| 231 | 14× | let m = scan(&data[..], Some("x86_64".into()), |c, b| { |
| 232 | 14× | stored.insert(c.hash().to_string(), b.to_vec()); |
| 233 | 14× | Ok(()) |
| 234 | 14× | }) |
| 235 | 1× | .unwrap(); |
| 236 | 1× | m.check().unwrap(); |
| 237 | 1× | assert_eq!(m.hash, blake3::hash(&data).to_hex().to_string(), "named as pipelines pin it"); |
| 238 | 1× | assert_eq!(m.size, data.len() as u64); |
| 239 | 1× | assert_eq!( |
| 240 | 14× | m.chunks.iter().map(|c| c.len()).collect::<Vec<_>>(), |
| 241 | 14× | oxy_core::chunk::chunks(&data).iter().map(|c| c.len()).collect::<Vec<_>>(), |
| 242 | | "the same cuts as Oxygen's chunker over the whole" |
| 243 | | ); |
| 244 | 3× | assert!(m.chunks.iter().any(|c| c.is_zeros())); |
| 245 | 1× | let dir = tempfile::tempdir().unwrap(); |
| 246 | 1× | let out = dir.path().join("img.ext4"); |
| 247 | 1× | let mut fetched = 0; |
| 248 | 3× | assemble(&m, &out, |c| { |
| 249 | 3× | fetched += 1; |
| 250 | 3× | assert!(!c.is_zeros()); |
| 251 | 3× | Ok(stored[c.hash()].clone()) |
| 252 | 3× | }) |
| 253 | 1× | .unwrap(); |
| 254 | 1× | assert_eq!(std::fs::read(&out).unwrap(), data); |
| 255 | 1× | assert!(fetched <= m.needed().len() + 1); |
| 256 | |
|
| 257 | | |
| 258 | 1× | let mut changed = data.clone(); |
| 259 | 1× | changed[100] ^= 0xff; |
| 260 | 14× | let m2 = scan(&changed[..], None, |_, _| Ok(())).unwrap(); |
| 261 | 14× | let have: HashSet<&str> = m.chunks.iter().map(|c| c.hash()).collect(); |
| 262 | 3× | let new = m2.needed().into_iter().filter(|h| !have.contains(h)).count(); |
| 263 | 1× | assert!(new <= 2, "{new} new chunks for a one-byte change"); |
| 264 | 1× | } |
| 265 | |
|
| 266 | | #[test] |
| 267 | 1× | fn bad_bytes_and_bad_manifests_are_refused() { |
| 268 | 1× | let data = image(2); |
| 269 | 14× | let m = scan(&data[..], None, |_, _| Ok(())).unwrap(); |
| 270 | 1× | let dir = tempfile::tempdir().unwrap(); |
| 271 | 1× | let out = dir.path().join("img.ext4"); |
| 272 | 1× | let e = assemble(&m, &out, |c| Ok(vec![1u8; c.len()])).unwrap_err(); |
| 273 | 1× | assert!(e.to_string().contains("other bytes"), "{e}"); |
| 274 | 1× | assert!(!out.exists() && !out.with_extension("part").exists()); |
| 275 | 1× | let mut short = m.clone(); |
| 276 | 1× | short.size += 1; |
| 277 | 1× | assert!(short.check().unwrap_err().to_string().contains("chunks hold")); |
| 278 | 1× | let mut odd = m.clone(); |
| 279 | 1× | odd.hash = "nope".into(); |
| 280 | 1× | assert!(odd.check().is_err()); |
| 281 | 1× | } |
| 282 | | } |