// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Tests of the protobuf conversion logic

use super::specs::types::*;
use super::*;
use crate::{AzksValue, Direction};
use rand::{thread_rng, Rng};

// ================= Test helpers ================= //

fn random_hash() -> [u8; 32] {
    thread_rng().gen::<[u8; 32]>()
}

fn random_azks_element() -> crate::AzksElement {
    crate::AzksElement {
        label: random_label(),
        value: AzksValue(random_hash()),
    }
}

fn random_label() -> crate::NodeLabel {
    crate::NodeLabel {
        label_val: random_hash(),
        label_len: thread_rng().gen::<u32>(),
    }
}

// ================= Test cases ================= //

#[test]
fn test_convert_nodelabel() {
    let original = random_label();

    let protobuf: NodeLabel = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_azks_element() {
    let original = random_azks_element();

    let protobuf: AzksElement = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_layer_proof() {
    let original = crate::SiblingProof {
        label: random_label(),
        siblings: [random_azks_element()],
        direction: Direction::Right,
    };

    let protobuf: SiblingProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_membership_proof() {
    let original = crate::MembershipProof {
        label: random_label(),
        hash_val: AzksValue(random_hash()),
        sibling_proofs: vec![crate::SiblingProof {
            label: random_label(),
            siblings: [random_azks_element()],
            direction: Direction::Right,
        }],
    };

    let protobuf: MembershipProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_non_membership_proof() {
    let original = crate::NonMembershipProof {
        label: random_label(),
        longest_prefix: random_label(),
        longest_prefix_children: [random_azks_element(), random_azks_element()],
        longest_prefix_membership_proof: crate::MembershipProof {
            label: random_label(),
            hash_val: AzksValue(random_hash()),
            sibling_proofs: vec![crate::SiblingProof {
                label: random_label(),
                siblings: [random_azks_element()],
                direction: Direction::Right,
            }],
        },
    };

    let protobuf: NonMembershipProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_lookup_proof() {
    let mut rng = thread_rng();
    let original = crate::LookupProof {
        epoch: rng.gen(),
        value: crate::AkdValue(random_hash().to_vec()),
        version: rng.gen(),
        existence_vrf_proof: random_hash().to_vec(),
        existence_proof: crate::MembershipProof {
            label: random_label(),
            hash_val: AzksValue(random_hash()),
            sibling_proofs: vec![crate::SiblingProof {
                label: random_label(),
                siblings: [random_azks_element()],
                direction: Direction::Right,
            }],
        },
        marker_vrf_proof: random_hash().to_vec(),
        marker_proof: crate::MembershipProof {
            label: random_label(),
            hash_val: AzksValue(random_hash()),
            sibling_proofs: vec![crate::SiblingProof {
                label: random_label(),
                siblings: [random_azks_element()],
                direction: Direction::Right,
            }],
        },
        freshness_vrf_proof: random_hash().to_vec(),
        freshness_proof: crate::NonMembershipProof {
            label: random_label(),
            longest_prefix: random_label(),
            longest_prefix_children: [random_azks_element(), random_azks_element()],
            longest_prefix_membership_proof: crate::MembershipProof {
                label: random_label(),
                hash_val: AzksValue(random_hash()),
                sibling_proofs: vec![crate::SiblingProof {
                    label: random_label(),
                    siblings: [random_azks_element()],
                    direction: Direction::Right,
                }],
            },
        },
        commitment_nonce: random_hash().to_vec(),
    };

    let protobuf: LookupProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_update_proof() {
    let mut rng = thread_rng();
    let original = crate::UpdateProof {
        epoch: rng.gen(),
        value: crate::AkdValue(random_hash().to_vec()),
        version: rng.gen(),
        existence_vrf_proof: random_hash().to_vec(),
        existence_proof: crate::MembershipProof {
            label: random_label(),
            hash_val: AzksValue(random_hash()),
            sibling_proofs: vec![crate::SiblingProof {
                label: random_label(),
                siblings: [random_azks_element()],
                direction: Direction::Right,
            }],
        },
        previous_version_vrf_proof: Some(random_hash().to_vec()),
        previous_version_proof: Some(crate::MembershipProof {
            label: random_label(),
            hash_val: AzksValue(random_hash()),
            sibling_proofs: vec![crate::SiblingProof {
                label: random_label(),
                siblings: [random_azks_element()],
                direction: Direction::Right,
            }],
        }),
        commitment_nonce: random_hash().to_vec(),
    };

    let protobuf: UpdateProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_history_proof() {
    fn non_membership_proof() -> crate::NonMembershipProof {
        crate::NonMembershipProof {
            label: random_label(),
            longest_prefix: random_label(),
            longest_prefix_children: [random_azks_element(), random_azks_element()],
            longest_prefix_membership_proof: crate::MembershipProof {
                label: random_label(),
                hash_val: AzksValue(random_hash()),
                sibling_proofs: vec![crate::SiblingProof {
                    label: random_label(),
                    siblings: [random_azks_element()],
                    direction: Direction::Right,
                }],
            },
        }
    }

    fn upd_proof() -> crate::UpdateProof {
        let mut rng = thread_rng();
        crate::UpdateProof {
            epoch: rng.gen(),
            value: crate::AkdValue(random_hash().to_vec()),
            version: rng.gen(),
            existence_vrf_proof: random_hash().to_vec(),
            existence_proof: crate::MembershipProof {
                label: random_label(),
                hash_val: AzksValue(random_hash()),
                sibling_proofs: vec![crate::SiblingProof {
                    label: random_label(),
                    siblings: [random_azks_element()],
                    direction: Direction::Right,
                }],
            },
            previous_version_vrf_proof: Some(random_hash().to_vec()),
            previous_version_proof: Some(crate::MembershipProof {
                label: random_label(),
                hash_val: AzksValue(random_hash()),
                sibling_proofs: vec![crate::SiblingProof {
                    label: random_label(),
                    siblings: [random_azks_element()],
                    direction: Direction::Right,
                }],
            }),
            commitment_nonce: random_hash().to_vec(),
        }
    }

    let original = crate::HistoryProof {
        update_proofs: vec![upd_proof(), upd_proof(), upd_proof()],
        until_marker_vrf_proofs: vec![
            random_hash().to_vec(),
            random_hash().to_vec(),
            random_hash().to_vec(),
        ],
        non_existence_until_marker_proofs: vec![
            non_membership_proof(),
            non_membership_proof(),
            non_membership_proof(),
        ],
        future_marker_vrf_proofs: vec![
            random_hash().to_vec(),
            random_hash().to_vec(),
            random_hash().to_vec(),
        ],
        non_existence_of_future_marker_proofs: vec![
            non_membership_proof(),
            non_membership_proof(),
            non_membership_proof(),
        ],
    };

    let protobuf: HistoryProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_convert_single_append_only_proof() {
    let inserted = [
        random_azks_element(),
        random_azks_element(),
        random_azks_element(),
    ];
    let unchanged_nodes = [
        random_azks_element(),
        random_azks_element(),
        random_azks_element(),
        random_azks_element(),
        random_azks_element(),
        random_azks_element(),
    ];
    let original = crate::SingleAppendOnlyProof {
        inserted: inserted.to_vec(),
        unchanged_nodes: unchanged_nodes.to_vec(),
    };

    let protobuf: SingleAppendOnlyProof = (&original).into();
    assert_eq!(original, (&protobuf).try_into().unwrap());
}

#[test]
fn test_minimum_encoding_label_bytes() {
    let full_label: [u8; 32] = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 1,
    ];

    let half_label: [u8; 32] = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0,
    ];

    let zero_label: [u8; 32] = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0,
    ];

    let min_full_label = encode_minimum_label(&full_label);
    let min_half_label = encode_minimum_label(&half_label);
    let min_zero_label = encode_minimum_label(&zero_label);

    assert_eq!(32, min_full_label.len());
    assert_eq!(16, min_half_label.len());
    assert_eq!(0, min_zero_label.len());

    assert_eq!(full_label, decode_minimized_label(&min_full_label));
    assert_eq!(half_label, decode_minimized_label(&min_half_label));
    assert_eq!(zero_label, decode_minimized_label(&min_zero_label));
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! @generated code

include!(concat!(env!("OUT_DIR"), "/protos/mod.rs"));


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! @generated code

include!(concat!(env!("OUT_DIR"), "/protos/types.rs"));


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains all of the structs which need to be constructed
//! to verify any of the following AKD proofs
//!
//! 1. Lookup
//! 2: Key history
//! 3: Audit (append-only)

use crate::hash::Digest;
#[cfg(feature = "serde_serialization")]
use crate::utils::serde_helpers::{
    azks_value_hex_deserialize, azks_value_hex_serialize, bytes_deserialize_hex,
    bytes_serialize_hex,
};
use crate::ARITY;

#[cfg(feature = "nostd")]
use alloc::string::{String, ToString};
#[cfg(feature = "nostd")]
use alloc::vec::Vec;
#[cfg(feature = "nostd")]
use core::cmp::{Ord, Ordering, PartialOrd};
#[cfg(feature = "rand")]
use rand::{CryptoRng, Rng};
#[cfg(not(feature = "nostd"))]
use std::cmp::{Ord, Ordering, PartialOrd};

pub mod node_label;
pub use node_label::*;

// ============================================
// Traits
// ============================================

/// Retrieve the in-memory size of a structure
pub trait SizeOf {
    /// Retrieve the in-memory size of a structure
    fn size_of(&self) -> usize;
}

// ============================================
// Typedefs and constants
// ============================================

/// Whether or not a node is marked as stale or fresh
/// Stale nodes are no longer active because a newer
/// version exists to replace them.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum VersionFreshness {
    /// Represents not being the most recent version
    Stale = 0u8,
    /// Corresponds to the most recent version
    Fresh = 1u8,
}

/// This type is used to indicate whether or not
/// one label is a prefix of another, and if so,
/// whether the longer string has a 0 after the prefix,
/// or a 1 after the prefix. If the first label is equal
/// to the second, or not a prefix of the second, then it
/// is considered invalid.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
#[repr(u8)]
pub enum PrefixOrdering {
    /// Corresponds to a [Direction::Left]
    WithZero = 0u8,
    /// Corresponds to a [Direction::Right]
    WithOne = 1u8,
    /// First label is either equal to the second, or
    /// simply not a prefix of the second
    Invalid = u8::MAX,
}

impl SizeOf for PrefixOrdering {
    fn size_of(&self) -> usize {
        // The size of the enum is 24 bytes. The extra 8 bytes are used to store a 64-bit
        // discriminator that is used to identify the variant currently saved in the enum.
        24usize
    }
}

impl From<Bit> for PrefixOrdering {
    fn from(bit: Bit) -> Self {
        match bit {
            Bit::Zero => Self::WithZero,
            Bit::One => Self::WithOne,
        }
    }
}

/// This type is used to indicate a direction for a
/// particular node relative to its parent. We use
/// 0 to represent "left" and 1 to represent "right".
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
#[repr(u8)]
pub enum Direction {
    /// Left
    Left = 0u8,
    /// Right
    Right = 1u8,
}

impl SizeOf for Direction {
    fn size_of(&self) -> usize {
        // The size of the enum is 24 bytes. The extra 8 bytes are used to store a 64-bit
        // discriminator that is used to identify the variant currently saved in the enum.
        24usize
    }
}

impl From<Bit> for Direction {
    fn from(bit: Bit) -> Self {
        match bit {
            Bit::Zero => Self::Left,
            Bit::One => Self::Right,
        }
    }
}

impl core::convert::TryFrom<PrefixOrdering> for Direction {
    type Error = String;
    fn try_from(prefix_ordering: PrefixOrdering) -> Result<Self, Self::Error> {
        match prefix_ordering {
            PrefixOrdering::WithZero => Ok(Direction::Left),
            PrefixOrdering::WithOne => Ok(Direction::Right),
            _ => Err("Could not convert from PrefixOrdering to Direction".to_string()),
        }
    }
}

impl Direction {
    /// Returns the opposite of the direction
    pub fn other(&self) -> Self {
        match self {
            Direction::Left => Direction::Right,
            Direction::Right => Direction::Left,
        }
    }
}

/// The label of a particular entry in the AKD
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct AkdLabel(
    #[cfg_attr(
        feature = "serde_serialization",
        serde(serialize_with = "bytes_serialize_hex")
    )]
    #[cfg_attr(
        feature = "serde_serialization",
        serde(deserialize_with = "bytes_deserialize_hex")
    )]
    pub Vec<u8>,
);

impl SizeOf for AkdLabel {
    fn size_of(&self) -> usize {
        self.0.len()
    }
}

impl core::ops::Deref for AkdLabel {
    type Target = Vec<u8>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl core::ops::DerefMut for AkdLabel {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl core::convert::From<&str> for AkdLabel {
    fn from(s: &str) -> Self {
        Self(s.as_bytes().to_vec())
    }
}

impl core::convert::From<&String> for AkdLabel {
    fn from(s: &String) -> Self {
        Self(s.as_bytes().to_vec())
    }
}

impl AkdLabel {
    #[cfg(feature = "rand")]
    /// Gets a random label
    pub fn random<R: CryptoRng + Rng>(rng: &mut R) -> Self {
        let mut bytes = [0u8; 32];
        rng.fill_bytes(&mut bytes);
        Self(bytes.to_vec())
    }
}

/// The value of a particular entry in the AKD
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct AkdValue(
    #[cfg_attr(
        feature = "serde_serialization",
        serde(serialize_with = "bytes_serialize_hex")
    )]
    #[cfg_attr(
        feature = "serde_serialization",
        serde(deserialize_with = "bytes_deserialize_hex")
    )]
    pub Vec<u8>,
);

