ic_host_artifacts/artifact/chunks/
mod.rs1#[cfg(test)]
4mod tests;
5
6use super::{ArtifactError, ArtifactIdentity, Sha256Digest, reserve_bounded, visit_reader};
7use sha2::{Digest, Sha256};
8use std::{collections::TryReserveError, fmt, io::Read, num::NonZeroUsize};
9
10#[derive(Debug)]
12pub enum ChunkDigestError {
13 Input(ArtifactError),
15 ChunkLimit {
17 limit: usize,
19 },
20 Allocation(TryReserveError),
22}
23
24impl fmt::Display for ChunkDigestError {
25 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26 match self {
27 Self::Input(source) => write!(f, "chunk input failed: {source}"),
28 Self::ChunkLimit { limit } => write!(f, "input exceeds {limit} chunks"),
29 Self::Allocation(source) => write!(f, "chunk digest allocation failed: {source}"),
30 }
31 }
32}
33
34impl std::error::Error for ChunkDigestError {
35 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
36 match self {
37 Self::Input(source) => Some(source),
38 Self::Allocation(source) => Some(source),
39 Self::ChunkLimit { .. } => None,
40 }
41 }
42}
43
44impl From<ArtifactError> for ChunkDigestError {
45 fn from(source: ArtifactError) -> Self {
46 Self::Input(source)
47 }
48}
49
50pub fn chunk_digests(
72 mut reader: impl Read,
73 chunk_bytes: NonZeroUsize,
74 max_bytes: u64,
75 max_chunks: usize,
76) -> Result<(Vec<Sha256Digest>, ArtifactIdentity), ChunkDigestError> {
77 let mut digests = Vec::new();
78 let mut whole = Sha256::new();
79 let mut chunk = Sha256::new();
80 let mut filled = 0;
81 let bytes = visit_reader::<ChunkDigestError>(&mut reader, max_bytes, |mut buffer| {
82 whole.update(buffer);
83 while !buffer.is_empty() {
84 if filled == 0 {
85 if digests.len() == max_chunks {
86 return Err(ChunkDigestError::ChunkLimit { limit: max_chunks });
87 }
88 reserve_bounded(&mut digests, 1, max_chunks)
89 .map_err(ChunkDigestError::Allocation)?;
90 }
91 let count = buffer.len().min(chunk_bytes.get() - filled);
92 chunk.update(&buffer[..count]);
93 filled += count;
94 buffer = &buffer[count..];
95 if filled == chunk_bytes.get() {
96 digests.push(Sha256Digest::from_bytes(chunk.finalize_reset().into()));
97 filled = 0;
98 }
99 }
100 Ok(())
101 })?;
102 if filled != 0 {
103 digests.push(Sha256Digest::from_bytes(chunk.finalize().into()));
105 }
106 Ok((
107 digests,
108 ArtifactIdentity {
109 bytes,
110 sha256: Sha256Digest::from_bytes(whole.finalize().into()),
111 },
112 ))
113}