mod wire;
use super::{
artifacts::ArtifactChecksumRecord,
ic_request::{IcRequestError, MAX_IC_ARGUMENT_BYTES, management_request_digest},
ic_snapshot_metadata::IcSnapshotMetadataReply,
};
use candid::Principal;
use ic_management_canister_types::{ReadCanisterSnapshotDataArgs, SnapshotDataKind};
use std::fmt;
use thiserror::Error;
pub const MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES: usize = 1024 * 1024;
pub const MAX_IC_SNAPSHOT_DATA_REPLY_BYTES: usize = 2 * 1024 * 1024;
pub struct IcSnapshotDataRequest<'metadata> {
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
kind: SnapshotDataKind,
target_bytes: Vec<u8>,
arguments: Vec<u8>,
}
impl<'metadata> IcSnapshotDataRequest<'metadata> {
pub fn new(
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
kind: SnapshotDataKind,
) -> Result<Self, IcSnapshotDataError> {
validate_kind(metadata, &kind)?;
let target = Principal::from_text(metadata.request().target())
.map_err(|_| IcRequestError::InvalidTarget)?;
let arguments = candid::encode_one(ReadCanisterSnapshotDataArgs {
canister_id: target,
snapshot_id: metadata.request().snapshot_id().to_vec(),
kind: kind.clone(),
})
.map_err(|error| IcRequestError::Encoding(error.to_string()))?;
if arguments.len() > MAX_IC_ARGUMENT_BYTES {
return Err(IcRequestError::ArgumentsTooLarge.into());
}
Ok(Self {
metadata,
kind,
target_bytes: target.as_slice().to_vec(),
arguments,
})
}
#[must_use]
pub const fn metadata(&self) -> &'metadata IcSnapshotMetadataReply<'metadata> {
self.metadata
}
#[must_use]
pub const fn kind(&self) -> &SnapshotDataKind {
&self.kind
}
#[must_use]
pub fn target(&self) -> &str {
self.metadata.request().target()
}
#[must_use]
pub fn snapshot_id(&self) -> &[u8] {
self.metadata.request().snapshot_id()
}
#[must_use]
pub fn arguments(&self) -> &[u8] {
&self.arguments
}
#[must_use]
pub const fn receiver(&self) -> &'static str {
"aaaaa-aa"
}
#[must_use]
pub const fn method(&self) -> &'static str {
"read_canister_snapshot_data"
}
#[must_use]
pub fn digest(&self) -> ArtifactChecksumRecord {
management_request_digest(&self.target_bytes, self.method(), &self.arguments)
}
}
pub(crate) fn validate_kind(
metadata: &IcSnapshotMetadataReply<'_>,
kind: &SnapshotDataKind,
) -> Result<(), IcSnapshotDataError> {
let values = metadata.metadata();
match kind {
SnapshotDataKind::WasmModule { offset, size } => {
range(*offset, *size, values.wasm_module_size)?;
}
SnapshotDataKind::WasmMemory { offset, size } => {
range(*offset, *size, values.wasm_memory_size)?;
}
SnapshotDataKind::StableMemory { offset, size } => {
range(*offset, *size, values.stable_memory_size)?;
}
SnapshotDataKind::WasmChunk { hash } => {
if hash.len() != 32 {
return Err(IcSnapshotDataError::InvalidChunkHash);
}
if !values
.wasm_chunk_store
.iter()
.any(|chunk| chunk.hash == *hash)
{
return Err(IcSnapshotDataError::ChunkNotInMetadata);
}
}
}
Ok(())
}
fn range(offset: u64, size: u64, total: u64) -> Result<(), IcSnapshotDataError> {
let admitted_size = usize::try_from(size)
.is_ok_and(|size| (1..=MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES).contains(&size));
if !admitted_size {
return Err(IcSnapshotDataError::InvalidRangeSize);
}
if offset.checked_add(size).is_none_or(|end| end > total) {
return Err(IcSnapshotDataError::RangeOutsideMetadata);
}
Ok(())
}
impl fmt::Debug for IcSnapshotDataRequest<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IcSnapshotDataRequest")
.field("target", &self.target())
.field("argument_bytes", &self.arguments.len())
.finish_non_exhaustive()
}
}
pub struct IcSnapshotDataReply<'request, 'metadata> {
request: &'request IcSnapshotDataRequest<'metadata>,
chunk: Vec<u8>,
chunk_checksum: ArtifactChecksumRecord,
payload_checksum: ArtifactChecksumRecord,
}
impl<'request, 'metadata> IcSnapshotDataReply<'request, 'metadata> {
pub fn decode(
request: &'request IcSnapshotDataRequest<'metadata>,
bytes: &[u8],
) -> Result<Self, IcSnapshotDataError> {
if bytes.len() > MAX_IC_SNAPSHOT_DATA_REPLY_BYTES {
return Err(IcSnapshotDataError::ReplyTooLarge);
}
let chunk = wire::decode(bytes)?;
let chunk_checksum = ArtifactChecksumRecord::from_bytes(&chunk);
match request.kind() {
SnapshotDataKind::WasmModule { size, .. }
| SnapshotDataKind::WasmMemory { size, .. }
| SnapshotDataKind::StableMemory { size, .. } => {
if u64::try_from(chunk.len()).ok() != Some(*size) {
return Err(IcSnapshotDataError::LengthMismatch);
}
}
SnapshotDataKind::WasmChunk { hash } => {
if chunk_checksum.hash() != crate::hash::hex_bytes(hash) {
return Err(IcSnapshotDataError::ChunkHashMismatch);
}
}
}
Ok(Self {
request,
chunk,
chunk_checksum,
payload_checksum: ArtifactChecksumRecord::from_bytes(bytes),
})
}
#[must_use]
pub const fn request(&self) -> &'request IcSnapshotDataRequest<'metadata> {
self.request
}
#[must_use]
pub fn chunk(&self) -> &[u8] {
&self.chunk
}
#[must_use]
pub const fn chunk_checksum(&self) -> &ArtifactChecksumRecord {
&self.chunk_checksum
}
#[must_use]
pub const fn payload_checksum(&self) -> &ArtifactChecksumRecord {
&self.payload_checksum
}
#[must_use]
pub fn digest(&self) -> ArtifactChecksumRecord {
let mut bytes = b"ic-backup/ic-snapshot-data-reply/v1\0".to_vec();
bytes.extend_from_slice(self.request.metadata().digest().hash().as_bytes());
bytes.extend_from_slice(self.request.digest().hash().as_bytes());
bytes.extend_from_slice(self.payload_checksum.hash().as_bytes());
ArtifactChecksumRecord::from_bytes(&bytes)
}
}
impl fmt::Debug for IcSnapshotDataReply<'_, '_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IcSnapshotDataReply")
.field("request", &self.request)
.field("chunk_bytes", &self.chunk.len())
.field("payload_checksum", &self.payload_checksum)
.finish_non_exhaustive()
}
}
#[derive(Debug, Error, Eq, PartialEq)]
pub enum IcSnapshotDataError {
#[error(transparent)]
Request(#[from] IcRequestError),
#[error("snapshot range size must be 1..={MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES}")]
InvalidRangeSize,
#[error("snapshot range is outside retained metadata")]
RangeOutsideMetadata,
#[error("snapshot chunk identity must be 32 bytes")]
InvalidChunkHash,
#[error("snapshot chunk identity is absent from retained metadata")]
ChunkNotInMetadata,
#[error("snapshot data reply exceeds {MAX_IC_SNAPSHOT_DATA_REPLY_BYTES} bytes")]
ReplyTooLarge,
#[error("invalid or unsupported snapshot data reply")]
InvalidReply,
#[error("snapshot data length differs from the requested range")]
LengthMismatch,
#[error("snapshot data differs from the requested chunk hash")]
ChunkHashMismatch,
}
#[cfg(test)]
mod tests;