assetpack_core/
file_reader.rs

1use std::{
2  collections::BTreeSet,
3  io::{self, BufReader, Cursor, Write},
4};
5
6use sha3::{Digest, Sha3_256};
7
8use crate::{
9  Error, Hash32, ObjectKind, ObjectSource, RecipeData, Result, TransformDecoderRegistry, parse_recipe, recipe::preflight_recipe_limits,
10};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct FileReadLimits {
14  pub max_chunk_count: u64,
15  pub max_stored_stream_bytes: u64,
16  pub max_original_file_bytes: u64,
17}
18
19impl Default for FileReadLimits {
20  fn default() -> Self {
21    Self {
22      max_chunk_count: 1_000_000,
23      max_stored_stream_bytes: 256 * 1024 * 1024,
24      max_original_file_bytes: 256 * 1024 * 1024,
25    }
26  }
27}
28
29pub struct FileReader<'a, S: ?Sized> {
30  source: &'a S,
31  decoders: &'a TransformDecoderRegistry,
32  limits: FileReadLimits,
33}
34
35impl<'a, S: ObjectSource + ?Sized> FileReader<'a, S> {
36  pub fn new(source: &'a S, decoders: &'a TransformDecoderRegistry, limits: FileReadLimits) -> Self {
37    Self { source, decoders, limits }
38  }
39
40  pub async fn resolve_closure(&self, recipe_hash: Hash32) -> Result<BTreeSet<Hash32>> {
41    let recipe = self.load_recipe(recipe_hash).await?;
42    let mut hashes = BTreeSet::from([recipe_hash]);
43    for (hash, expected_len) in &recipe.chunks {
44      let object = self.source.read_object(hash).await?.ok_or(Error::ObjectNotFound)?;
45      self.validate_chunk(&object.hash, object.kind, object.bytes.len() as u64, u64::from(*expected_len))?;
46      hashes.insert(*hash);
47    }
48    Ok(hashes)
49  }
50
51  pub async fn read_stored_stream(&self, recipe_hash: Hash32) -> Result<Vec<u8>> {
52    let recipe = self.load_recipe(recipe_hash).await?;
53    self.read_stored_for_recipe(&recipe).await
54  }
55
56  pub async fn restore_file(&self, recipe_hash: Hash32, output: &mut dyn Write) -> Result<()> {
57    let recipe = self.load_recipe(recipe_hash).await?;
58    let stored = self.read_stored_for_recipe(&recipe).await?;
59    self.decode_recipe(&recipe, stored, output)
60  }
61
62  pub async fn read_file(&self, recipe_hash: Hash32) -> Result<Vec<u8>> {
63    let recipe = self.load_recipe(recipe_hash).await?;
64    let stored = self.read_stored_for_recipe(&recipe).await?;
65    let mut output = FallibleVecWriter::default();
66    let result = self.decode_recipe(&recipe, stored, &mut output);
67    if let Some(requested) = output.allocation_failure {
68      return Err(Error::FileReadAllocationFailed {
69        buffer: "original file",
70        requested,
71      });
72    }
73    result?;
74    Ok(output.bytes)
75  }
76
77  async fn load_recipe(&self, recipe_hash: Hash32) -> Result<RecipeData> {
78    let object = self.source.read_object(&recipe_hash).await?.ok_or(Error::ObjectNotFound)?;
79    if object.kind != ObjectKind::Recipe {
80      return Err(Error::ObjectKindMismatch {
81        hash: recipe_hash,
82        expected: ObjectKind::Recipe,
83        actual: object.kind,
84      });
85    }
86    let (original_file_size, chunk_count) =
87      preflight_recipe_limits(&object.bytes).map_err(|error| Error::InvalidRecipe(error.to_string()))?;
88    self.check_limit("chunk count", self.limits.max_chunk_count, u64::from(chunk_count))?;
89    self.check_limit("original file bytes", self.limits.max_original_file_bytes, original_file_size)?;
90    let recipe = parse_recipe(&object.bytes).map_err(|error| Error::InvalidRecipe(error.to_string()))?;
91    self.check_limit(
92      "stored stream bytes",
93      self.limits.max_stored_stream_bytes,
94      recipe.stored_stream_size,
95    )?;
96    Ok(recipe)
97  }
98
99  async fn read_stored_for_recipe(&self, recipe: &RecipeData) -> Result<Vec<u8>> {
100    let mut stored = Vec::new();
101    for (hash, expected_len) in &recipe.chunks {
102      let object = self.source.read_object(hash).await?.ok_or(Error::ObjectNotFound)?;
103      self.validate_chunk(&object.hash, object.kind, object.bytes.len() as u64, u64::from(*expected_len))?;
104      let requested = stored
105        .len()
106        .checked_add(object.bytes.len())
107        .and_then(|length| u64::try_from(length).ok())
108        .ok_or(Error::FileReadAllocationFailed {
109          buffer: "stored stream",
110          requested: recipe.stored_stream_size,
111        })?;
112      stored
113        .try_reserve(object.bytes.len())
114        .map_err(|_| Error::FileReadAllocationFailed {
115          buffer: "stored stream",
116          requested,
117        })?;
118      stored.extend_from_slice(&object.bytes);
119    }
120    Ok(stored)
121  }
122
123  fn decode_recipe(&self, recipe: &RecipeData, stored: Vec<u8>, output: &mut dyn Write) -> Result<()> {
124    let decoder = self.decoders.get(recipe.transform_id, recipe.transform_version)?;
125    let mut input = BufReader::new(Cursor::new(stored));
126    let mut verified = VerifyingWriter::new(output, recipe.original_file_size);
127    let result = decoder.decode(&mut input, &mut verified);
128    if verified.exceeded {
129      return self.original_mismatch(recipe, &verified);
130    }
131    result?;
132    let actual_hash = verified.hash();
133    if verified.written != recipe.original_file_size || actual_hash != recipe.original_file_hash {
134      return Err(Error::OriginalFileMismatch {
135        expected_size: recipe.original_file_size,
136        actual_size: verified.written,
137        expected_hash: recipe.original_file_hash,
138        actual_hash,
139      });
140    }
141    Ok(())
142  }
143
144  fn original_mismatch(&self, recipe: &RecipeData, writer: &VerifyingWriter<'_>) -> Result<()> {
145    Err(Error::OriginalFileMismatch {
146      expected_size: recipe.original_file_size,
147      actual_size: writer.written,
148      expected_hash: recipe.original_file_hash,
149      actual_hash: writer.hash(),
150    })
151  }
152
153  fn validate_chunk(&self, hash: &Hash32, kind: ObjectKind, actual_len: u64, expected_len: u64) -> Result<()> {
154    if kind != ObjectKind::Chunk {
155      return Err(Error::ObjectKindMismatch {
156        hash: *hash,
157        expected: ObjectKind::Chunk,
158        actual: kind,
159      });
160    }
161    if actual_len != expected_len {
162      return Err(Error::RecipeChunkLengthMismatch {
163        hash: *hash,
164        expected: expected_len,
165        actual: actual_len,
166      });
167    }
168    Ok(())
169  }
170
171  fn check_limit(&self, limit: &'static str, maximum: u64, actual: u64) -> Result<()> {
172    if actual > maximum {
173      return Err(self.limit_error(limit, maximum, actual));
174    }
175    Ok(())
176  }
177
178  fn limit_error(&self, limit: &'static str, maximum: u64, actual: u64) -> Error {
179    Error::FileReadLimitExceeded { limit, maximum, actual }
180  }
181}
182
183#[derive(Default)]
184struct FallibleVecWriter {
185  bytes: Vec<u8>,
186  allocation_failure: Option<u64>,
187}
188
189impl Write for FallibleVecWriter {
190  fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
191    let requested = self
192      .bytes
193      .len()
194      .checked_add(bytes.len())
195      .and_then(|length| u64::try_from(length).ok())
196      .ok_or_else(|| io::Error::new(io::ErrorKind::OutOfMemory, "original file length exceeds address space"))?;
197    if self.bytes.try_reserve(bytes.len()).is_err() {
198      self.allocation_failure = Some(requested);
199      return Err(io::Error::new(
200        io::ErrorKind::OutOfMemory,
201        "unable to allocate original file buffer",
202      ));
203    }
204    self.bytes.extend_from_slice(bytes);
205    Ok(bytes.len())
206  }
207
208  fn flush(&mut self) -> io::Result<()> {
209    Ok(())
210  }
211}
212
213struct VerifyingWriter<'a> {
214  output: &'a mut dyn Write,
215  hasher: Sha3_256,
216  maximum: u64,
217  written: u64,
218  exceeded: bool,
219}
220
221impl<'a> VerifyingWriter<'a> {
222  fn new(output: &'a mut dyn Write, maximum: u64) -> Self {
223    Self {
224      output,
225      hasher: Sha3_256::new(),
226      maximum,
227      written: 0,
228      exceeded: false,
229    }
230  }
231
232  fn hash(&self) -> Hash32 {
233    Hash32::new(self.hasher.clone().finalize().into())
234  }
235}
236
237impl Write for VerifyingWriter<'_> {
238  fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
239    let next = self.written.checked_add(bytes.len() as u64).ok_or_else(|| {
240      self.exceeded = true;
241      io::Error::new(io::ErrorKind::FileTooLarge, "decoded file length overflow")
242    })?;
243    if next > self.maximum {
244      self.exceeded = true;
245      return Err(io::Error::new(io::ErrorKind::FileTooLarge, "decoded file exceeds declared size"));
246    }
247    let written = self.output.write(bytes)?;
248    self.hasher.update(&bytes[..written]);
249    self.written += written as u64;
250    Ok(written)
251  }
252
253  fn flush(&mut self) -> io::Result<()> {
254    self.output.flush()
255  }
256}
257
258#[cfg(test)]
259mod tests {
260  use std::{
261    collections::BTreeMap,
262    sync::{
263      Arc,
264      atomic::{AtomicUsize, Ordering},
265    },
266  };
267
268  use super::*;
269  use crate::{TRANSFORM_ID_NONE, TRANSFORM_VERSION_NONE, VerifiedObject, build_recipe};
270
271  #[derive(Default)]
272  struct MemorySource(BTreeMap<Hash32, VerifiedObject>);
273
274  impl ObjectSource for MemorySource {
275    async fn read_object(&self, hash: &Hash32) -> Result<Option<VerifiedObject>> {
276      Ok(self.0.get(hash).cloned())
277    }
278  }
279
280  fn source_for(
281    chunks: &[&[u8]],
282    order: &[usize],
283    original_size: u64,
284    original_hash: Hash32,
285    transform: (u16, u16),
286  ) -> (MemorySource, Hash32) {
287    let mut source = MemorySource::default();
288    let hashes = chunks
289      .iter()
290      .map(|bytes| {
291        let hash = Hash32::sha3_256(bytes);
292        source.0.insert(
293          hash,
294          VerifiedObject {
295            hash,
296            kind: ObjectKind::Chunk,
297            bytes: bytes.to_vec(),
298          },
299        );
300        hash
301      })
302      .collect::<Vec<_>>();
303    let recipe_chunks = order
304      .iter()
305      .map(|index| (hashes[*index], chunks[*index].len() as u32))
306      .collect::<Vec<_>>();
307    let recipe = build_recipe(original_size, &recipe_chunks, original_hash, transform.0, transform.1);
308    let recipe_hash = Hash32::sha3_256(&recipe);
309    source.0.insert(
310      recipe_hash,
311      VerifiedObject {
312        hash: recipe_hash,
313        kind: ObjectKind::Recipe,
314        bytes: recipe,
315      },
316    );
317    (source, recipe_hash)
318  }
319
320  #[tokio::test]
321  async fn none_reader_handles_empty_single_multi_and_duplicate_chunks() {
322    for (chunks, order) in [
323      (vec![], vec![]),
324      (vec![b"alpha".as_slice()], vec![0]),
325      (vec![b"alpha".as_slice(), b"beta".as_slice()], vec![0, 1]),
326      (vec![b"alpha".as_slice(), b"beta".as_slice()], vec![0, 1, 0]),
327    ] {
328      let expected = order.iter().flat_map(|index| chunks[*index]).copied().collect::<Vec<_>>();
329      let (source, recipe_hash) = source_for(
330        &chunks,
331        &order,
332        expected.len() as u64,
333        Hash32::sha3_256(&expected),
334        (TRANSFORM_ID_NONE, TRANSFORM_VERSION_NONE),
335      );
336      let registry = TransformDecoderRegistry::default();
337      let reader = FileReader::new(&source, &registry, FileReadLimits::default());
338      assert_eq!(reader.read_stored_stream(recipe_hash).await.unwrap(), expected);
339      assert_eq!(reader.read_file(recipe_hash).await.unwrap(), expected);
340      let mut restored = Vec::new();
341      reader.restore_file(recipe_hash, &mut restored).await.unwrap();
342      assert_eq!(restored, expected);
343      assert_eq!(reader.resolve_closure(recipe_hash).await.unwrap().len(), chunks.len() + 1);
344    }
345  }
346
347  #[tokio::test]
348  async fn errors_are_classified_by_file_reader_boundary() {
349    #[derive(Clone, Copy)]
350    enum ExpectedError {
351      NotFound,
352      KindMismatch,
353      InvalidRecipe,
354      ChunkLengthMismatch,
355    }
356
357    let expected = b"alpha";
358    type ErrorCase = (&'static str, fn(&mut MemorySource, Hash32, Hash32), ExpectedError);
359    let cases: [ErrorCase; 6] = [
360      (
361        "missing recipe",
362        |source, recipe_hash, _| {
363          source.0.remove(&recipe_hash);
364        },
365        ExpectedError::NotFound,
366      ),
367      (
368        "missing chunk",
369        |source, _, chunk_hash| {
370          source.0.remove(&chunk_hash);
371        },
372        ExpectedError::NotFound,
373      ),
374      (
375        "recipe kind is chunk",
376        |source, recipe_hash, _| {
377          source.0.get_mut(&recipe_hash).unwrap().kind = ObjectKind::Chunk;
378        },
379        ExpectedError::KindMismatch,
380      ),
381      (
382        "recipe bytes are invalid",
383        |source, recipe_hash, _| {
384          source.0.get_mut(&recipe_hash).unwrap().bytes = vec![0xff];
385        },
386        ExpectedError::InvalidRecipe,
387      ),
388      (
389        "chunk kind is recipe",
390        |source, _, chunk_hash| {
391          source.0.get_mut(&chunk_hash).unwrap().kind = ObjectKind::Recipe;
392        },
393        ExpectedError::KindMismatch,
394      ),
395      (
396        "chunk length mismatch",
397        |source, _, chunk_hash| {
398          source.0.get_mut(&chunk_hash).unwrap().bytes.push(0);
399        },
400        ExpectedError::ChunkLengthMismatch,
401      ),
402    ];
403    let registry = TransformDecoderRegistry::default();
404    for (label, mutate, expected_error) in cases {
405      let (mut source, recipe_hash) = source_for(
406        &[expected.as_slice()],
407        &[0],
408        expected.len() as u64,
409        Hash32::sha3_256(expected),
410        (TRANSFORM_ID_NONE, TRANSFORM_VERSION_NONE),
411      );
412      let chunk_hash = parse_recipe(&source.0[&recipe_hash].bytes).unwrap().chunks[0].0;
413      mutate(&mut source, recipe_hash, chunk_hash);
414      let reader = FileReader::new(&source, &registry, FileReadLimits::default());
415      for result in [
416        reader.read_file(recipe_hash).await.map(|_| ()),
417        reader.resolve_closure(recipe_hash).await.map(|_| ()),
418      ] {
419        let error = result.unwrap_err();
420        let matched = matches!(
421          (&error, expected_error),
422          (Error::ObjectNotFound, ExpectedError::NotFound)
423            | (Error::ObjectKindMismatch { .. }, ExpectedError::KindMismatch)
424            | (Error::InvalidRecipe(_), ExpectedError::InvalidRecipe)
425            | (Error::RecipeChunkLengthMismatch { .. }, ExpectedError::ChunkLengthMismatch)
426        );
427        assert!(matched, "{label} returned unexpected error: {error}");
428      }
429    }
430  }
431
432  #[tokio::test]
433  async fn limits_and_transform_or_original_mismatch_are_structured() {
434    struct CountingDecoder(Arc<AtomicUsize>);
435    impl crate::TransformDecoder for CountingDecoder {
436      fn id(&self) -> u16 {
437        88
438      }
439      fn version(&self) -> u16 {
440        1
441      }
442      fn decode(&self, input: &mut dyn io::BufRead, output: &mut dyn Write) -> Result<()> {
443        self.0.fetch_add(1, Ordering::Relaxed);
444        io::copy(input, output)?;
445        Ok(())
446      }
447    }
448
449    let bytes = b"alpha";
450    let original_hash = Hash32::sha3_256(bytes);
451    let calls = Arc::new(AtomicUsize::new(0));
452    let registry = TransformDecoderRegistry::new([Arc::new(CountingDecoder(calls.clone())) as Arc<dyn crate::TransformDecoder>]).unwrap();
453    let exact_limits = FileReadLimits {
454      max_chunk_count: 1,
455      max_stored_stream_bytes: bytes.len() as u64,
456      max_original_file_bytes: bytes.len() as u64,
457    };
458    let (source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64, original_hash, (88, 1));
459    let reader = FileReader::new(&source, &registry, exact_limits);
460    assert_eq!(reader.read_file(recipe_hash).await.unwrap(), bytes);
461    assert_eq!(calls.load(Ordering::Relaxed), 1);
462
463    let declared_size = 8 * 1024 * 1024 * 1024;
464    let (source, recipe_hash) = source_for(&[], &[], declared_size, Hash32::sha3_256([]), (0, 0));
465    let reader = FileReader::new(&source, &registry, FileReadLimits::default());
466    assert!(matches!(
467      reader.read_file(recipe_hash).await,
468      Err(Error::FileReadLimitExceeded {
469        limit: "original file bytes",
470        actual,
471        ..
472      }) if actual == declared_size
473    ));
474    let reader = FileReader::new(
475      &source,
476      &registry,
477      FileReadLimits {
478        max_original_file_bytes: declared_size,
479        ..FileReadLimits::default()
480      },
481    );
482    assert!(matches!(
483      reader.read_file(recipe_hash).await,
484      Err(Error::OriginalFileMismatch {
485        expected_size,
486        actual_size: 0,
487        ..
488      }) if expected_size == declared_size
489    ));
490
491    let cases = [
492      (
493        FileReadLimits {
494          max_chunk_count: 0,
495          ..FileReadLimits::default()
496        },
497        "chunk count",
498      ),
499      (
500        FileReadLimits {
501          max_stored_stream_bytes: 4,
502          ..FileReadLimits::default()
503        },
504        "stored stream bytes",
505      ),
506      (
507        FileReadLimits {
508          max_original_file_bytes: 4,
509          ..FileReadLimits::default()
510        },
511        "original file bytes",
512      ),
513    ];
514    for (limits, expected_limit) in cases {
515      let before = calls.load(Ordering::Relaxed);
516      let (source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64, original_hash, (88, 1));
517      let reader = FileReader::new(&source, &registry, limits);
518      assert!(matches!(
519        reader.read_file(recipe_hash).await,
520        Err(Error::FileReadLimitExceeded { limit, .. }) if limit == expected_limit
521      ));
522      assert_eq!(calls.load(Ordering::Relaxed), before);
523    }
524
525    let (mut source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64, original_hash, (0, 0));
526    source.0.get_mut(&recipe_hash).unwrap().bytes[45..49].copy_from_slice(&u32::MAX.to_le_bytes());
527    let reader = FileReader::new(
528      &source,
529      &registry,
530      FileReadLimits {
531        max_chunk_count: 0,
532        ..FileReadLimits::default()
533      },
534    );
535    assert!(matches!(
536      reader.read_file(recipe_hash).await,
537      Err(Error::FileReadLimitExceeded { limit: "chunk count", actual, .. }) if actual == u64::from(u32::MAX)
538    ));
539    assert_eq!(calls.load(Ordering::Relaxed), 1);
540
541    let (source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64, original_hash, (TRANSFORM_ID_NONE, 4));
542    let reader = FileReader::new(&source, &registry, FileReadLimits::default());
543    assert!(matches!(
544      reader.read_file(recipe_hash).await,
545      Err(Error::UnsupportedTransform {
546        id: TRANSFORM_ID_NONE,
547        version: 4
548      })
549    ));
550
551    let (source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64, Hash32::sha3_256(b"other"), (0, 0));
552    let none_registry = TransformDecoderRegistry::default();
553    let reader = FileReader::new(&source, &none_registry, FileReadLimits::default());
554    assert_eq!(reader.read_stored_stream(recipe_hash).await.unwrap(), bytes);
555    assert!(matches!(
556      reader.read_file(recipe_hash).await,
557      Err(Error::OriginalFileMismatch { .. })
558    ));
559    let mut restored = Vec::new();
560    assert!(matches!(
561      reader.restore_file(recipe_hash, &mut restored).await,
562      Err(Error::OriginalFileMismatch { .. })
563    ));
564
565    let (source, recipe_hash) = source_for(&[bytes], &[0], bytes.len() as u64 + 1, original_hash, (0, 0));
566    let reader = FileReader::new(&source, &none_registry, FileReadLimits::default());
567    assert!(matches!(
568      reader.read_file(recipe_hash).await,
569      Err(Error::OriginalFileMismatch { expected_size, actual_size, .. }) if expected_size == bytes.len() as u64 + 1 && actual_size == bytes.len() as u64
570    ));
571  }
572
573  #[test]
574  fn registry_accepts_decoder_trait_objects_without_writer_features() {
575    struct Decoder;
576    impl crate::TransformDecoder for Decoder {
577      fn id(&self) -> u16 {
578        91
579      }
580      fn version(&self) -> u16 {
581        2
582      }
583      fn decode(&self, input: &mut dyn io::BufRead, output: &mut dyn Write) -> Result<()> {
584        io::copy(input, output)?;
585        Ok(())
586      }
587    }
588    let registry = TransformDecoderRegistry::new([Arc::new(Decoder) as Arc<dyn crate::TransformDecoder>]).unwrap();
589    assert_eq!(registry.get(91, 2).unwrap().version(), 2);
590  }
591}
assetpack_core/file_reader.rs

assetpack_core/
file_reader.rs
