magi-code 0.96.1

use super::*;

pub(super) const COMMITTED_UNDURABLE_WARNING: &str =
    "mutation committed but durability is uncertain";
pub(super) const COMMITTED_CLEANUP_WARNING: &str =
    "mutation committed but recovery cleanup is incomplete";
pub(super) const TARGET_CHANGED_DURING_REWIND: &str = "target changed during rewind";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum MutationOutcome {
    Committed { warning: Option<&'static str> },
    Conflict(&'static str),
    Failed(&'static str),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ValidatedRemoval {
    Removed,
    Missing,
    Preserved,
}

pub(super) static REWIND_TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);

fn move_entry_to_quarantine(parent: &CapDir, source: &OsStr, kind: &str) -> io::Result<OsString> {
    move_entry_to_quarantine_with_attempts(parent, source, kind, 128, |_| {})
}

pub(super) fn move_entry_to_quarantine_with_attempts(
    parent: &CapDir,
    source: &OsStr,
    kind: &str,
    attempt_limit: u64,
    mut before_move: impl FnMut(&OsStr),
) -> io::Result<OsString> {
    let stamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let pid = std::process::id();
    for attempt in 0..attempt_limit {
        let sequence = REWIND_TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
        let name = OsString::from(format!(
            ".magi-rewind-{kind}-{pid}-{stamp}-{sequence}-{attempt}"
        ));
        before_move(name.as_os_str());
        // The native operation checks the destination and moves the source
        // atomically. Never replace a raced destination or use check-then-rename.
        match rename_noreplace(parent, source, name.as_os_str()) {
            Ok(()) => return Ok(name),
            // A generated name can still be occupied. Try another name, but
            // leave the occupied entry untouched.
            Err(error) if error.kind() == ErrorKind::AlreadyExists => {}
            Err(error) => return Err(error),
        }
    }
    Err(io::Error::new(
        ErrorKind::AlreadyExists,
        "could not allocate a rewind quarantine name",
    ))
}

fn move_target_to_quarantine(target: &RewindTarget) -> io::Result<OsString> {
    reject_final_symlink(&target.parent, &target.name)?;
    move_entry_to_quarantine(&target.parent, &target.name, "quarantine")
}

fn remove_validated_entry(
    parent: &CapDir,
    source: &OsStr,
    expected_hash: &str,
) -> ValidatedRemoval {
    let recovery_name = match move_entry_to_quarantine(parent, source, "cleanup") {
        Ok(name) => name,
        Err(error) if error.kind() == ErrorKind::NotFound => return ValidatedRemoval::Missing,
        Err(_) => return ValidatedRemoval::Preserved,
    };
    let snapshot = match snapshot_named(parent, &recovery_name) {
        Ok(Some(snapshot)) => snapshot,
        Ok(None) => return ValidatedRemoval::Missing,
        Err(_) => return ValidatedRemoval::Preserved,
    };
    if snapshot.hash != expected_hash {
        return ValidatedRemoval::Preserved;
    }
    match parent.remove_file_or_symlink(Path::new(recovery_name.as_os_str())) {
        Ok(()) => ValidatedRemoval::Removed,
        Err(_) => ValidatedRemoval::Preserved,
    }
}

fn recover_quarantine(
    parent: &CapDir,
    quarantine_name: &OsStr,
    target_name: &OsStr,
    quarantine_hash: &str,
) {
    let Ok(Some(snapshot)) = snapshot_named(parent, quarantine_name) else {
        return;
    };
    if snapshot.hash != quarantine_hash {
        return;
    }
    // The native operation checks the destination and rolls the quarantine
    // entry back atomically. On conflict or unsupported platforms, retain both
    // entries rather than replacing or cleaning either one.
    let _ = rename_noreplace(parent, quarantine_name, target_name);
}

fn finish_committed(
    parent: &CapDir,
    inject_final_sync_failure: bool,
    cleanup_incomplete: bool,
) -> MutationOutcome {
    let final_sync_failed = inject_final_sync_failure || sync_capability_directory(parent).is_err();
    let warning = if final_sync_failed {
        Some(COMMITTED_UNDURABLE_WARNING)
    } else if cleanup_incomplete {
        Some(COMMITTED_CLEANUP_WARNING)
    } else {
        None
    };
    MutationOutcome::Committed { warning }
}

fn create_rewind_temp_file(
    parent: &CapDir,
    target_mode: Option<u32>,
) -> anyhow::Result<(OsString, CapFile)> {
    #[cfg(not(unix))]
    let _ = target_mode;
    let stamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let pid = std::process::id();
    for attempt in 0..128_u64 {
        let sequence = REWIND_TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
        let name = OsString::from(format!(
            ".magi-rewind-{pid}-{stamp}-{sequence}-{attempt}.tmp"
        ));
        let mut options = CapOpenOptions::new();
        options
            .write(true)
            .create_new(true)
            .truncate(true)
            .follow(FollowSymlinks::No);
        #[cfg(unix)]
        options.mode(target_mode.unwrap_or(0o666));
        match parent.open_with(Path::new(name.as_os_str()), &options) {
            Ok(file) => return Ok((name, file)),
            Err(error) if error.kind() == ErrorKind::AlreadyExists => {}
            Err(error) => return Err(error.into()),
        }
    }
    anyhow::bail!("could not allocate a unique rewind temporary file")
}

pub(super) fn prepare_rewind_temp_file(
    parent: &CapDir,
    bytes: &[u8],
    target_mode: Option<u32>,
    inject_write_failure: bool,
) -> Result<(OsString, String), &'static str> {
    let (name, mut file) =
        create_rewind_temp_file(parent, target_mode).map_err(|_| "write failed")?;
    let expected_hash = sha256_hex(bytes);
    let prepare_result = (|| -> anyhow::Result<()> {
        file.write_all(bytes)?;
        if inject_write_failure {
            anyhow::bail!("injected write failure");
        }
        file.flush()?;
        #[cfg(unix)]
        if let Some(mode) = target_mode {
            file.set_permissions(cap_std::fs::Permissions::from_mode(mode))?;
        }
        file.sync_all()?;
        Ok(())
    })();
    drop(file);
    if prepare_result.is_err() {
        // The entry is removed only after it has been moved and verified. A
        // replacement of the temporary name is retained as recovery evidence.
        let _ = remove_validated_entry(parent, &name, &expected_hash);
        return Err("write failed");
    }
    Ok((name, expected_hash))
}

pub(super) struct TempFileCleanupGuard<'a> {
    pub(super) parent: &'a CapDir,
    pub(super) name: Option<OsString>,
    expected_hash: String,
}