impl SizeOf for AkdValue {
    fn size_of(&self) -> usize {
        self.0.len()
    }
}

impl core::ops::Deref for AkdValue {
    type Target = Vec<u8>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl core::ops::DerefMut for AkdValue {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl core::convert::From<&str> for AkdValue {
    fn from(s: &str) -> Self {
        Self(s.as_bytes().to_vec())
    }
}

impl core::convert::From<&String> for AkdValue {
    fn from(s: &String) -> Self {
        Self(s.as_bytes().to_vec())
    }
}

impl AkdValue {
    #[cfg(feature = "rand")]
    /// Gets a random value for a AKD
    pub fn random<R: CryptoRng + Rng>(rng: &mut R) -> Self {
        let mut bytes = [0u8; 32];
        rng.fill_bytes(&mut bytes);
        Self(bytes.to_vec())
    }
}

/// The value to be hashed every time an empty node's hash is to be considered
pub const EMPTY_VALUE: [u8; 1] = [0u8];

/// A "tombstone" is a false value in an AKD ValueState denoting that a real value has been removed (e.g. data rentention policies).
/// Should a tombstone be encountered, we have to assume that the hash of the value is correct, and we move forward without being able to
/// verify the raw value. We utilize an empty array to save space in the storage layer
///
/// See [GitHub issue #130](https://github.com/novifinancial/akd/issues/130) for more context
pub const TOMBSTONE: &[u8] = &[];

// ============================================
// Structs
// ============================================

/// The value associated with an element of the AZKS
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct AzksValue(pub Digest);

/// Used to denote an azks value that has been hashed together with an epoch
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct AzksValueWithEpoch(pub Digest);

/// Represents an element to be inserted into the AZKS. This
/// is a pair consisting of a label ([NodeLabel]) and a value.
/// The purpose of the directory publish is to convert an
/// insertion set of ([AkdLabel], [AkdValue]) tuples into a
/// set of [AzksElement]s, which are then inserted into
/// the AZKS.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct AzksElement {
    /// The label of the node
    pub label: NodeLabel,
    /// The associated hash of the node
    #[cfg_attr(
        feature = "serde_serialization",
        serde(serialize_with = "azks_value_hex_serialize")
    )]
    #[cfg_attr(
        feature = "serde_serialization",
        serde(deserialize_with = "azks_value_hex_deserialize")
    )]
    pub value: AzksValue,
}

impl SizeOf for AzksElement {
    fn size_of(&self) -> usize {
        self.label.size_of() + self.value.0.len()
    }
}

impl PartialOrd for AzksElement {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for AzksElement {
    fn cmp(&self, other: &Self) -> Ordering {
        self.label.cmp(&other.label)
    }
}

/// Represents a specific level of the tree with the parental sibling and the direction
/// of the parent for use in tree hash calculations
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct SiblingProof {
    /// The parent's label
    pub label: NodeLabel,
    /// Sibling of the parent that is not on the path
    pub siblings: [AzksElement; 1],
    /// The direction
    pub direction: Direction,
}

/// Merkle proof of membership of a [`NodeLabel`] with a particular hash
/// value in the tree at a given epoch
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct MembershipProof {
    /// The node label
    pub label: NodeLabel,
    /// The hash of the value
    pub hash_val: AzksValue,
    /// The parents of the node in question
    pub sibling_proofs: Vec<SiblingProof>,
}

/// Merkle Patricia proof of non-membership for a [`NodeLabel`] in the tree
/// at a given epoch.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct NonMembershipProof {
    /// The label in question
    pub label: NodeLabel,
    /// The longest prefix in the tree
    pub longest_prefix: NodeLabel,
    /// The children of the longest prefix
    pub longest_prefix_children: [AzksElement; ARITY],
    /// The membership proof of the longest prefix
    pub longest_prefix_membership_proof: MembershipProof,
}

/// Proof that a given label was at a particular state at the given epoch.
/// This means we need to show that the state and version we are claiming for this node must have been:
/// * committed in the tree,
/// * not too far ahead of the most recent marker version,
/// * not stale when served.
/// This proof is sent in response to a lookup query for a particular key.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct LookupProof {
    /// The epoch of this record
    pub epoch: u64,
    /// The plaintext value in question
    pub value: AkdValue,
    /// The version of the record
    pub version: u64,
    /// VRF proof for the label corresponding to this version
    pub existence_vrf_proof: Vec<u8>,
    /// Record existence proof
    pub existence_proof: MembershipProof,
    /// VRF proof for the marker preceding (less than or equal to) this version
    pub marker_vrf_proof: Vec<u8>,
    /// Existence at specific marker
    pub marker_proof: MembershipProof,
    /// VRF proof for the label corresponding to this version being stale
    pub freshness_vrf_proof: Vec<u8>,
    /// Freshness proof (non member at previous epoch)
    pub freshness_proof: NonMembershipProof,
    /// Proof for commitment value derived from raw AkdLabel and AkdValue
    pub commitment_nonce: Vec<u8>,
}

/// A vector of UpdateProofs are sent as the proof to a history query for a particular key.
/// For each version of the value associated with the key, the verifier must check that:
/// * the version was included in the claimed epoch,
/// * the previous version was retired at this epoch,
/// * the version did not exist prior to this epoch,
/// * the next few versions (up until the next marker), did not exist at this epoch,
/// * the future marker versions did  not exist at this epoch.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct UpdateProof {
    /// Epoch of this update
    pub epoch: u64,
    /// Value at this update
    pub value: AkdValue,
    /// Version at this update
    pub version: u64,
    /// VRF proof for the label for the current version
    pub existence_vrf_proof: Vec<u8>,
    /// Membership proof to show that the key was included in this epoch
    pub existence_proof: MembershipProof,
    /// VRF proof for the label for the previous version which became stale
    pub previous_version_vrf_proof: Option<Vec<u8>>,
    /// Proof that previous value was set to old at this epoch
    pub previous_version_proof: Option<MembershipProof>,
    /// Nonce for commitment value derived from raw AkdLabel and AkdValue
    pub commitment_nonce: Vec<u8>,
}

/// This proof is just an array of [`UpdateProof`]s.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct HistoryProof {
    /// The update proofs in the key history
    pub update_proofs: Vec<UpdateProof>,
    /// VRF Proofs for the labels of the values until the next marker version
    pub until_marker_vrf_proofs: Vec<Vec<u8>>,
    /// Proof that the values until the next marker version did not exist at this time
    pub non_existence_until_marker_proofs: Vec<NonMembershipProof>,
    /// VRF proofs for the labels of future marker entries
    pub future_marker_vrf_proofs: Vec<Vec<u8>>,
    /// Proof that future markers did not exist
    pub non_existence_of_future_marker_proofs: Vec<NonMembershipProof>,
}

/// The payload that is outputted as a result of successful verification of
/// a [LookupProof] or [HistoryProof]. This includes the fields containing the
/// epoch that the leaf was published in, the version corresponding to the value,
/// and the value itself.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
pub struct VerifyResult {
    /// The epoch of this record
    pub epoch: u64,
    /// Version at this update
    pub version: u64,
    /// The plaintext value associated with the record
    pub value: AkdValue,
}

