ic-platform 0.5.0

use ic_certification::{
    Certificate, Delegation, HashTree, HashTreeNode, LookupResult, SubtreeLookupResult, leaf,
};
use ic_principal::Principal;
use ic_verify_bls_signature::verify_bls_signature;
use serde::{Deserialize, Serialize};
use serde_bytes::ByteBuf;
use sha2::{Digest, Sha256};

pub const IC_STATE_ROOT_DOMAIN_SEPARATOR: &[u8; 14] = b"\x0Dic-state-root";

use super::creation::{CanisterSigPublicKey, extract_raw_root_pk_from_der};

/// Verifies that `signature` is a valid canister signature on `message`.
/// https://internetcomputer.org/docs/current/references/ic-interface-spec#canister-signatures
pub fn verify_canister_sig(
    message: &[u8],
    signature_cbor: &[u8],
    public_key_der: &[u8],
    ic_root_public_key_raw: &[u8],
) -> Result<(), String> {
    let signature = parse_signature_cbor(signature_cbor)?;
    let public_key = CanisterSigPublicKey::try_from(public_key_der)
        .map_err(|e| format!("failed to parse canister sig public key: {e}"))?;
    let certificate =
        check_certified_data_and_get_certificate(&signature, &public_key.canister_id)?;
    check_sig_path(&signature, &public_key, message)?;
    verify_certificate(&certificate, public_key.canister_id, ic_root_public_key_raw)
}

// Check that signature.certificate's tree contains for the canister identified by
// signing_canister_id an entry for certified_data that matches signature.tree.digest.
fn check_certified_data_and_get_certificate(
    signature: &CanisterSignature,
    signing_canister_id: &Principal,
) -> Result<Certificate, String> {
    let certificate = parse_certificate_cbor(&signature.certificate)?;
    let cert_data_path = [
        "canister".as_bytes(),
        signing_canister_id.as_slice(),
        "certified_data".as_bytes(),
    ];
    let SubtreeLookupResult::Found(cert_data_leaf) =
        certificate.tree.lookup_subtree(&cert_data_path)
    else {
        return Err("certified_data entry not found".to_string());
    };
    if cert_data_leaf != leaf(signature.tree.digest()) {
        return Err("certified_data doesn't match sig tree digest".to_string());
    }
    Ok(certificate)
}

// Check that signature.tree contains an empty leaf at correct "sig"-path,
// where the path is determined by hashes of canister_sig_pk.seed and msg.
fn check_sig_path(
    signature: &CanisterSignature,
    canister_sig_pk: &CanisterSigPublicKey,
    msg: &[u8],
) -> Result<(), String> {
    let seed_hash = hash_sha256(&canister_sig_pk.seed);
    let msg_hash = hash_sha256(msg);
    let sig_path = ["sig".as_bytes(), &seed_hash, &msg_hash];
    let SubtreeLookupResult::Found(sig_leaf) = signature.tree.lookup_subtree(&sig_path) else {
        return Err("signature entry not found".to_string());
    };
    if sig_leaf != leaf(b"") {
        return Err("signature entry is not an empty leaf".to_string());
    }
    Ok(())
}

#[derive(Debug, Serialize, Deserialize)]
struct CanisterSignature {
    pub certificate: ByteBuf,
    pub tree: HashTree,
}

fn parse_signature_cbor(signature_cbor: &[u8]) -> Result<CanisterSignature, String> {
    // 0xd9d9f7 (cf. https://tools.ietf.org/html/rfc7049#section-2.4.5) is the
    // self-describing CBOR tag required to be present by the interface spec.
    if signature_cbor.len() < 3 || signature_cbor[0..3] != [0xd9, 0xd9, 0xf7] {
        return Err("signature CBOR doesn't have a self-describing tag".to_string());
    }
    serde_cbor::from_slice::<CanisterSignature>(signature_cbor)
        .map_err(|e| format!("failed to parse signature CBOR: {e}"))
}

fn parse_certificate_cbor(certificate_cbor: &[u8]) -> Result<Certificate, String> {
    // 0xd9d9f7 (cf. https://tools.ietf.org/html/rfc7049#section-2.4.5) is the
    // self-describing CBOR tag required to be present by the interface spec.
    if certificate_cbor.len() < 3 || certificate_cbor[0..3] != [0xd9, 0xd9, 0xf7] {
        return Err("certificate CBOR doesn't have a self-describing tag".to_string());
    }
    serde_cbor::from_slice::<Certificate>(certificate_cbor)
        .map_err(|e| format!("failed to parse certificate CBOR: {e}"))
}

