mod plan;
mod session;
pub use plan::IcSnapshotDataUploadPreparationError;
use super::{
DownloadJournalGuard, File, IcSnapshotArtifactError, REGIONS, member,
verification::open_regular_child, with_artifact_directory,
};
use crate::model::{
ic_snapshot_data::{MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES, validate_kind},
ic_snapshot_metadata::IcSnapshotMetadataReply,
ic_snapshot_upload::{
IcSnapshotDataUploadPlan, IcSnapshotUploadError, IcSnapshotUploadRequest,
},
operation_plan::OperationPlanRecord,
};
use ic_platform::management_canister_types::SnapshotDataKind;
use std::io::{Read, Seek, SeekFrom};
use thiserror::Error;
impl DownloadJournalGuard<'_> {
pub(crate) fn upload_session<'a, 'workflow, 'request, 'source>(
&'a self,
snapshot: &'a str,
upload: &'a IcSnapshotDataUploadPlan<'workflow, 'request, 'source>,
) -> Result<
session::UploadSession<'a, 'workflow, 'request, 'source>,
IcSnapshotUploadArtifactError,
> {
session::UploadSession::new(self, snapshot, upload)
}
pub fn prepare_ic_snapshot_upload_metadata<'source>(
&self,
plan: &'source OperationPlanRecord,
snapshot: &str,
metadata: &'source IcSnapshotMetadataReply<'source>,
) -> Result<IcSnapshotUploadRequest<'source>, IcSnapshotUploadArtifactError> {
let checksum = self.verify_ic_snapshot_artifact(plan, snapshot, metadata)?;
Ok(IcSnapshotUploadRequest::metadata(
plan, metadata, &checksum,
)?)
}
pub fn prepare_ic_snapshot_upload_data<'source>(
&self,
plan: &OperationPlanRecord,
snapshot: &str,
metadata_upload: &IcSnapshotUploadRequest<'source>,
snapshot_id: &[u8],
source_kind: SnapshotDataKind,
) -> Result<IcSnapshotUploadRequest<'source>, IcSnapshotUploadArtifactError> {
let metadata = metadata_upload.source();
metadata_upload.validate_data_destination(snapshot_id)?;
validate_kind(metadata, &source_kind).map_err(IcSnapshotUploadError::from)?;
self.with_upload_source(plan, snapshot, metadata_upload, |directory| {
let chunk = read_chunk(directory, metadata, &source_kind)?;
Ok(IcSnapshotUploadRequest::data(
metadata_upload,
snapshot_id,
source_kind,
&chunk,
)?)
})
}
fn with_upload_source<T>(
&self,
plan: &OperationPlanRecord,
snapshot: &str,
metadata_upload: &IcSnapshotUploadRequest<'_>,
read: impl FnOnce(&File) -> Result<T, IcSnapshotUploadArtifactError>,
) -> Result<T, IcSnapshotUploadArtifactError> {
let metadata = metadata_upload.source();
self.require_upload_source(plan, snapshot, metadata_upload)?;
let entry = self
.record
.artifact(metadata.request().target(), snapshot)
.map_err(super::DownloadJournalError::from)
.map_err(IcSnapshotArtifactError::from)?;
with_artifact_directory(self.layout.root(), entry.artifact_path(), read, || {
self.require_upload_source(plan, snapshot, metadata_upload)
})
}
fn require_upload_source(
&self,
plan: &OperationPlanRecord,
snapshot: &str,
upload: &IcSnapshotUploadRequest<'_>,
) -> Result<(), IcSnapshotUploadArtifactError> {
if plan.digest() != upload.source_plan().digest() {
return Err(IcSnapshotUploadArtifactError::SourceMismatch);
}
let actual = self.verify_ic_snapshot_artifact(plan, snapshot, upload.source())?;
if actual != *upload.source_checksum() {
return Err(IcSnapshotUploadArtifactError::SourceMismatch);
}
Ok(())
}
}
fn read_chunk(
directory: &File,
metadata: &IcSnapshotMetadataReply<'_>,
kind: &SnapshotDataKind,
) -> Result<Vec<u8>, IcSnapshotArtifactError> {
let (name, length, offset, size) = match member::select(kind)? {
member::Member::Region {
region,
offset,
size,
} => (
REGIONS[region.index()].to_owned(),
Some(metadata.region_sizes()[region.index()]),
offset,
Some(size),
),
member::Member::Chunk { name } => (name, None, 0, None),
};
let maximum = length.unwrap_or(MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES as u64);
let (mut file, original) = open_regular_child(directory, &name, length, maximum)?;
file.seek(SeekFrom::Start(offset))?;
let size = size.unwrap_or(original.len());
read_chunk_bytes(file, size)
}
fn read_chunk_bytes(reader: impl Read, size: u64) -> Result<Vec<u8>, IcSnapshotArtifactError> {
#[cfg(test)]
let reader = super::read_count::Reader(reader);
let limit = usize::try_from(size)
.ok()
.filter(|size| *size <= MAX_IC_SNAPSHOT_DATA_CHUNK_BYTES)
.ok_or(IcSnapshotArtifactError::FileShape)?;
ic_host_artifacts::artifact::read_exact_reader(reader.take(size), limit).map_err(|error| {
use ic_host_artifacts::artifact::ArtifactError;
match error {
ArtifactError::LimitExceeded { .. } | ArtifactError::LengthMismatch { .. } => {
IcSnapshotArtifactError::FileShape
}
error => IcSnapshotArtifactError::Io(error.into()),
}
})
}
#[derive(Debug, Error)]
pub enum IcSnapshotUploadArtifactError {
#[error("snapshot upload source checksum differs")]
SourceMismatch,
#[error(transparent)]
Artifact(#[from] IcSnapshotArtifactError),
#[error(transparent)]
Upload(#[from] IcSnapshotUploadError),
}