/// Proof that no leaves were deleted from the initial epoch.
/// This means that unchanged_nodes should hash to the initial root hash
/// and the vec of inserted is the set of leaves inserted between these epochs.
/// If we built the tree using the nodes in inserted and the nodes in unchanged_nodes
/// as the leaves with the correct epoch of insertion,
/// it should result in the final root hash.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
pub struct SingleAppendOnlyProof {
    /// The inserted nodes & digests
    pub inserted: Vec<AzksElement>,
    /// The unchanged nodes & digests
    pub unchanged_nodes: Vec<AzksElement>,
}

/// Proof that no leaves were deleted from the initial epoch.
/// This is done using a list of SingleAppendOnly proofs, one proof
/// for each epoch between the initial epoch and final epochs which are
/// being audited.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Deserialize, serde::Serialize)
)]
pub struct AppendOnlyProof {
    /// Proof for a single epoch being append-only
    pub proofs: Vec<SingleAppendOnlyProof>,
    /// Epochs over which this audit is being performed
    pub epochs: Vec<u64>,
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains the specifics for NodeLabel only, other types don't have the
//! same level of detail and aren't broken into sub-modules

use crate::{PrefixOrdering, SizeOf};

#[cfg(feature = "serde_serialization")]
use crate::utils::serde_helpers::{bytes_deserialize_hex, bytes_serialize_hex};
#[cfg(feature = "nostd")]
use alloc::vec::Vec;

#[cfg(test)]
mod tests;

/// The label used for an empty node
pub const EMPTY_LABEL: NodeLabel = NodeLabel {
    label_val: [1u8; 32], // FIXME(#344): Let's use a more concise value here, like [1u8, 0u8, ..., 0u8]
    label_len: 0,
};

/// Represents the label of a AKD node
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[cfg_attr(
    feature = "serde_serialization",
    derive(serde::Serialize, serde::Deserialize)
)]
pub struct NodeLabel {
    #[cfg_attr(
        feature = "serde_serialization",
        serde(serialize_with = "bytes_serialize_hex")
    )]
    #[cfg_attr(
        feature = "serde_serialization",
        serde(deserialize_with = "bytes_deserialize_hex")
    )]
    /// Stores a binary string as a 32-byte array of `u8`s
    pub label_val: [u8; 32],
    /// len keeps track of how long the binary string is in bits
    pub label_len: u32,
}

impl SizeOf for NodeLabel {
    fn size_of(&self) -> usize {
        self.label_val.len() + core::mem::size_of::<u32>()
    }
}

impl PartialOrd for NodeLabel {
    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for NodeLabel {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        // `label_len`, `label_val`
        let len_cmp = self.label_len.cmp(&other.label_len);
        if let core::cmp::Ordering::Equal = len_cmp {
            self.label_val.cmp(&other.label_val)
        } else {
            len_cmp
        }
    }
}

impl core::fmt::Display for NodeLabel {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "(0x{}, {})", hex::encode(self.label_val), self.label_len)
    }
}

#[derive(Debug, PartialEq, Eq)]
#[repr(u8)]
pub(crate) enum Bit {
    Zero = 0u8,
    One = 1u8,
}

impl NodeLabel {
    /// Returns the value of the [NodeLabel]
    pub fn value(&self) -> Vec<u8> {
        // FIXME(#344): We shouldn't need to actually hash the label, it is redundant. Change this to:
        // self.to_bytes()
        crate::hash::hash(&self.to_bytes()).to_vec()
    }

    pub(crate) fn to_bytes(self) -> Vec<u8> {
        [&self.label_len.to_be_bytes(), &self.label_val[..]].concat()
    }

    /// Outputs whether or not self is a prefix of the other [NodeLabel]
    pub fn is_prefix_of(&self, other: &Self) -> bool {
        if self.label_len > other.label_len {
            return false;
        }
        (0..self.label_len).all(|i| self.get_bit_at(i) == other.get_bit_at(i))
    }

    /// Takes as input a pointer to the caller and another [NodeLabel],
    /// returns a NodeLabel that is the longest common prefix of the two.
    pub fn get_longest_common_prefix(&self, other: NodeLabel) -> Self {
        let shorter_len = if self.label_len < other.label_len {
            self.label_len
        } else {
            other.label_len
        };

        let mut prefix_len = 0;
        while prefix_len < shorter_len
            && self.get_bit_at(prefix_len) == other.get_bit_at(prefix_len)
        {
            prefix_len += 1;
        }
        if *self == EMPTY_LABEL || other == EMPTY_LABEL {
            return EMPTY_LABEL;
        }
        self.get_prefix(prefix_len)
    }

    /// Returns the bit at a specified index, and a 0 on an out of range index
    ///
    /// Note that this is calculated from the right, for example:
    /// let mut label = [0u8; 32];
    /// label[0] = 0b10100000u8;
    /// We should get outputs as follows:
    /// * label.get_bit_at(0) = 1
    /// * label.get_bit_at(1) = 0
    /// * label.get_bit_at(2) = 1
    /// * label.get_bit_at(3) = 0
    /// * label.get_bit_at(4) = 0
    /// * label.get_bit_at(5) = 0
    /// * label.get_bit_at(6) = 0
    /// * label.get_bit_at(7) = 0
    fn get_bit_at(&self, index: u32) -> Bit {
        if index >= self.label_len {
            return Bit::Zero;
        }

        let usize_index: usize = index as usize;
        let index_full_blocks = usize_index / 8;
        let index_remainder = usize_index % 8;
        if (self.label_val[index_full_blocks] >> (7 - index_remainder)) & 1 == 0 {
            Bit::Zero
        } else {
            Bit::One
        }
    }

    /// Returns the prefix of a specified length, and the entire value on an out of range length
    pub fn get_prefix(&self, len: u32) -> Self {
        if len >= self.label_len {
            return *self;
        }
        if len == 0 {
            return Self {
                label_val: [0u8; 32],
                label_len: 0,
            };
        }

        let usize_len: usize = (len - 1) as usize;
        let len_remainder = usize_len % 8;
        let len_div = usize_len / 8;

        let mut out_val = [0u8; 32];
        out_val[..len_div].clone_from_slice(&self.label_val[..len_div]);
        out_val[len_div] = (self.label_val[len_div] >> (7 - len_remainder)) << (7 - len_remainder);

        Self {
            label_val: out_val,
            label_len: len,
        }
    }

    /// Creates a new NodeLabel representing the root.
    pub fn root() -> Self {
        Self::new([0u8; 32], 0)
    }

    /// Creates a new [NodeLabel] with the given value and len (in bits).
    pub fn new(val: [u8; 32], len: u32) -> Self {
        NodeLabel {
            label_val: val,
            label_len: len,
        }
    }

    /// Gets the length of a NodeLabel in bits.
    pub fn get_len(&self) -> u32 {
        self.label_len
    }

    /// Gets the value of a NodeLabel.
    pub fn get_val(&self) -> [u8; 32] {
        self.label_val
    }