fn verify_certificate(
    certificate: &Certificate,
    signing_canister_id: Principal,
    root_public_key_raw: &[u8],
) -> Result<(), String> {
    let bls_pk_raw = match &certificate.delegation {
        Some(delegation) => {
            verify_delegation(delegation, signing_canister_id, root_public_key_raw)?
        }
        _ => root_public_key_raw.to_vec(),
    };
    check_bls_signature(certificate, &bls_pk_raw)
}

fn verify_delegation(
    delegation: &Delegation,
    signing_canister_id: Principal,
    root_public_key: &[u8],
) -> Result<Vec<u8>, String> {
    let cert: Certificate = parse_certificate_cbor(&delegation.certificate)
        .map_err(|e| format!("invalid delegation certificate: {e}"))?;

    // disallow nested delegations
    if cert.delegation.is_some() {
        return Err("multiple delegations not allowed".to_string());
    }

    check_bls_signature(&cert, root_public_key)?;

    // Try new structure first: /canister_ranges/<subnet_id>/<range_key>
    let canister_ranges_path = ["canister_ranges".as_bytes(), delegation.subnet_id.as_ref()];
    let canister_in_ranges = match cert.tree.lookup_subtree(&canister_ranges_path) {
        SubtreeLookupResult::Found(subnet_tree) => {
            match find_leaf_for_principal(subnet_tree.as_ref(), signing_canister_id.as_slice()) {
                Some(range_data) => {
                    let subnet_ranges: Vec<(Principal, Principal)> =
                        serde_cbor::from_slice(range_data)
                            .map_err(|e| format!("invalid canister range: {e}"))?;
                    principal_is_within_ranges(&signing_canister_id, &subnet_ranges)
                }
                None => false,
            }
        }
        SubtreeLookupResult::Absent | SubtreeLookupResult::Unknown => {
            let old_canister_ranges_path = [
                "subnet".as_bytes(),
                delegation.subnet_id.as_ref(),
                "canister_ranges".as_bytes(),
            ];
            match cert.tree.lookup_path(&old_canister_ranges_path) {
                LookupResult::Found(old_range_data) => {
                    let canister_ranges: Vec<(Principal, Principal)> =
                        serde_cbor::from_slice(old_range_data)
                            .map_err(|e| format!("invalid canister range: {e}"))?;
                    principal_is_within_ranges(&signing_canister_id, &canister_ranges)
                }
                _ => {
                    return Err("canister_ranges-entry not found".to_string());
                }
            }
        }
    };

    if !canister_in_ranges {
        return Err("signing canister id not in canister_ranges".to_string());
    }

    // lookup the public key delegated to
    let public_key_path = [
        "subnet".as_bytes(),
        delegation.subnet_id.as_ref(),
        "public_key".as_bytes(),
    ];
    let LookupResult::Found(subnet_public_key_der) = cert.tree.lookup_path(&public_key_path) else {
        return Err("subnet public key not found".to_string());
    };
    extract_raw_root_pk_from_der(subnet_public_key_der)
}

fn principal_is_within_ranges(principal: &Principal, ranges: &[(Principal, Principal)]) -> bool {
    ranges
        .iter()
        .any(|(start, end)| (start..=end).contains(&principal))
}

/// Walks the one-level-deep subtree under `/canister_ranges/<subnet_id>/` to
/// find the shard whose key (label) is the largest one that is <= `target`.
/// Per the IC spec, shard keys are the starting canister ID of the first range
/// in each shard, sorted in ascending order, so the correct shard for any
/// canister ID is the one with the greatest key not exceeding that ID.
///
/// Returns the leaf data (CBOR-encoded ranges) of the matching shard, or `None`
/// if no label <= `target` exists in the tree (which includes pruned/empty
/// subtrees).
fn find_leaf_for_principal<'a>(node: &'a HashTreeNode, target: &[u8]) -> Option<&'a [u8]> {
    match node {
        HashTreeNode::Empty() | HashTreeNode::Pruned(_) => None,
        HashTreeNode::Leaf(data) => Some(data),
        HashTreeNode::Labeled(label, node) => {
            if label.as_bytes() <= target {
                find_leaf_for_principal(node, target)
            } else {
                None
            }
        }
        HashTreeNode::Fork(children) => {
            let (left, right) = children.as_ref();
            // We want the *largest* label <= target (greatest lower bound).
            // In a well-formed hash tree, every label in the right subtree is
            // strictly greater than every label in the left, so a match in the
            // right subtree is always >= any match in the left. Try the right
            // subtree first and only fall back to the left if it has no match.
            find_leaf_for_principal(right, target).or_else(|| find_leaf_for_principal(left, target))
        }
    }
}