impl<'a> TempFileCleanupGuard<'a> {
    pub(super) fn new(parent: &'a CapDir, name: OsString, expected_hash: String) -> Self {
        Self {
            parent,
            name: Some(name),
            expected_hash,
        }
    }

    fn disarm(&mut self) {
        self.name = None;
    }
}

impl Drop for TempFileCleanupGuard<'_> {
    fn drop(&mut self) {
        if let Some(name) = self.name.take() {
            let _ = remove_validated_entry(self.parent, &name, &self.expected_hash);
        }
    }
}

fn atomic_replace_target(
    target: &RewindTarget,
    bytes: &[u8],
    expected_current_hash: &str,
    target_mode: Option<u32>,
    inject_write_failure: bool,
    inject_final_sync_failure: bool,
    after_quarantine_hashed: &mut dyn FnMut(),
) -> MutationOutcome {
    let (temp_name, temp_hash) =
        match prepare_rewind_temp_file(&target.parent, bytes, target_mode, inject_write_failure) {
            Ok(prepared) => prepared,
            Err(reason) => return MutationOutcome::Failed(reason),
        };
    let mut temp_cleanup =
        TempFileCleanupGuard::new(&target.parent, temp_name.clone(), temp_hash.clone());

    let quarantine_name = match move_target_to_quarantine(target) {
        Ok(name) => name,
        Err(error) if matches!(error.kind(), ErrorKind::NotFound | ErrorKind::AlreadyExists) => {
            return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
        }
        Err(_) => return MutationOutcome::Failed("write failed"),
    };
    let quarantine_snapshot = match snapshot_named(&target.parent, &quarantine_name) {
        Ok(Some(snapshot)) => snapshot,
        Ok(None) | Err(_) => return MutationOutcome::Failed("hash check failed"),
    };
    if quarantine_snapshot.hash != expected_current_hash {
        recover_quarantine(
            &target.parent,
            &quarantine_name,
            &target.name,
            &quarantine_snapshot.hash,
        );
        return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
    }
    after_quarantine_hashed();

    match target.parent.hard_link(
        Path::new(temp_name.as_os_str()),
        &target.parent,
        Path::new(target.name.as_os_str()),
    ) {
        Ok(()) => {}
        Err(error) if error.kind() == ErrorKind::AlreadyExists => {
            recover_quarantine(
                &target.parent,
                &quarantine_name,
                &target.name,
                &quarantine_snapshot.hash,
            );
            return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
        }
        Err(_) => {
            recover_quarantine(
                &target.parent,
                &quarantine_name,
                &target.name,
                &quarantine_snapshot.hash,
            );
            return MutationOutcome::Failed("write failed");
        }
    }

    let temp_cleanup_incomplete = !matches!(
        remove_validated_entry(&target.parent, &temp_name, &temp_hash),
        ValidatedRemoval::Removed | ValidatedRemoval::Missing
    );
    temp_cleanup.disarm();
    let quarantine_cleanup_incomplete = !matches!(
        remove_validated_entry(&target.parent, &quarantine_name, &quarantine_snapshot.hash,),
        ValidatedRemoval::Removed | ValidatedRemoval::Missing
    );
    finish_committed(
        &target.parent,
        inject_final_sync_failure,
        temp_cleanup_incomplete || quarantine_cleanup_incomplete,
    )
}

