ic-backup 0.24.6

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

use super::{IcSnapshotMetadataError, MAX_IC_SNAPSHOT_CHUNKS, MAX_IC_SNAPSHOT_GLOBALS};
use crate::model::ic_wire::{self, DecoderLimits};
use candid::CandidType;
use ic_management_canister_types::{
    CanisterTimer, ChunkHash, OnLowWasmMemoryHookStatus, ReadCanisterSnapshotMetadataResult,
    SnapshotMetadataGlobal, SnapshotSource,
};
use serde::{Deserialize, Deserializer, de};
use std::{collections::BTreeSet, fmt};

#[derive(CandidType, Deserialize)]
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)]
struct WireChunk {
    #[serde(deserialize_with = "bytes32")]
    hash: Vec<u8>,
}

#[derive(CandidType, Deserialize)]
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() -> candid::types::Type {
        candid::types::TypeInner::Reserved.into()
    }
    fn idl_serialize<S: candid::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>,
{
    ic_wire::bounded_sequence::deserialize::<_, T, LIMIT>(
        d,
        |f, limit| write!(f, "at most {limit} metadata elements"),
        "metadata sequence exceeds bound",
    )
}

pub(super) fn decode(
    bytes: &[u8],
) -> Result<ReadCanisterSnapshotMetadataResult, IcSnapshotMetadataError> {
    let value: WireMetadata = ic_wire::decode_one(
        bytes,
        DecoderLimits {
            decoding: 2 * 1024 * 1024,
            skipping: 0,
            type_table: 64,
            header: Some(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(candid::Reserved),
            WireSource::MetadataUpload(_) => SnapshotSource::MetadataUpload(candid::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,
    })
}