ic-testkit 0.76.0

//! Convert a native snapshot to metadata-upload semantics without shared files.

use super::{PocketIc, PocketIcCallFailure, Principal, transport};
use ic_platform::{
    candid::contract::{self, ArgumentDecoder, CandidType, DecoderConfig},
    management_canister_types::{
        ReadCanisterSnapshotDataArgs, ReadCanisterSnapshotDataResult,
        ReadCanisterSnapshotMetadataArgs, ReadCanisterSnapshotMetadataResult, SnapshotDataKind,
        SnapshotDataOffset, UploadCanisterSnapshotDataArgs, UploadCanisterSnapshotMetadataArgs,
        UploadCanisterSnapshotMetadataResult,
    },
    pocket_ic::common::rest::RawEffectivePrincipal,
};

// One MiB leaves room for Candid framing within the management message limit.
// Total work follows the already captured canister state; no whole heap buffer,
// filesystem, application-specific timer reset or additional clock policy.
const CHUNK_BYTES: u64 = 1024 * 1024;

pub(super) enum Failure {
    Call(PocketIcCallFailure),
    Codec(contract::Error),
    Invalid(&'static str),
}

impl From<PocketIcCallFailure> for Failure {
    fn from(error: PocketIcCallFailure) -> Self {
        Self::Call(error)
    }
}

fn call<A: CandidType, R: for<'de> ArgumentDecoder<'de>>(
    pic: &PocketIc,
    canister_id: Principal,
    sender: Option<Principal>,
    method: &str,
    args: A,
) -> Result<R, Failure> {
    let args = contract::encode_args((args,)).map_err(Failure::Codec)?;
    let reply = transport::capture_call(|| {
        pic.update_call_with_effective_principal(
            Principal::management_canister(),
            RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()),
            sender.unwrap_or_else(Principal::anonymous),
            method,
            args,
        )
    })?;
    decode_reply(&reply)
}

fn decode_reply<R: for<'de> ArgumentDecoder<'de>>(reply: &[u8]) -> Result<R, Failure> {
    if reply.len() > 2_000_000 {
        return Err(Failure::Invalid(
            "snapshot reply exceeds the management message limit",
        ));
    }
    let mut config = DecoderConfig::new();
    config.set_full_error_message(false);
    contract::decode_args_with_config(reply, &config).map_err(Failure::Codec)
}

pub(super) fn upload(
    pic: &PocketIc,
    canister_id: Principal,
    sender: Option<Principal>,
    original: &[u8],
    pending: &mut Vec<Vec<u8>>,
) -> Result<Vec<u8>, Failure> {
    let (metadata,): (ReadCanisterSnapshotMetadataResult,) = call(
        pic,
        canister_id,
        sender,
        "read_canister_snapshot_metadata",
        ReadCanisterSnapshotMetadataArgs {
            canister_id,
            snapshot_id: original.to_vec(),
        },
    )?;
    require_complete_system_state(&metadata)?;
    let (uploaded,): (UploadCanisterSnapshotMetadataResult,) = call(
        pic,
        canister_id,
        sender,
        "upload_canister_snapshot_metadata",
        UploadCanisterSnapshotMetadataArgs {
            canister_id,
            replace_snapshot: None,
            wasm_module_size: metadata.wasm_module_size,
            // None denotes globals outside the exported/mutable snapshot set.
            globals: metadata.globals.into_iter().flatten().collect(),
            wasm_memory_size: metadata.wasm_memory_size,
            stable_memory_size: metadata.stable_memory_size,
            certified_data: metadata.certified_data,
            global_timer: metadata.global_timer,
            on_low_wasm_memory_hook_status: metadata.on_low_wasm_memory_hook_status,
        },
    )?;
    let uploaded = uploaded.snapshot_id;
    pending.push(uploaded.clone());
    let transfer = SnapshotTransfer {
        pic,
        canister_id,
        sender,
        original,
        uploaded: &uploaded,
    };
    transfer.copy_region(metadata.wasm_module_size, Region::WasmModule)?;
    transfer.copy_region(metadata.wasm_memory_size, Region::WasmMemory)?;
    transfer.copy_region(metadata.stable_memory_size, Region::StableMemory)?;
    for chunk in metadata.wasm_chunk_store {
        let data = transfer.read(SnapshotDataKind::WasmChunk {
            hash: chunk.hash.clone(),
        })?;
        if data.len() as u64 > CHUNK_BYTES
            || ic_host_artifacts::artifact::Sha256Digest::compute(&data)
                .as_bytes()
                .as_slice()
                != chunk.hash.as_slice()
        {
            return Err(Failure::Invalid(
                "snapshot chunk store entry exceeds the protocol limit or has the wrong hash",
            ));
        }
        transfer.write(SnapshotDataOffset::WasmChunk, data)?;
    }
    Ok(uploaded)
}