pub(super) fn delete_target(
    target: &RewindTarget,
    expected_current_hash: &str,
    inject_delete_failure: bool,
    inject_final_sync_failure: bool,
    after_quarantine_hashed: &mut impl FnMut(),
) -> MutationOutcome {
    let quarantine_name = match move_target_to_quarantine(target) {
        Ok(name) => name,
        Err(error) if matches!(error.kind(), ErrorKind::NotFound | ErrorKind::AlreadyExists) => {
            return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
        }
        Err(_) => return MutationOutcome::Failed("delete failed"),
    };
    let quarantine_snapshot = match snapshot_named(&target.parent, &quarantine_name) {
        Ok(Some(snapshot)) => snapshot,
        Ok(None) | Err(_) => return MutationOutcome::Failed("hash check failed"),
    };
    if quarantine_snapshot.hash != expected_current_hash {
        recover_quarantine(
            &target.parent,
            &quarantine_name,
            &target.name,
            &quarantine_snapshot.hash,
        );
        return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
    }
    after_quarantine_hashed();
    if inject_delete_failure {
        recover_quarantine(
            &target.parent,
            &quarantine_name,
            &target.name,
            &quarantine_snapshot.hash,
        );
        return MutationOutcome::Failed("delete failed");
    }

    match remove_validated_entry(&target.parent, &quarantine_name, &quarantine_snapshot.hash) {
        ValidatedRemoval::Removed => {
            finish_committed(&target.parent, inject_final_sync_failure, false)
        }
        ValidatedRemoval::Missing => MutationOutcome::Failed("delete failed"),
        ValidatedRemoval::Preserved => MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND),
    }
}

pub(super) fn install_new_target(
    target: &RewindTarget,
    bytes: &[u8],
    inject_write_failure: bool,
    inject_final_sync_failure: bool,
) -> MutationOutcome {
    let (temp_name, temp_hash) =
        match prepare_rewind_temp_file(&target.parent, bytes, None, inject_write_failure) {
            Ok(prepared) => prepared,
            Err(reason) => return MutationOutcome::Failed(reason),
        };
    let mut temp_cleanup =
        TempFileCleanupGuard::new(&target.parent, temp_name.clone(), temp_hash.clone());
    match target.parent.hard_link(
        Path::new(temp_name.as_os_str()),
        &target.parent,
        Path::new(target.name.as_os_str()),
    ) {
        Ok(()) => {}
        Err(error) if error.kind() == ErrorKind::AlreadyExists => {
            return MutationOutcome::Conflict(TARGET_CHANGED_DURING_REWIND);
        }
        Err(_) => return MutationOutcome::Failed("write failed"),
    }
    let cleanup_incomplete = !matches!(
        remove_validated_entry(&target.parent, &temp_name, &temp_hash),
        ValidatedRemoval::Removed | ValidatedRemoval::Missing
    );
    temp_cleanup.disarm();
    finish_committed(
        &target.parent,
        inject_final_sync_failure,
        cleanup_incomplete,
    )
}

pub(super) struct RewindBlobOptions<'a> {
    pub(super) expected_current_hash: Option<&'a str>,
    pub(super) target_mode: Option<u32>,
    pub(super) injected_failure: Option<&'a str>,
    pub(super) inject_final_sync_failure: bool,
    pub(super) after_quarantine_hashed: &'a mut dyn FnMut(),
}

pub(super) fn write_rewind_blob(
    store: &CheckpointStore,
    blob: &SnapshotBlobRef,
    target: &RewindTarget,
    options: RewindBlobOptions<'_>,
) -> MutationOutcome {
    let RewindBlobOptions {
        expected_current_hash,
        target_mode,
        injected_failure,
        inject_final_sync_failure,
        after_quarantine_hashed,
    } = options;
    let bytes = match read_rewind_blob(store, blob, injected_failure) {
        Ok(bytes) => bytes,
        Err(reason) => return MutationOutcome::Failed(reason),
    };
    match expected_current_hash {
        Some(expected_current_hash) => atomic_replace_target(
            target,
            &bytes,
            expected_current_hash,
            target_mode,
            injected_failure == Some("write"),
            inject_final_sync_failure,
            after_quarantine_hashed,
        ),
        None => install_new_target(
            target,
            &bytes,
            injected_failure == Some("write"),
            inject_final_sync_failure,
        ),
    }
}