    /// Gets the prefix ordering of other with respect to self, if self is a prefix of other.
    /// If self is not a prefix of other, then this returns [PrefixOrdering::Invalid].
    pub fn get_prefix_ordering(&self, other: Self) -> PrefixOrdering {
        if self.get_len() >= other.get_len() {
            return PrefixOrdering::Invalid;
        }
        if other.get_prefix(self.get_len()) != *self {
            return PrefixOrdering::Invalid;
        }
        PrefixOrdering::from(other.get_bit_at(self.get_len()))
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Tests for node labels

use super::*;
#[cfg(feature = "nostd")]
use alloc::vec;
use rand::{thread_rng, Rng};

// ================= Test helpers ================= //

fn random_label() -> crate::NodeLabel {
    let mut rng = thread_rng();
    crate::NodeLabel {
        label_val: rng.gen::<[u8; 32]>(),
        label_len: 256,
    }
}

// Creates a byte array of 32 bytes from a u64
// Note that this representation is big-endian, and
// places the bits to the front of the output byte_array.
fn byte_arr_from_u64(input_int: u64) -> [u8; 32] {
    let mut output_arr = [0u8; 32];
    let input_arr = input_int.to_be_bytes();
    output_arr[..8].clone_from_slice(&input_arr[..8]);
    output_arr
}

// Creates a byte array of 32 bytes from a u64
// Note that this representation is little-endian, and
// places the bits to the front of the output byte_array.
fn byte_arr_from_u64_le(input_int: u64) -> [u8; 32] {
    let mut output_arr = [0u8; 32];
    let input_arr = input_int.to_be_bytes();
    output_arr[..8].clone_from_slice(&input_arr[..8]);
    output_arr
}

/// This test tests get_bit_at on a small label of len 4.
/// The label is logically equal to the binary string "1010"
/// and should return the corresponding bits.
#[test]
pub fn test_get_bit_at_small() {
    let val = 0b1010u64 << 60;
    let expected = vec![Bit::One, Bit::Zero, Bit::One, Bit::Zero];
    let label = NodeLabel::new(byte_arr_from_u64(val), 4);
    for (index, item) in expected.iter().enumerate().take(4) {
        assert!(
            *item == label.get_bit_at(index as u32),
            "get_bit_at({}) wrong for the 4 digit label 0b1010! Expected {:?} and got {:?}",
            index,
            *item,
            label.get_bit_at(index as u32)
        )
    }
    for index in 4u32..256u32 {
        assert_eq!(
            label.get_bit_at(index),
            Bit::Zero,
            "Index {index} should be 0 in a label of length 4 but it doesn't!"
        );
    }
}

/// In this test, we have a label of length 256, logically equal to
/// 1 followed by 255 0s. We want to make sure its 0th bit is read out as 1.
#[test]
pub fn test_get_bit_at_medium_1() {
    let val = 0b1u64 << 63;
    let expected = Bit::One;
    let label = NodeLabel::new(byte_arr_from_u64(val), 256);
    let computed = label.get_bit_at(0);
    assert!(
        expected == computed,
        "{}",
        "get_bit_at(2) wrong for the 4 digit label 10! Expected {expected:?} and got {computed:?}"
    )
}

// In this test, we have a label of length 256, logically equal to
// 1 followed by 255 0s. We want to make sure its 190th bit is read out as 0.
// We have this because the string itself has only one non-zero bit and we still want
// to check beyond the 0th index.
#[test]
pub fn test_get_bit_at_medium_2() {
    let val = 0b1u64 << 63;
    let expected = Bit::Zero;
    let label = NodeLabel::new(byte_arr_from_u64(val), 256);
    let computed = label.get_bit_at(190);
    assert!(
        expected == computed,
        "{}",
        "get_bit_at(2) wrong for the 4 digit label 10! Expected {expected:?} and got {computed:?}"
    )
}

/// This test creates a label of length 256 logically equal to
/// "0000 0000 0000 0000 1010 0000" followed by all 0s. We know that the
/// first non-zero bit is at position 16, and we want to check that.
#[test]
pub fn test_get_bit_at_large() {
    let mut val = [0u8; 32];
    // 128u8 = 0b1000 0000u8 and 32u8 = 0b10 0000u8, hence their
    // sum is "1010 0000"
    val[2] = 128u8 + 32u8;
    // create the label
    let label = NodeLabel::new(val, 256);
    // val[2] is positions 16-23 (both included),
    // so we want to check everything till there.
    let expected_raw = vec![
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0,
    ];
    let expected = expected_raw
        .iter()
        .map(|x| if *x == 0 { Bit::Zero } else { Bit::One })
        .collect::<Vec<Bit>>();

    // the vector expected covers the first 24 indices.
    for (index, item) in expected.iter().enumerate().take(24) {
        let index_32 = index as u32;
        assert!(
            *item == label.get_bit_at(index_32),
            "get_bit_at({}) wrong for the 256 digit label 0000 0000 0000 0000 1010 0000! Expected {:?} and got {:?}",
            index,
            *item,
            label.get_bit_at(index_32)
        )
    }
    // Everything after the first 24 indixes is 0
    for index in 24..256 {
        let index_32 = index as u32;
        assert!(
            Bit::Zero == label.get_bit_at(index_32),
            "get_bit_at({}) wrong for the 256 digit label 0000 0000 0000 0000 1010 0000! Expected {:?} and got {:?}",
            index,
            Bit::Zero,
            label.get_bit_at(index_32)
        )
    }
}

/// This test is testing our helper function byte_arr_from_u64, which
/// we mainly use for testing. Still we want it to be correct!
/// We call it "small" since it only tests what would
/// result in 1 non-zero byte.
#[test]
pub fn test_byte_arr_from_u64_small() {
    // This val is 2 copies of "10" followed by all 0s.
    // This should be converted into the byte array of all 0s
    // but with the first two byte 0b10100000u8.
    let val = 0b1010u64 << 60;
    let mut expected = [0u8; 32];
    expected[0] = 0b10100000u8;
    let computed = byte_arr_from_u64(val);
    assert!(
        expected == computed,
        "{}",
        "Byte from u64 conversion wrong for small u64! Expected {expected:?} and got {computed:?}"
    )
}

/// This test is testing our helper function byte_arr_from_u64, which
/// we mainly use for testing. Still we want it to be correct!
/// It is only testing for 2 non-zero bytes.
#[test]
pub fn test_byte_arr_from_u64_medium() {
    // This val is 6 copies of "10" followed by all 0s.
    // This should be converted into the byte array of all 0s
    // but with the first two bytes 0b10101010u8 and 0b10100000u8.
    let val = 0b101010101010u64 << 52;
    let mut expected = [0u8; 32];
    expected[0] = 0b10101010u8;
    expected[1] = 0b10100000u8;
    let computed = byte_arr_from_u64(val);
    assert!(
        expected == computed,
        "{}", "Byte from u64 conversion wrong for medium, ~2 byte u64! Expected {expected:?} and got {computed:?}"
    )
}

/// This test is testing our helper function byte_arr_from_u64, which
/// we mainly use for testing. Still we want it to be correct!
/// It is only testing for 3 non-zero bytes.
#[test]
pub fn test_byte_arr_from_u64_larger() {
    // This string was hand-generated for testing so that
    // all three non-zero bytes were distinct.
    let val = 0b01011010101101010101010u64 << 41;
    let mut expected = [0u8; 32];
    expected[0] = 0b01011010u8;
    expected[1] = 0b10110101u8;
    expected[2] = 0b01010100u8;

    let computed = byte_arr_from_u64(val);
    assert!(
        expected == computed,
        "{}", "Byte from u64 conversion wrong for larger, ~3 byte u64! Expected {expected:?} and got {computed:?}"
    )
}

/// Test two NodeLabels for equality, when their leading bit is 1.
#[test]
pub fn test_node_label_equal_leading_one() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    assert!(
        label_1 == label_2,
        "Identical labels with leading one not found equal!"
    )
}

/// Test two NodeLabels for equality, when their leading bit is 0.
#[test]
pub fn test_node_label_equal_leading_zero() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(100000000u64 << 55), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 9u32);
    assert!(
        label_1 == label_2,
        "Identical labels with leading zero not found equal!"
    )
}

/// Test two NodeLabels for inequality, when their leading bit is 1.
#[test]
pub fn test_node_label_unequal_values() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(10000000u64), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(110000000u64), 9u32);
    assert!(label_1 != label_2, "Unequal labels found equal!")
}

/// Test two NodeLabels for inequality due to differing length, when their leading bit is 1.
#[test]
pub fn test_node_label_equal_values_unequal_len() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 9u32);
    assert!(
        label_1 != label_2,
        "Identical labels with unequal lengths not found equal!"
    )
}

/// Test for get_longest_common_prefix between a label and itself being itself. Leading 1.
#[test]
pub fn test_node_label_lcp_with_self_leading_one() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    let expected = NodeLabel::new(byte_arr_from_u64(10000000u64 << 56), 8u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with self with leading one, not equal to itself!"
    )
}

/// Test for get_longest_common_prefix between a label and a zero-length node label.
#[test]
pub fn test_node_label_lcp_with_zero_length_label() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0u64), 0u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0u64), 2u32);
    let expected = label_1;
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with zero-length label, not equal to zero-length label!"
    );
    assert!(
        label_2.get_longest_common_prefix(label_1) == expected,
        "Longest common substring with zero-length label, not equal to zero-length label!"
    );
}

/// Test for get_longest_common_prefix between a label and its prefix.
#[test]
pub fn test_node_label_lcp_with_prefix_label() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(01u64 << 62), 2u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(01u64 << 62), 3u32);
    let expected = label_1;
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with prefix label, not equal to prefix label!"
    );
    assert!(
        label_2.get_longest_common_prefix(label_1) == expected,
        "Longest common substring with prefix label, not equal to prefix label!"
    );
}

/// Test for get_longest_common_prefix between a label and itself being itself. Leading 0.
#[test]
pub fn test_node_label_lcp_with_self_leading_zero() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b1000000u64 << 56), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b1000000u64 << 56), 9u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1000000u64 << 56), 9u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with self with leading zero, not equal to itself!"
    )
}

/// Test for get_longest_common_prefix between a label and a prefix of this label. Leading 1.
#[test]
pub fn test_node_label_lcp_self_prefix_leading_one() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b1000u64 << 60), 4u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 56), 8u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1000u64 << 60), 4u32);
    let computed = label_1.get_longest_common_prefix(label_2);
    assert!(
        computed == expected,
        "{}", "Longest common substring with self with leading one, not equal to itself! Expected: {expected:?}, Got: {computed:?}"
    )
}

/// Test for get_longest_common_prefix between a label and a prefix of this label. Leading 0.
#[test]
pub fn test_node_label_lcp_self_prefix_leading_zero() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 55), 7u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 55), 9u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 55), 7u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with self with leading zero, not equal to itself!"
    )
}

/// Test for get_longest_common_prefix between two labels starting at the bit 1.
#[test]
pub fn test_node_label_lcp_other_one() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 56), 8u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b11000000u64 << 56), 8u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1u64 << 63), 1u32);
    let computed = label_1.get_longest_common_prefix(label_2);
    assert!(
        computed == expected,
        "{}", "Longest common substring with other with leading one, not equal to expected! Expected: {expected:?}, Computed: {computed:?}"
    )
}

/// Test for get_longest_common_prefix between two labels starting at the bits 01.
#[test]
pub fn test_node_label_lcp_other_zero() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 55), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b11000000u64 << 55), 9u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1u64 << 62), 2u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with other with leading zero, not equal to expected!"
    )
}

/// Test for get_longest_common_prefix between two labels which have no common prefix.
#[test]
pub fn test_node_label_lcp_empty() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b10000000u64 << 55), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b11000000u64 << 56), 8u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b0u64), 0u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring should be empty!"
    )
}

/// Test for get_longest_common_prefix between two labels starting at the bits 1101.
#[test]
pub fn test_node_label_lcp_some_leading_one() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b11010000u64 << 56), 8u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b11011000u64 << 56), 8u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1101u64 << 60), 4u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with other with leading one, not equal to expected!"
    )
}

/// Test for get_longest_common_prefix between two labels starting at the bits 01101.
#[test]
pub fn test_node_label_lcp_some_leading_zero() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b11010000u64 << 55), 9u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b11011000u64 << 55), 9u32);
    let expected = NodeLabel::new(byte_arr_from_u64(0b1101u64 << 59), 5u32);
    assert!(
        label_1.get_longest_common_prefix(label_2) == expected,
        "Longest common substring with other with leading zero, not equal to expected!"
    )
}

/// This test tests get_dir by manually computing the prefix and the bit
/// immediately following the prefix of that length.
#[test]
pub fn test_get_dir_large() {
    for i in 0..257 {
        let label_1 = random_label();
        let pos = i;
        // if prefix is of len 256, this will get the entire random string
        let label_2 = label_1.get_prefix(pos);
        // if the prefix is of length pos, then we want to get the prefix in that position, since the
        // label's value is indexed on 0, so the bit following the prefix of len "pos" is at position pos.
        let mut expected = PrefixOrdering::from(label_1.get_bit_at(pos));
        if pos == 256 {
            expected = PrefixOrdering::Invalid;
        }
        let computed = label_2.get_prefix_ordering(label_1);
        assert!(
            computed == expected,
            "{}",
            "Direction not equal to expected. Node = {label_1:?}, prefix = {label_2:?}"
        )
    }
}

