use std::fs;
use std::ops::ControlFlow;
use std::path::{Path, PathBuf};
use super::{
AcquisitionOptions, AcquisitionWriter, Seal, WriteOptions, checkpoint_count, checkpoint_path,
ensure_output_absent, fingerprint, normalize_errors, normalize_first, read_fixed,
require_directory, require_file_type, staged_path, validate_segment_budget,
};
use crate::publication::{OutputLock, acquisition_path};
use crate::{AcquisitionError, EwfError, Image, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum AcquisitionOperationPhase {
ValidatingSegments,
RehashingMedia,
Publishing,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct AcquisitionOperationProgress {
pub phase: AcquisitionOperationPhase,
pub bytes_processed: u64,
pub bytes_total: u64,
pub segments_processed: usize,
pub segments_total: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct AcquisitionCheckpoint {
pub source_size: u64,
pub checkpoint_bytes: u64,
pub chunk_size: u64,
pub sealed_segments: usize,
pub stored_bytes: u64,
pub ready_to_finish: bool,
pub publication_started: bool,
pub segment_hashes_validated: bool,
pub acquisition_errors: Vec<AcquisitionError>,
pub substituted_sectors: u64,
}
pub(super) struct CheckpointData {
pub(super) state: PathBuf,
pub(super) sealed: Vec<Seal>,
pub(super) report: AcquisitionCheckpoint,
pub(super) image: Option<Image>,
}
impl AcquisitionWriter {
pub fn inspect_checkpoint(
first: impl AsRef<Path>,
options: &AcquisitionOptions,
source_identity: [u8; 32],
) -> Result<AcquisitionCheckpoint> {
inspect(first.as_ref(), options, source_identity, false, &mut |_| {
ControlFlow::Continue(())
})
}
pub fn validate_checkpoint(
first: impl AsRef<Path>,
options: &AcquisitionOptions,
source_identity: [u8; 32],
mut on_progress: impl FnMut(AcquisitionOperationProgress) -> ControlFlow<()>,
) -> Result<AcquisitionCheckpoint> {
inspect(
first.as_ref(),
options,
source_identity,
true,
&mut on_progress,
)
}
}
fn inspect(
first: &Path,
options: &AcquisitionOptions,
source_identity: [u8; 32],
validate: bool,
callback: &mut impl FnMut(AcquisitionOperationProgress) -> ControlFlow<()>,
) -> Result<AcquisitionCheckpoint> {
let write_options = options.writer_options()?;
let first = normalize_first(first)?;
validate_segment_budget(&first, options)?;
let _lock = OutputLock::acquire(&first)?;
Ok(load_checkpoint(
&first,
options,
&write_options,
source_identity,
validate,
callback,
)?
.report)
}
pub(super) fn load_checkpoint(
first: &Path,
options: &AcquisitionOptions,
write_options: &WriteOptions,
source_identity: [u8; 32],
validate: bool,
callback: &mut impl FnMut(AcquisitionOperationProgress) -> ControlFlow<()>,
) -> Result<CheckpointData> {
crate::publication::ensure_no_publication(first)?;
let state = acquisition_path(first)?;
require_directory(&state)?;
require_directory(&state.join("scratch"))?;
if read_fixed::<32>(&state.join("config"))?
!= fingerprint(first, options, write_options, source_identity)?
{
return Err(EwfError::Malformed(
"acquisition identity or options changed".into(),
));
}
let count = checkpoint_count(&state)?;
let mut sealed = Vec::with_capacity(count);
let mut paths = Vec::with_capacity(count);
let mut previous = 0;
let mut stored_bytes = 0_u64;
let chunk_size = u64::from(options.bytes_per_sector) * u64::from(options.sectors_per_chunk);
for index in 1..=count {
let seal = Seal::read(&checkpoint_path(&state, index))?;
if seal.end <= previous
|| seal.end > options.source_size
|| (seal.end < options.source_size && !seal.end.is_multiple_of(chunk_size))
|| (seal.end - previous).div_ceil(chunk_size) > u64::from(options.chunks_per_segment)
{
return Err(EwfError::Malformed(
"invalid acquisition checkpoint geometry".into(),
));
}
let path = staged_path(first, &state, index)?;
let metadata = fs::symlink_metadata(&path)?;
require_file_type(&metadata)?;
if metadata.len() != seal.size {
return Err(EwfError::Malformed(
"sealed acquisition segment changed".into(),
));
}
stored_bytes = stored_bytes
.checked_add(seal.size)
.ok_or_else(|| EwfError::Malformed("acquisition stored-size overflow".into()))?;
previous = seal.end;
paths.push(path);
sealed.push(seal);
}
let publishing = state.join("publishing").try_exists()?;
if publishing {
read_fixed::<0>(&state.join("publishing"))?;
if previous != options.source_size {
return Err(EwfError::Malformed(
"incomplete acquisition publication".into(),
));
}
} else {
ensure_output_absent(first)?;
}
if validate {
let mut progress = operation_progress(
AcquisitionOperationPhase::ValidatingSegments,
stored_bytes,
count,
);
notify(progress, callback)?;
for (path, seal) in paths.iter().zip(&sealed) {
seal.validate_with_progress(path, &mut progress, callback)?;
}
}
let (image, errors) = if paths.is_empty() {
(None, Vec::new())
} else {
let image = Image::open_acquisition_prefix(paths, options.source_size, previous)?;
if image.chunk_size() != chunk_size
|| image.info().acquisition_complete != (previous == options.source_size)
{
return Err(EwfError::Malformed(
"acquisition segment geometry changed".into(),
));
}
let errors = normalize_errors(
image.acquisition_errors(),
previous,
options.bytes_per_sector,
)?;
(Some(image), errors)
};
let report = AcquisitionCheckpoint {
source_size: options.source_size,
checkpoint_bytes: previous,
chunk_size,
sealed_segments: count,
stored_bytes,
ready_to_finish: previous == options.source_size,
publication_started: publishing,
segment_hashes_validated: validate,
substituted_sectors: errors.iter().map(|range| range.sector_count).sum(),
acquisition_errors: errors,
};
Ok(CheckpointData {
state,
sealed,
report,
image,
})
}
pub(super) fn operation_progress(
phase: AcquisitionOperationPhase,
total: u64,
segments: usize,
) -> AcquisitionOperationProgress {
AcquisitionOperationProgress {
phase,
bytes_processed: 0,
bytes_total: total,
segments_processed: 0,
segments_total: segments,
}
}
pub(super) fn notify(
progress: AcquisitionOperationProgress,
callback: &mut impl FnMut(AcquisitionOperationProgress) -> ControlFlow<()>,
) -> Result<()> {
if callback(progress).is_break() {
Err(EwfError::Aborted)
} else {
Ok(())
}
}