const fn require_complete_system_state(
    metadata: &ReadCanisterSnapshotMetadataResult,
) -> Result<(), Failure> {
    if metadata.global_timer.is_none() || metadata.on_low_wasm_memory_hook_status.is_none() {
        return Err(Failure::Invalid(
            "snapshot omitted timer or low-memory hook state",
        ));
    }
    Ok(())
}

#[derive(Clone, Copy)]
enum Region {
    WasmModule,
    WasmMemory,
    StableMemory,
}

impl Region {
    const fn at(self, offset: u64, size: u64) -> (SnapshotDataKind, SnapshotDataOffset) {
        match self {
            Self::WasmModule => (
                SnapshotDataKind::WasmModule { offset, size },
                SnapshotDataOffset::WasmModule { offset },
            ),
            Self::WasmMemory => (
                SnapshotDataKind::WasmMemory { offset, size },
                SnapshotDataOffset::WasmMemory { offset },
            ),
            Self::StableMemory => (
                SnapshotDataKind::StableMemory { offset, size },
                SnapshotDataOffset::StableMemory { offset },
            ),
        }
    }
}

struct SnapshotTransfer<'a> {
    pic: &'a PocketIc,
    canister_id: Principal,
    sender: Option<Principal>,
    original: &'a [u8],
    uploaded: &'a [u8],
}

impl SnapshotTransfer<'_> {
    fn read(&self, kind: SnapshotDataKind) -> Result<Vec<u8>, Failure> {
        let (data,): (ReadCanisterSnapshotDataResult,) = call(
            self.pic,
            self.canister_id,
            self.sender,
            "read_canister_snapshot_data",
            ReadCanisterSnapshotDataArgs {
                canister_id: self.canister_id,
                snapshot_id: self.original.to_vec(),
                kind,
            },
        )?;
        Ok(data.chunk)
    }

    fn write(&self, kind: SnapshotDataOffset, chunk: Vec<u8>) -> Result<(), Failure> {
        call(
            self.pic,
            self.canister_id,
            self.sender,
            "upload_canister_snapshot_data",
            UploadCanisterSnapshotDataArgs {
                canister_id: self.canister_id,
                snapshot_id: self.uploaded.to_vec(),
                kind,
                chunk,
            },
        )
    }

    fn copy_region(&self, bytes: u64, region: Region) -> Result<(), Failure> {
        let mut offset = 0;
        while offset < bytes {
            let size = CHUNK_BYTES.min(bytes - offset);
            let (read_kind, write_kind) = region.at(offset, size);
            let data = self.read(read_kind)?;
            if data.len() as u64 != size {
                return Err(Failure::Invalid(
                    "snapshot returned an incomplete data region",
                ));
            }
            self.write(write_kind, data)?;
            offset += size;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{Failure, decode_reply, require_complete_system_state};
    use ic_platform::candid::contract;
    use ic_platform::management_canister_types::{
        CanisterTimer, OnLowWasmMemoryHookStatus, ReadCanisterSnapshotMetadataResult,
    };

    #[test]
    fn omitted_timer_or_hook_is_refused_instead_of_silently_preserved_on_restore() {
        let mut metadata = ReadCanisterSnapshotMetadataResult {
            source: None,
            taken_at_timestamp: 0,
            wasm_module_size: 0,
            globals: vec![],
            wasm_memory_size: 0,
            stable_memory_size: 0,
            wasm_chunk_store: vec![],
            canister_version: 0,
            certified_data: vec![],
            global_timer: Some(CanisterTimer::Inactive),
            on_low_wasm_memory_hook_status: Some(OnLowWasmMemoryHookStatus::ConditionNotSatisfied),
        };
        assert!(require_complete_system_state(&metadata).is_ok());
        metadata.global_timer = None;
        assert!(matches!(
            require_complete_system_state(&metadata),
            Err(Failure::Invalid(_))
        ));
        metadata.global_timer = Some(CanisterTimer::Inactive);
        metadata.on_low_wasm_memory_hook_status = None;
        assert!(matches!(
            require_complete_system_state(&metadata),
            Err(Failure::Invalid(_))
        ));
    }

    #[test]
    fn snapshot_replies_refuse_oversize_malformed_and_wrong_shapes() {
        assert!(matches!(
            decode_reply::<()>(&vec![0; 2_000_001]),
            Err(Failure::Invalid(_))
        ));
        assert!(matches!(
            decode_reply::<()>(b"malformed"),
            Err(Failure::Codec(_))
        ));
        let empty = contract::encode_args(()).unwrap();
        assert!(decode_reply::<()>(&empty).is_ok());
        assert!(matches!(
            decode_reply::<(u64,)>(&empty),
            Err(Failure::Codec(_))
        ));
    }
}