const SHA256_DIGEST_LEN: usize = 32;
fn hash_sha256(data: &[u8]) -> [u8; SHA256_DIGEST_LEN] {
    let mut hash = Sha256::default();
    hash.update(data);
    <[u8; SHA256_DIGEST_LEN]>::from(hash.finalize())
}

fn check_bls_signature(certificate: &Certificate, signing_pk_raw: &[u8]) -> Result<(), String> {
    let sig = certificate.signature.as_slice();
    let root_hash = certificate.tree.digest();
    let mut msg = vec![];
    msg.extend_from_slice(IC_STATE_ROOT_DOMAIN_SEPARATOR);
    msg.extend_from_slice(&root_hash);
    if verify_bls_signature(sig, &msg, signing_pk_raw).is_err() {
        return Err("invalid BLS signature".to_string());
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use ic_certification::{empty, fork, labeled, leaf, pruned};

    #[test]
    fn find_leaf_empty_tree_returns_none() {
        let tree = empty();
        assert_eq!(find_leaf_for_principal(tree.as_ref(), b"\x10"), None);
    }

    #[test]
    fn find_leaf_pruned_tree_returns_none() {
        let tree = pruned([0_u8; 32]);
        assert_eq!(find_leaf_for_principal(tree.as_ref(), b"\x10"), None);
    }

    #[test]
    fn find_leaf_single_shard_target_above_label() {
        let tree = labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec()));
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x20"),
            Some(b"data_A".as_ref())
        );
    }

    #[test]
    fn find_leaf_single_shard_target_equals_label() {
        let tree = labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec()));
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x10"),
            Some(b"data_A".as_ref())
        );
    }

    #[test]
    fn find_leaf_single_shard_target_below_label() {
        let tree = labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec()));
        assert_eq!(find_leaf_for_principal(tree.as_ref(), b"\x05"), None);
    }

    #[test]
    fn find_leaf_two_shards_target_in_first() {
        let tree = fork(
            labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
            labeled(b"\x30".to_vec(), leaf(b"data_B".to_vec())),
        );
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x20"),
            Some(b"data_A".as_ref())
        );
    }

    #[test]
    fn find_leaf_two_shards_target_in_second() {
        let tree = fork(
            labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
            labeled(b"\x30".to_vec(), leaf(b"data_B".to_vec())),
        );
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x40"),
            Some(b"data_B".as_ref())
        );
    }

    #[test]
    fn find_leaf_two_shards_target_equals_second_label() {
        let tree = fork(
            labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
            labeled(b"\x30".to_vec(), leaf(b"data_B".to_vec())),
        );
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x30"),
            Some(b"data_B".as_ref())
        );
    }

    #[test]
    fn find_leaf_two_shards_target_before_all() {
        let tree = fork(
            labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
            labeled(b"\x30".to_vec(), leaf(b"data_B".to_vec())),
        );
        assert_eq!(find_leaf_for_principal(tree.as_ref(), b"\x05"), None);
    }

    #[test]
    fn find_leaf_three_shards_target_in_middle() {
        let tree = fork(
            fork(
                labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
                labeled(b"\x20".to_vec(), leaf(b"data_B".to_vec())),
            ),
            labeled(b"\x30".to_vec(), leaf(b"data_C".to_vec())),
        );
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\x25"),
            Some(b"data_B".as_ref())
        );
    }

    #[test]
    fn find_leaf_three_shards_target_in_last() {
        let tree = fork(
            fork(
                labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
                labeled(b"\x20".to_vec(), leaf(b"data_B".to_vec())),
            ),
            labeled(b"\x30".to_vec(), leaf(b"data_C".to_vec())),
        );
        assert_eq!(
            find_leaf_for_principal(tree.as_ref(), b"\xff"),
            Some(b"data_C".as_ref())
        );
    }

    #[test]
    fn find_leaf_three_shards_target_before_all() {
        let tree = fork(
            fork(
                labeled(b"\x10".to_vec(), leaf(b"data_A".to_vec())),
                labeled(b"\x20".to_vec(), leaf(b"data_B".to_vec())),
            ),
            labeled(b"\x30".to_vec(), leaf(b"data_C".to_vec())),
        );
        assert_eq!(find_leaf_for_principal(tree.as_ref(), b"\x05"), None);
    }
}