ic-backup 0.31.2

//! Decode-only pinned wire fields with bounded allocation and no skipped work.

use super::MAX_IC_SNAPSHOT_METADATA_BYTES;
use super::{IcSnapshotMetadataError, MAX_IC_SNAPSHOT_CHUNKS, MAX_IC_SNAPSHOT_GLOBALS};
use crate::model::bounded_sequence;
use ic_platform::candid::contract;
use ic_platform::candid::contract::CandidType;
use ic_platform::candid::{DecodeLimits, decode_exactly_one};
use ic_platform::management_canister_types::{
    CanisterTimer, ChunkHash, OnLowWasmMemoryHookStatus, ReadCanisterSnapshotMetadataResult,
    SnapshotMetadataGlobal, SnapshotSource,
};
use serde::{Deserialize, Deserializer, de};
use std::{collections::BTreeSet, fmt};

#[derive(CandidType, Deserialize)]
#[candid_path("ic_platform::candid::contract")]
struct WireMetadata {
    source: Option<WireSource>,
    taken_at_timestamp: u64,
    wasm_module_size: u64,
    #[serde(deserialize_with = "globals")]
    globals: Vec<Option<SnapshotMetadataGlobal>>,
    wasm_memory_size: u64,
    stable_memory_size: u64,
    #[serde(deserialize_with = "chunks")]
    wasm_chunk_store: Vec<WireChunk>,
    canister_version: u64,
    #[serde(deserialize_with = "bytes32")]
    certified_data: Vec<u8>,
    global_timer: Option<CanisterTimer>,
    on_low_wasm_memory_hook_status: Option<OnLowWasmMemoryHookStatus>,
}

#[derive(CandidType, Deserialize)]
#[candid_path("ic_platform::candid::contract")]
struct WireChunk {
    #[serde(deserialize_with = "bytes32")]
    hash: Vec<u8>,
}

#[derive(CandidType, Deserialize)]
#[candid_path("ic_platform::candid::contract")]
enum WireSource {
    #[serde(rename = "taken_from_canister")]
    TakenFromCanister(WireReserved),
    #[serde(rename = "metadata_upload")]
    MetadataUpload(WireReserved),
}

// The SDK's Reserved deserializer deliberately skips arbitrary values. Read the
// actual Candid reserved marker instead so recognized source tags need no skipped
// work and cannot silently discard an unknown payload under our zero-skip quota.
struct WireReserved;

impl CandidType for WireReserved {
    fn _ty() -> contract::types::Type {
        contract::types::TypeInner::Reserved.into()
    }
    fn idl_serialize<S: contract::types::Serializer>(&self, serializer: S) -> Result<(), S::Error> {
        serializer.serialize_null(())
    }
}

impl<'de> Deserialize<'de> for WireReserved {
    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        struct ReservedVisitor;
        impl de::Visitor<'_> for ReservedVisitor {
            type Value = WireReserved;
            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.write_str("the Candid reserved marker")
            }
            fn visit_byte_buf<E: de::Error>(self, bytes: Vec<u8>) -> Result<Self::Value, E> {
                if bytes == [3] {
                    Ok(WireReserved)
                } else {
                    Err(E::custom("invalid reserved marker"))
                }
            }
        }
        d.deserialize_any(ReservedVisitor)
    }
}

fn globals<'de, D: Deserializer<'de>>(
    d: D,
) -> Result<Vec<Option<SnapshotMetadataGlobal>>, D::Error> {
    sequence::<_, _, MAX_IC_SNAPSHOT_GLOBALS>(d)
}

fn chunks<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<WireChunk>, D::Error> {
    sequence::<_, _, MAX_IC_SNAPSHOT_CHUNKS>(d)
}

fn bytes32<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
    sequence::<_, _, 32>(d)
}

// Never allocate using an untrusted size hint. Each retained vector has one
// explicit bound; the decoder's separate quota also bounds value/type work.
fn sequence<'de, D, T, const LIMIT: usize>(d: D) -> Result<Vec<T>, D::Error>
where
    D: Deserializer<'de>,
    T: Deserialize<'de>,
{
    bounded_sequence::deserialize_wire::<_, T, LIMIT>(
        d,
        format_args!("at most {LIMIT} metadata elements"),
        "metadata sequence exceeds bound",
    )
}

pub(super) fn decode(
    bytes: &[u8],
) -> Result<ReadCanisterSnapshotMetadataResult, IcSnapshotMetadataError> {
    let value: WireMetadata = decode_exactly_one(
        bytes,
        DecodeLimits {
            reply_bytes: MAX_IC_SNAPSHOT_METADATA_BYTES,
            decoding_quota: 2 * 1024 * 1024,
            skipping_quota: 0,
            type_entries: 64,
            header_bytes: 16 * 1024,
        },
    )
    .map_err(|_| IcSnapshotMetadataError::InvalidReply)?;
    if value
        .globals
        .iter()
        .flatten()
        .any(|global| matches!(global, SnapshotMetadataGlobal::V128(n) if n.0.bits() > 128))
    {
        return Err(IcSnapshotMetadataError::InvalidGlobal);
    }
    let mut hashes = BTreeSet::new();
    for chunk in &value.wasm_chunk_store {
        if chunk.hash.len() != 32 || !hashes.insert(chunk.hash.as_slice()) {
            return Err(IcSnapshotMetadataError::InvalidChunkHash);
        }
    }
    Ok(ReadCanisterSnapshotMetadataResult {
        source: value.source.map(|source| match source {
            WireSource::TakenFromCanister(_) => {
                SnapshotSource::TakenFromCanister(contract::Reserved)
            }
            WireSource::MetadataUpload(_) => SnapshotSource::MetadataUpload(contract::Reserved),
        }),
        taken_at_timestamp: value.taken_at_timestamp,
        wasm_module_size: value.wasm_module_size,
        globals: value.globals,
        wasm_memory_size: value.wasm_memory_size,
        stable_memory_size: value.stable_memory_size,
        wasm_chunk_store: value
            .wasm_chunk_store
            .into_iter()
            .map(|c| ChunkHash { hash: c.hash })
            .collect(),
        canister_version: value.canister_version,
        certified_data: value.certified_data,
        global_timer: value.global_timer,
        on_low_wasm_memory_hook_status: value.on_low_wasm_memory_hook_status,
    })
}