/// This test just serves as another example of get_dir and this time we want to use little endian encoding
/// since we are using more complex u64 values.
#[test]
pub fn test_get_dir_example() {
    // 23 in little endian is 10111 and 10049430782486799941u64 begins with
    // the prefix 00110100, hence, label_1 is not a prefix of label_2.
    let label_1 = NodeLabel::new(byte_arr_from_u64_le(10049430782486799941u64), 64u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64_le(23u64), 5u32);
    let expected = PrefixOrdering::Invalid;
    let computed = label_2.get_prefix_ordering(label_1);
    assert!(
        computed == expected,
        "{}", "Direction not equal to expected. Node = {label_1:?}, prefix = {label_2:?}, computed = {computed:?}"
    )
}

/// This test gets a prefix for a hard-coded random string and makes sure it is equal to a hand-computed value.
#[test]
pub fn test_get_prefix_small() {
    let label_1 = NodeLabel::new(
        byte_arr_from_u64(0b1000101101110110110000000000110101110001000000000110011001000101u64),
        64u32,
    );
    let prefix_len = 10u32;
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b1000101101u64 << 54), prefix_len);
    let computed = label_1.get_prefix(prefix_len);
    assert!(
        computed == label_2,
        "{}", "Direction not equal to expected. Node = {label_1:?}, prefix = {label_2:?}, computed = {computed:?}"
    )
}

#[test]
pub fn test_is_prefix_of() {
    let label_1 = NodeLabel::new(byte_arr_from_u64(0b01u64 << 62), 4u32);
    let label_2 = NodeLabel::new(byte_arr_from_u64(0b010u64 << 61), 5u32);
    let label_3 = NodeLabel::new(byte_arr_from_u64(0b0u64), 4u32);

    // empty label is prefix of all labels
    assert_eq!(EMPTY_LABEL.is_prefix_of(&label_1), true);
    assert_eq!(EMPTY_LABEL.is_prefix_of(&label_2), true);
    assert_eq!(EMPTY_LABEL.is_prefix_of(&label_3), true);

    // every label is a prefix of itself
    assert_eq!(label_1.is_prefix_of(&label_1), true);
    assert_eq!(label_2.is_prefix_of(&label_2), true);
    assert_eq!(label_3.is_prefix_of(&label_3), true);

    // valid prefixes
    assert_eq!(label_1.is_prefix_of(&label_2), true);

    // invalid prefixes
    assert_eq!(label_1.is_prefix_of(&label_3), false);
    assert_eq!(label_2.is_prefix_of(&label_1), false);
    assert_eq!(label_2.is_prefix_of(&label_3), false);
    assert_eq!(label_3.is_prefix_of(&label_1), false);
    assert_eq!(label_3.is_prefix_of(&label_2), false);
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Base functionality for verification operations (membership, non-membership, etc)

use super::VerificationError;

use crate::crypto::{
    compute_parent_hash_from_children, compute_root_hash_from_val, get_hash_from_label_input,
    hash_leaf_with_commitment, hash_leaf_with_value,
};
use crate::ecvrf::{Proof, VrfError};
use crate::hash::Digest;
use crate::{
    AkdLabel, AkdValue, AzksValue, Direction, MembershipProof, NodeLabel, NonMembershipProof,
    VersionFreshness, EMPTY_LABEL,
};

#[cfg(feature = "nostd")]
use alloc::format;
#[cfg(feature = "nostd")]
use alloc::string::ToString;
use core::convert::TryFrom;

/// Verify the membership proof
pub fn verify_membership(
    root_hash: Digest,
    proof: &MembershipProof,
) -> Result<(), VerificationError> {
    let mut curr_val = proof.hash_val;
    let mut curr_label = proof.label;

    for sibling_proof in proof.sibling_proofs.iter().rev() {
        let sibling = sibling_proof.siblings[0];
        let (left_val, left_label, right_val, right_label) = match sibling_proof.direction {
            Direction::Left => (
                curr_val,
                curr_label.value(),
                sibling.value,
                sibling.label.value(),
            ),
            Direction::Right => (
                sibling.value,
                sibling.label.value(),
                curr_val,
                curr_label.value(),
            ),
        };
        curr_val =
            compute_parent_hash_from_children(&left_val, &left_label, &right_val, &right_label);
        curr_label = sibling_proof.label;
    }

    if compute_root_hash_from_val(&curr_val) == root_hash {
        Ok(())
    } else {
        Err(VerificationError::MembershipProof(format!(
            "Membership proof for label {:?} did not verify",
            proof.label
        )))
    }
}

/// Verifies the non-membership proof with respect to the root hash
pub fn verify_nonmembership(
    root_hash: Digest,
    proof: &NonMembershipProof,
) -> Result<(), VerificationError> {
    // Verify that the proof's label is not equal to either of the children's labels
    if proof.label == proof.longest_prefix_children[0].label
        || proof.label == proof.longest_prefix_children[1].label
    {
        return Err(VerificationError::NonMembershipProof(
            "Proof's label is equal to one of the children's labels".to_string(),
        ));
    }

    // Verify that the proof's label is a prefix of proof.longest_prefix
    if !proof.longest_prefix.is_prefix_of(&proof.label) {
        return Err(VerificationError::NonMembershipProof(
            "Proof's label is not a prefix of longest_prefix".to_string(),
        ));
    }

    // Verify that proof.longest_prefix is the longest common prefix of the children
    let mut lcp_children = proof.longest_prefix_children[0]
        .label
        .get_longest_common_prefix(proof.longest_prefix_children[1].label);
    if lcp_children == EMPTY_LABEL {
        // This is a special case that only occurs when the lcp is the root node and
        // it is missing one of its children
        lcp_children = NodeLabel::root();
    }
    if proof.longest_prefix != lcp_children {
        return Err(VerificationError::NonMembershipProof(
            "longest_prefix != computed lcp".to_string(),
        ));
    }

    let lcp_hash = compute_parent_hash_from_children(
        &proof.longest_prefix_children[0].value,
        &proof.longest_prefix_children[0].label.value(),
        &proof.longest_prefix_children[1].value,
        &proof.longest_prefix_children[1].label.value(),
    );
    if lcp_children != proof.longest_prefix_membership_proof.label
        || lcp_hash != proof.longest_prefix_membership_proof.hash_val
    {
        return Err(VerificationError::NonMembershipProof(
            "lcp_hash != longest_prefix_hash".to_string(),
        ));
    }
    verify_membership(root_hash, &proof.longest_prefix_membership_proof)?;

    Ok(())
}

/// This function is called to verify that a given [NodeLabel] is indeed
/// the VRF for a given version (fresh or stale) for a [AkdLabel].
/// Hence, it also takes as input the server's public key.
fn verify_label(
    vrf_public_key: &[u8],
    akd_label: &AkdLabel,
    freshness: VersionFreshness,
    version: u64,
    vrf_proof: &[u8],
    node_label: NodeLabel,
) -> Result<(), VerificationError> {
    let vrf_pk = crate::ecvrf::VRFPublicKey::try_from(vrf_public_key)?;
    let hashed_label = get_hash_from_label_input(akd_label, freshness, version);

    // VRF proof verification (returns VRF hash output)
    let proof = Proof::try_from(vrf_proof)?;
    vrf_pk.verify(&proof, &hashed_label)?;
    let output: crate::ecvrf::Output = (&proof).into();

    if NodeLabel::new(output.to_truncated_bytes(), 256) != node_label {
        return Err(VerificationError::Vrf(VrfError::Verification(
            "Expected first 32 bytes of the proof output did NOT match the supplied label"
                .to_string(),
        )));
    }
    Ok(())
}

pub(crate) fn verify_existence(
    vrf_public_key: &[u8],
    root_hash: Digest,
    akd_label: &AkdLabel,
    freshness: VersionFreshness,
    version: u64,
    vrf_proof: &[u8],
    membership_proof: &MembershipProof,
) -> Result<(), VerificationError> {
    verify_label(
        vrf_public_key,
        akd_label,
        freshness,
        version,
        vrf_proof,
        membership_proof.label,
    )?;
    verify_membership(root_hash, membership_proof)?;
    Ok(())
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn verify_existence_with_val(
    vrf_public_key: &[u8],
    root_hash: Digest,
    akd_label: &AkdLabel,
    akd_value: &AkdValue,
    epoch: u64,
    commitment_nonce: &[u8],
    freshness: VersionFreshness,
    version: u64,
    vrf_proof: &[u8],
    membership_proof: &MembershipProof,
) -> Result<(), VerificationError> {
    if hash_leaf_with_value(akd_value, epoch, commitment_nonce).0 != membership_proof.hash_val.0 {
        return Err(VerificationError::MembershipProof(
            "Hash of plaintext value did not match existence proof hash".to_string(),
        ));
    }
    verify_existence(
        vrf_public_key,
        root_hash,
        akd_label,
        freshness,
        version,
        vrf_proof,
        membership_proof,
    )?;

    Ok(())
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn verify_existence_with_commitment(
    vrf_public_key: &[u8],
    root_hash: Digest,
    akd_label: &AkdLabel,
    commitment: AzksValue,
    epoch: u64,
    freshness: VersionFreshness,
    version: u64,
    vrf_proof: &[u8],
    membership_proof: &MembershipProof,
) -> Result<(), VerificationError> {
    if hash_leaf_with_commitment(commitment, epoch).0 != membership_proof.hash_val.0 {
        return Err(VerificationError::MembershipProof(
            "Hash of plaintext value did not match existence proof hash".to_string(),
        ));
    }
    verify_existence(
        vrf_public_key,
        root_hash,
        akd_label,
        freshness,
        version,
        vrf_proof,
        membership_proof,
    )?;

    Ok(())
}

pub(crate) fn verify_nonexistence(
    vrf_public_key: &[u8],
    root_hash: Digest,
    akd_label: &AkdLabel,
    freshness: VersionFreshness,
    version: u64,
    vrf_proof: &[u8],
    nonmembership_proof: &NonMembershipProof,
) -> Result<(), VerificationError> {
    verify_label(
        vrf_public_key,
        akd_label,
        freshness,
        version,
        vrf_proof,
        nonmembership_proof.label,
    )?;
    verify_nonmembership(root_hash, nonmembership_proof)?;
    Ok(())
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Verification of key history proofs

use super::base::{
    verify_existence, verify_existence_with_commitment, verify_existence_with_val,
    verify_nonexistence,
};
use super::VerificationError;

use crate::crypto::stale_azks_value;
use crate::hash::Digest;
use crate::{AkdLabel, HistoryProof, UpdateProof, VerifyResult, VersionFreshness};
#[cfg(feature = "nostd")]
use alloc::format;
#[cfg(feature = "nostd")]
use alloc::string::ToString;
#[cfg(feature = "nostd")]
use alloc::vec::Vec;

/// Parameters for customizing how history proof verification proceeds
#[derive(Copy, Clone)]
pub enum HistoryVerificationParams {
    /// No customization to the verification procedure
    Default,
    /// Allows for the encountering of missing (tombstoned) values
    /// instead of attempting to check if their hash matches the leaf node
    /// hash
    AllowMissingValues,
}

impl Default for HistoryVerificationParams {
    fn default() -> Self {
        Self::Default
    }
}

/// Verifies a key history proof, given the corresponding sequence of hashes.
/// Returns a vector of whether the validity of a hash could be verified.
/// When false, the value <=> hash validity at the position could not be
/// verified because the value has been removed ("tombstoned") from the storage layer.
pub fn key_history_verify(
    vrf_public_key: &[u8],
    root_hash: Digest,
    current_epoch: u64,
    akd_key: AkdLabel,
    proof: HistoryProof,
    params: HistoryVerificationParams,
) -> Result<Vec<VerifyResult>, VerificationError> {
    let mut results = Vec::new();
    let mut last_version = 0;

    let num_proofs = proof.update_proofs.len();

    // Make sure the update proofs are non-empty
    if num_proofs == 0 {
        return Err(VerificationError::HistoryProof(format!(
            "No update proofs included in the proof of user {akd_key:?} at epoch {current_epoch:?}!"
        )));
    }

    // Check that the sent proofs are for a contiguous sequence of decreasing versions
    for count in 0..num_proofs {
        if count > 0 {
            // Make sure this proof is for a version 1 more than the previous one.
            if proof.update_proofs[count].version + 1 != proof.update_proofs[count - 1].version {
                return Err(VerificationError::HistoryProof(format!("Why did you give me consecutive update proofs without version numbers decrementing by 1? Version {} = {}; version {} = {}",
                count, proof.update_proofs[count].version,
                count-1, proof.update_proofs[count-1].version
                )));
            }
        }
    }

    // Verify all individual update proofs
    let mut maybe_previous_update_epoch = None;
    for update_proof in proof.update_proofs.into_iter() {
        // Get the highest version sent among the update proofs.
        last_version = if update_proof.version > last_version {
            update_proof.version
        } else {
            last_version
        };

        if let Some(previous_update_epoch) = maybe_previous_update_epoch {
            // Make sure this this epoch is more than the previous epoch you checked
            if update_proof.epoch > previous_update_epoch {
                return Err(VerificationError::HistoryProof(format!(
                    "Why are your versions decreasing in updates and epochs not?!,
                    epoch = {}, previous epoch = {}",
                    update_proof.epoch, previous_update_epoch
                )));
            }
        }
        maybe_previous_update_epoch = Some(update_proof.epoch);
        let result =
            verify_single_update_proof(root_hash, vrf_public_key, update_proof, &akd_key, params)?;
        results.push(result);
    }

    // Get the least and greatest marker entries for the current version
    let next_marker = crate::utils::get_marker_version_log2(last_version) + 1;
    let final_marker = crate::utils::get_marker_version_log2(current_epoch);

    // ***** Future checks below ***************************
    // Verify the non-existence of future entries, up to the next marker
    for (i, version) in (last_version + 1..(1 << next_marker)).enumerate() {
        verify_nonexistence(
            vrf_public_key,
            root_hash,
            &akd_key,
            VersionFreshness::Fresh,
            version,
            &proof.until_marker_vrf_proofs[i],
            &proof.non_existence_until_marker_proofs[i],
        ).map_err(|_|
            VerificationError::HistoryProof(format!("Non-existence of next few proof of user {:?}'s version {:?} at epoch {:?} does not verify",
                    &akd_key, version, current_epoch)))?;
    }

    // Verify the VRFs and non-membership proofs for future markers
    for (i, pow) in (next_marker..final_marker + 1).enumerate() {
        let version = 1 << pow;
        verify_nonexistence(
            vrf_public_key,
            root_hash,
            &akd_key,
            VersionFreshness::Fresh,
            version,
            &proof.future_marker_vrf_proofs[i],
            &proof.non_existence_of_future_marker_proofs[i],
        ).map_err(|_|
            VerificationError::HistoryProof(format!("Non-existence of future marker proof of user {akd_key:?}'s version {version:?} at epoch {current_epoch:?} does not verify")))?;
    }

    Ok(results)
}

/// Verifies a single update proof
fn verify_single_update_proof(
    root_hash: Digest,
    vrf_public_key: &[u8],
    proof: UpdateProof,
    akd_label: &AkdLabel,
    params: HistoryVerificationParams,
) -> Result<VerifyResult, VerificationError> {
    // Verify the VRF and membership proof for the corresponding label for the version being updated to.
    match (params, &proof.value) {
        (HistoryVerificationParams::AllowMissingValues, bytes) if bytes.0 == crate::TOMBSTONE => {
            // A tombstone was encountered, we need to just take the
            // hash of the value at "face value" since we don't have
            // the real value available
            verify_existence(
                vrf_public_key,
                root_hash,
                akd_label,
                VersionFreshness::Fresh,
                proof.version,
                &proof.existence_vrf_proof,
                &proof.existence_proof,
            )?;
        }
        (_, akd_value) => {
            // No tombstone so hash the value found, and compare to the existence proof's value
            verify_existence_with_val(
                vrf_public_key,
                root_hash,
                akd_label,
                akd_value,
                proof.epoch,
                &proof.commitment_nonce,
                VersionFreshness::Fresh,
                proof.version,
                &proof.existence_vrf_proof,
                &proof.existence_proof,
            )?;
        }
    };

    let verify_result = VerifyResult {
        epoch: proof.epoch,
        version: proof.version,
        value: proof.value,
    };

    if proof.version <= 1 {
        // There is no previous version, so we can just return here
        return Ok(verify_result);
    }

    // ***** PART 2 ***************************
    // Verify the membership proof the for stale label of the previous version

    let previous_version_proof = proof.previous_version_proof.as_ref().ok_or_else(|| {
        VerificationError::HistoryProof("Missing membership proof for previous version".to_string())
    })?;
    let previous_version_vrf_proof =
        proof.previous_version_vrf_proof.as_ref().ok_or_else(|| {
            VerificationError::HistoryProof("Missing VRF proof for previous version".to_string())
        })?;

    verify_existence_with_commitment(
        vrf_public_key,
        root_hash,
        akd_label,
        stale_azks_value(),
        proof.epoch,
        VersionFreshness::Stale,
        proof.version - 1,
        previous_version_vrf_proof,
        previous_version_proof,
    )?;

    Ok(verify_result)
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! Verification of lookup proofs

use super::base::{verify_existence, verify_existence_with_val, verify_nonexistence};
use super::VerificationError;

use crate::hash::Digest;
use crate::{AkdLabel, LookupProof, VerifyResult, VersionFreshness};

/// Verifies a lookup with respect to the root_hash
pub fn lookup_verify(
    vrf_public_key: &[u8],
    root_hash: Digest,
    akd_label: AkdLabel,
    proof: LookupProof,
) -> Result<VerifyResult, VerificationError> {
    verify_existence_with_val(
        vrf_public_key,
        root_hash,
        &akd_label,
        &proof.value,
        proof.epoch,
        &proof.commitment_nonce,
        VersionFreshness::Fresh,
        proof.version,
        &proof.existence_vrf_proof,
        &proof.existence_proof,
    )?;

    let marker_version = 1 << crate::utils::get_marker_version_log2(proof.version);
    verify_existence(
        vrf_public_key,
        root_hash,
        &akd_label,
        VersionFreshness::Fresh,
        marker_version,
        &proof.marker_vrf_proof,
        &proof.marker_proof,
    )?;

    verify_nonexistence(
        vrf_public_key,
        root_hash,
        &akd_label,
        VersionFreshness::Stale,
        proof.version,
        &proof.freshness_vrf_proof,
        &proof.freshness_proof,
    )?;

    Ok(VerifyResult {
        epoch: proof.epoch,
        version: proof.version,
        value: proof.value,
    })
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module contains verification calls for different proofs contained in the AKD crate

pub mod base;
pub mod history;
pub mod lookup;

#[cfg(feature = "nostd")]
use alloc::format;
#[cfg(feature = "nostd")]
use alloc::string::String;
#[cfg(feature = "nostd")]
use alloc::string::ToString;

/// Proof verification error types
#[derive(Debug, Eq, PartialEq)]
pub enum VerificationError {
    /// Error verifying a membership proof
    MembershipProof(String),
    /// Error verifying a non-membership proof
    NonMembershipProof(String),
    /// Error verifying a lookup proof
    LookupProof(String),
    /// Error verifying a history proof
    HistoryProof(String),
    /// Error verifying a VRF proof
    #[cfg(feature = "vrf")]
    Vrf(crate::ecvrf::VrfError),
    /// Error converting protobuf types during verification
    #[cfg(feature = "protobuf")]
    Serialization(crate::proto::ConversionError),
}

impl core::fmt::Display for VerificationError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let code = match &self {
            VerificationError::MembershipProof(err) => format!("(Membership proof) - {err}"),
            VerificationError::NonMembershipProof(err) => {
                format!("(Non-membership proof) - {err}")
            }
            VerificationError::LookupProof(err) => format!("(Lookup proof) - {err}"),
            VerificationError::HistoryProof(err) => format!("(History proof) - {err}"),
            #[cfg(feature = "vrf")]
            VerificationError::Vrf(vrf) => vrf.to_string(),
            #[cfg(feature = "protobuf")]
            VerificationError::Serialization(proto) => proto.to_string(),
        };
        write!(f, "Verification error {code}")
    }
}

#[cfg(feature = "vrf")]
impl From<crate::ecvrf::VrfError> for VerificationError {
    fn from(input: crate::ecvrf::VrfError) -> Self {
        VerificationError::Vrf(input)
    }
}

#[cfg(feature = "protobuf")]
impl From<crate::proto::ConversionError> for VerificationError {
    fn from(input: crate::proto::ConversionError) -> Self {
        VerificationError::Serialization(input)
    }
}

#[cfg(feature = "protobuf")]
impl From<protobuf::Error> for VerificationError {
    fn from(input: protobuf::Error) -> Self {
        let conv: crate::proto::ConversionError = input.into();
        conv.into()
    }
}

// Re-export the necessary verification functions
pub use base::{verify_membership, verify_nonmembership};
pub use history::{key_history_verify, HistoryVerificationParams};
pub use lookup::lookup_verify;


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module implements some common test functionality (i.e. proof generation and whatnot)
//! for use in test infra

use akd::directory::Directory;
use akd::ecvrf::HardCodedAkdVRF;
use akd::storage::memory::AsyncInMemoryDatabase;
use akd::storage::StorageManager;
use akd::AkdLabel;
use akd::AkdValue;
use akd::Digest;

/// Generate a aws-safe bucket/table name from the test function name
///
/// So these are for S3 storage bucket names.
/// To guarantee parallel tests don't conflict with each other,
/// each test uses a different bucket for its tests which is
/// derived from the function name. However buckets cannot
/// contain the "_" char and must be capped @ 63 chars.
/// It's only for generating independent test cases
macro_rules! alphanumeric_function_name {
    () => {{
        fn f() {}
        fn type_name_of<T>(_: T) -> &'static str {
            std::any::type_name::<T>()
        }
        let name = type_name_of(f);
        name[..name.len() - 3]
            .trim_end_matches("::{{closure}}")
            .split("::")
            .last()
            .unwrap()
            .replace("_", "")
            .chars()
            .take(63)
            .collect::<String>()
    }};
}

// Re-export the macro for cross-module use
pub(crate) use alphanumeric_function_name;

pub struct AuditInformation {
    pub proof: akd::AppendOnlyProof,
    pub phash: Digest,
    pub chash: Digest,
}

/// Generate N audit proofs from a new tree
///
/// n: number of proofs to generate
/// expensive: This flag denotes if we should have many or just a single update in each iteration. If true, makes the audit's very large (100K updates)
pub async fn generate_audit_proofs(
    n: usize,
    expensive: bool,
) -> Result<Vec<AuditInformation>, akd::errors::AkdError> {
    let db = AsyncInMemoryDatabase::new();
    let storage_manager = StorageManager::new_no_cache(db);
    let vrf = HardCodedAkdVRF {};
    let akd = Directory::<_, _>::new(storage_manager, vrf).await?;
    let mut proofs = vec![];
    // gather the hash for epoch "0" (init)
    let mut phash = akd.get_epoch_hash().await?.1;

    for _epoch in 1..=n {
        if expensive {
            // generate (n) epochs of updates on many users
            let data = (1..100000)
                .map(|item| {
                    let user = format!("user{item}");
                    let value = format!("value{_epoch}_{item}");
                    (AkdLabel::from(&user), AkdValue::from(&value))
                })
                .collect::<Vec<(AkdLabel, AkdValue)>>();
            akd.publish(data).await?;
        } else {
            // generate (n) epochs of updates on the same user
            let t_value = format!("certificate{_epoch}");
            akd.publish(vec![(AkdLabel::from("user"), AkdValue::from(&t_value))])
                .await?;
        }
        // generate the append-only proof
        let proof = akd.audit((_epoch - 1) as u64, _epoch as u64).await?;
        // retrieve the updated hash
        let chash = akd.get_epoch_hash().await?.1;

        // we have everything we need, pass it along
        let info = AuditInformation {
            phash,
            chash,
            proof,
        };
        proofs.push(info);

        // reset the "previous" hash for the next iteration
        phash = chash;
    }
    Ok(proofs)
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module provides a basic logger implementation to the console with
//! some additional information (file, line number, thread, etc)

use colored::*;
use log::Level;
use log::Metadata;
use log::Record;
use once_cell::sync::OnceCell;
use std::io::Write;
use std::sync::Once;
use tokio::time::Duration;
use tokio::time::Instant;

static EPOCH: OnceCell<Instant> = OnceCell::new();
static INIT_ONCE: Once = Once::new();

pub(crate) static LOGGER: ConsoleLogger = ConsoleLogger {
    level: log::Level::Debug,
};

pub struct ConsoleLogger {
    pub level: Level,
}

impl ConsoleLogger {
    pub(crate) fn format_log_record(io: &mut (dyn Write + Send), record: &Record, colored: bool) {
        let target = {
            if let Some(target_str) = record.target().split(':').last() {
                if let Some(line) = record.line() {
                    format!(" ({target_str}:{line})")
                } else {
                    format!(" ({target_str})")
                }
            } else {
                "".to_string()
            }
        };

        let toc = if let Some(epoch) = EPOCH.get() {
            Instant::now() - *epoch
        } else {
            Duration::from_millis(0)
        };

        let seconds = toc.as_secs();
        let hours = seconds / 3600;
        let minutes = (seconds / 60) % 60;
        let seconds = seconds % 60;
        let milliseconds = toc.subsec_millis();

        let msg = format!(
            "[{:02}:{:02}:{:02}.{:03}] {:6} {}{}",
            hours,
            minutes,
            seconds,
            milliseconds,
            record.level(),
            record.args(),
            target
        );
        if colored {
            let msg = match record.level() {
                Level::Trace | Level::Debug => msg.white(),
                Level::Info => msg.green(),
                Level::Warn => msg.yellow().bold(),
                Level::Error => msg.red().bold(),
            };
            let _ = writeln!(io, "{msg}");
        } else {
            let _ = writeln!(io, "{msg}");
        }
    }
}

impl log::Log for ConsoleLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &Record) {
        if !self.enabled(record.metadata()) {
            return;
        }
        let mut io = std::io::stdout();
        ConsoleLogger::format_log_record(&mut io, record, true);
    }

    fn flush(&self) {
        let _ = std::io::stdout().flush();
    }
}

/// Initialize the logger for console logging within test environments.
/// This is safe to call multiple times, but it will only initialize the logger
/// to the log-level _first_ set. If you want a specific log-level (e.g. Debug)
/// for a specific test, make sure to only run that single test after editing that
/// test's log-level.
///
/// The default level applied everywhere is Info
pub(crate) fn init_logger(level: Level) {
    EPOCH.get_or_init(Instant::now);

    INIT_ONCE.call_once(|| {
        log::set_logger(&LOGGER)
            .map(|()| log::set_max_level(level.to_level_filter()))
            .unwrap();
    });
}

/// Global test startup constructor. Only runs in the TEST profile. Each
/// crate which wants logging enabled in tests being run should make this call
/// itself.
#[cfg(test)]
#[ctor::ctor]
fn test_start() {
    init_logger(Level::Warn);
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! A tool to verify audit proofs from a public (ideally immutable) storage
//! medium. This tool is a read-evaluate-print-loop (REPL) interface, where
//! a user can retrieve the information necessary about all operations by typing
//! `help` into the REPL prompt.
//!
//! # Summary
//!
//! To startup the client, you can choose 1 of the supported storage mediums.
//! Presently the supported storage mediums are
//!
//! 1. AWS DynamoDB index backed by S3 storage
//! 2. AWS S3 only, without an index
//!
//! The application is started with all necessary flags to connect to the storage medium
//! of choice, and then the REPL will start allowing the user to interact as an auditor.
//!
//! # Examples
//!
//! Assuming the audit proofs are stored in an S3 bucket in the AWS region `us-east-2` named
//! "myproofs". To start the application, you can run
//!
//! ```bash
//! cargo run -p akd_local_auditor -- s3 --bucket myproofs --region us-east-2
//! ```
//!
//! ## Connection customization
//!
//! If you need to customize the connection to AWS, both data-layers support providing custom
//! endpoints as well as a access key and secret key for authentication.

pub mod auditor;
mod console_log;
pub mod storage;
pub mod ui;

#[cfg(test)]
pub(crate) mod common_test;

use anyhow::Result;
use clap::{Parser, Subcommand, ValueEnum};
use log::debug;

#[derive(ValueEnum, Clone, Debug)]
enum PublicLogLevel {
    Error,
    Warn,
    Info,
    Debug,
    Trace,
}

impl From<&PublicLogLevel> for log::Level {
    fn from(level: &PublicLogLevel) -> Self {
        match &level {
            PublicLogLevel::Error => log::Level::Error,
            PublicLogLevel::Warn => log::Level::Warn,
            PublicLogLevel::Info => log::Level::Info,
            PublicLogLevel::Debug => log::Level::Debug,
            PublicLogLevel::Trace => log::Level::Trace,
        }
    }
}

/// What mode to open the tool in
#[derive(Subcommand, Clone, Debug)]
pub enum Mode {
    /// Launch the user-interface tool (built with Iced)
    Ui,

    /// Amazon S3 compatible storage
    S3(storage::s3::S3ClapSettings),

    /// DynamoDB
    DynamoDb(storage::dynamodb::DynamoDbClapSettings),
}

/// AKD audit proof verification utility
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
pub struct Arguments {
    /// The logging level to use for console output
    #[clap(long, short, value_enum, ignore_case = true, default_value = "Info")]
    log_level: PublicLogLevel,

    /// Show the verification QR code in the terminal
    #[clap(long)]
    qr: bool,

    /// Storage configuration for audit proofs
    #[clap(subcommand)]
    command: Mode,
}

// MAIN //
#[tokio::main]
async fn main() -> Result<()> {
    let args = Arguments::parse();

    // initialize the logger
    let log_level: log::Level = (&args.log_level).into();
    console_log::init_logger(log_level);

    debug!("Parsed args: {:?}", args);

    match args.command {
        Mode::Ui => crate::ui::UserInterface::execute(None),
        other => {
            let storage: Box<dyn storage::AuditProofStorage> = match &other {
                Mode::S3(s3_settings) => {
                    let imp: storage::s3::S3AuditStorage = s3_settings.into();
                    Box::new(imp)
                }
                Mode::DynamoDb(dynamo_settings) => {
                    let imp: storage::dynamodb::DynamoDbAuditStorage = dynamo_settings.into();
                    Box::new(imp)
                }
                _ => panic!("Unsupported storage mode?"),
            };

            let command_processor = crate::auditor::AuditProcessor::new_repl_processor(storage);

            let mut repl = rustyrepl::Repl::<auditor::AuditArgs>::new(
                command_processor,
                Some(auditor::HISTORY_FILE.to_string()),
                Some("$ ".to_string()),
            )?;
            repl.process().await
        }
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module holds the auditor operations based on binary-encoded AuditProof blobs

use super::storage::{EpochSummary, ProofIndexCacheOption};

use akd::Digest;
use anyhow::{anyhow, bail, Result};
use clap::{Parser, Subcommand};
use log::{error, info, warn};
use rustyrepl::ReplCommandProcessor;
use std::sync::Arc;

pub(crate) const HISTORY_FILE: &str = ".akd_local_auditor_history";

pub(crate) fn format_qr_record(p_hash: Digest, c_hash: Digest, epoch: u64) -> Vec<u8> {
    let qr_msg = format!(
        "AKD_VERIFY:{}/{}/{}",
        epoch,
        hex::encode(c_hash),
        hex::encode(p_hash)
    );
    qr_msg.as_bytes().to_vec()
}

pub async fn audit_epoch(blob: akd::local_auditing::AuditBlob, qr: bool) -> Result<()> {
    // decode the proof
    let (epoch, p_hash, c_hash, proof) = blob.decode().map_err(|err| anyhow!("{:?}", err))?;

    // verify it
    if let Err(akd_error) = akd::auditor::audit_verify(
        vec![p_hash, c_hash],
        akd::AppendOnlyProof {
            proofs: vec![proof],
            epochs: vec![epoch],
        },
    )
    .await
    {
        warn!(
            "Audit proof for epoch {} failed to verify with error {}",
            epoch, akd_error
        );
    } else {
        // verification passed, generate the appropriate QR code
        info!("Audit proof for epoch {} has verified!", epoch);
        if qr {
            info!("Generating scan-able QR code for the verification on device");
            let qr_code_data = format_qr_record(p_hash, c_hash, epoch);
            if let Err(error) = qr2term::print_qr(qr_code_data) {
                error!("Error generating QR code {}", error);
                bail!("Error generating QR code {}", error);
            }
        }
    }
    Ok(())
}

/// Storage options for retrieving audit proofs
#[derive(Subcommand, Clone, Debug)]
pub enum AuditCommand {
    /// Audit a specific epoch
    Audit {
        /// The epoch to audit
        #[clap(long)]
        epoch: u64,
        /// Show a QR code of the audit
        #[clap(long)]
        qr: bool,
        /// Force refreshing the index of proofs (i.e. don't use the cache)
        #[clap(long)]
        force_refresh: bool,
    },
    /// Show the available epochs to audit
    ShowEpochs {
        /// Force refreshing the index of proofs (i.e. don't use the cache)
        #[clap(long)]
        force_refresh: bool,
    },

    /// Show the user interface for auditing
    Ui,
}

/// Audit operations supported by the client
#[derive(Parser, Clone, Debug)]
pub struct AuditArgs {
    #[clap(subcommand)]
    command: AuditCommand,
}

/// Audit processing crate
#[derive(Debug)]
pub struct AuditProcessor {
    pub storage: Arc<Box<dyn crate::storage::AuditProofStorage>>,
    pub last_summaries: Arc<tokio::sync::RwLock<Option<Vec<EpochSummary>>>>,
}

impl AuditProcessor {
    pub fn new_repl_processor(
        storage: Box<dyn crate::storage::AuditProofStorage>,
    ) -> Box<dyn rustyrepl::ReplCommandProcessor<AuditArgs>> {
        Box::new(crate::auditor::AuditProcessor {
            storage: Arc::new(storage),
            last_summaries: Arc::new(tokio::sync::RwLock::new(None)),
        })
    }
}

#[async_trait::async_trait]
impl ReplCommandProcessor<AuditArgs> for AuditProcessor {
    fn is_quit(&self, command: &str) -> bool {
        matches!(command, "quit" | "exit")
    }

    async fn process_command(&self, cmd: AuditArgs) -> Result<()> {
        let default_cache_control = self.storage.default_cache_control();

        match &cmd.command {
            AuditCommand::Audit {
                epoch,
                qr,
                force_refresh,
            } => {
                let cache_control = if *force_refresh {
                    ProofIndexCacheOption::NoCache
                } else {
                    default_cache_control
                };

                // did they already call show-epochs below? If so, can we utilize the cached result to save
                // another scan operation?
                let proofs = if cache_control == ProofIndexCacheOption::UseCache {
                    let summaries = &*(self.last_summaries.read().await);
                    if let Some(cached_result) = summaries {
                        cached_result.clone()
                    } else {
                        self.storage.list_proofs(cache_control).await?
                    }
                } else {
                    self.storage.list_proofs(cache_control).await?
                };

                let maybe_proof = proofs.iter().find(|proof| proof.name.epoch == *epoch);
                if let Some(epoch_summary) = maybe_proof {
                    let proof = self.storage.get_proof(epoch_summary).await?;
                    audit_epoch(proof, *qr).await?;
                }
            }
            AuditCommand::ShowEpochs { force_refresh } => {
                let cache_control = if *force_refresh {
                    ProofIndexCacheOption::NoCache
                } else {
                    default_cache_control
                };

                let mut proofs = self.storage.list_proofs(cache_control).await?;
                *self.last_summaries.write().await = Some(proofs.clone());
                display_audit_proofs_info(&mut proofs)?;
            }
            AuditCommand::Ui => crate::ui::UserInterface::execute(Some(self.storage.clone()))?,
        }
        Ok(())
    }
}

fn display_audit_proofs_info(info: &mut Vec<crate::storage::EpochSummary>) -> Result<()> {
    info.sort_by(|a, b| a.name.epoch.cmp(&b.name.epoch));
    if info.is_empty() {
        bail!("There are no epochs present in the storage repository");
    }

    let min = info.first().unwrap().clone();
    let max = info.last().unwrap().clone();
    let (maybe_broken_epoch, is_contiguous) =
        info.iter()
            .skip(1)
            .fold((min.clone(), true), |(previous_item, cont), item| {
                if !cont {
                    (previous_item, cont)
                } else {
                    (
                        item.clone(),
                        item.name.epoch == previous_item.name.epoch + 1,
                    )
                }
            });

    if !is_contiguous {
        bail!("The audit proofs appear to not be continguous. There's a break in the linear history at epoch {}", maybe_broken_epoch.name.epoch);
    }

    println!(
        "Audit history is available between epochs ({}) and ({}), inclusively.",
        min.name.epoch, max.name.epoch
    );

    Ok(())
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed licenses.

//! This module holds the generic storage interaction layer trait along with the underlying implemented
//! storage interactions. We need to implement the trait for a [Box] of the same trait in order to support
//! passing the boxed implementations around. This is required for the async nature of the command and Rust's
//! type inference engine

use akd::local_auditing::{AuditBlob, AuditBlobName};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use std::convert::TryFrom;

// ************************************ Implementations ************************************ //

pub mod dynamodb;
pub mod s3;

// ************************************ Trait and Type Definitions ************************************ //

pub trait CommonStorageClapSettings {
    fn bucket(&self) -> String;
    fn region(&self) -> String;
    fn s3_endpoint(&self) -> Option<String>;
    fn access_key(&self) -> Option<String>;
    fn secret_key(&self) -> Option<String>;

    fn set_bucket(&mut self, v: String);
    fn set_region(&mut self, v: String);
    fn set_s3_endpoint(&mut self, v: Option<String>);
    fn set_access_key(&mut self, v: Option<String>);
    fn set_secret_key(&mut self, v: Option<String>);
}

/// Represents the summary of an epoch, and a unique key referring to the raw object in native storage (if needed)
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
pub struct EpochSummary {
    /// The name of the audit-blob decomposed into parts
    pub name: AuditBlobName,
    /// Unique idenfier for the blob in question
    pub key: String,
}

impl TryFrom<&str> for EpochSummary {
    type Error = anyhow::Error;

    fn try_from(potential_key: &str) -> Result<Self, Self::Error> {
        let name = AuditBlobName::try_from(potential_key).map_err(|err| anyhow!("{:?}", err))?;

        Ok(Self {
            name,
            key: potential_key.to_string(),
        })
    }
}

/// Options for proof index lookup operations
#[derive(Debug, Eq, PartialEq, Clone)]
pub enum ProofIndexCacheOption {
    /// Don't utilize a cache
    NoCache,
    /// Utilize the underlying proof lookup cache
    UseCache,
}

/// Represents a storage of audit proofs and READ ONLY interaction to retrieve the proof objects
#[async_trait]
pub trait AuditProofStorage: Sync + Send + std::fmt::Debug {
    /// The default cache control option for proof listings for this storage implementation
    fn default_cache_control(&self) -> ProofIndexCacheOption;

    /// List the proofs in the storage medium.
    async fn list_proofs(&self, cache_control: ProofIndexCacheOption) -> Result<Vec<EpochSummary>>;

    /// Retrieve an audit proof from the storage medium. If the underlying storage implementation
    /// requires the epoch summaries in order to re-retrieve a specific epoch, it is up to that
    /// implementation to cache the information. Example: See the AWS S3 implementation
    async fn get_proof(&self, epoch: &EpochSummary) -> Result<AuditBlob>;
}

// We need to implement the trait for a Box<dyn Trait> in order to utilize a box further downstream
// which allows us to have multiple implementations of the AuditProofStorage without a separate logic/code
// path for each implementation
//
// See the discussion here: https://stackoverflow.com/questions/71933895/why-cant-boxdyn-trait-be-pased-to-a-function-with-mut-trait-as-parameter
#[async_trait]
impl<APS: ?Sized> AuditProofStorage for Box<APS>
where
    APS: AuditProofStorage,
{
    /// The default cache control option for proof listings for this storage implementation
    #[allow(unconditional_recursion)]
    fn default_cache_control(&self) -> ProofIndexCacheOption {
        self.default_cache_control()
    }

    /// List the proofs in the storage medium.
    async fn list_proofs(&self, cache_control: ProofIndexCacheOption) -> Result<Vec<EpochSummary>> {
        self.list_proofs(cache_control).await
    }

    /// Retrieve an audit proof from the storage medium. If the underlying storage implementation
    /// requires the epoch summaries in order to re-retrieve a specific epoch, it is up to that
    /// implementation to cache the information. Example: See the AWS S3 implementation
    async fn get_proof(&self, epoch: &EpochSummary) -> Result<AuditBlob> {
        self.get_proof(epoch).await
    }
}