// Smoldot
// Copyright (C) 2019-2022  Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! AURA consensus.
//!
//! AURA, for Authority Round, is one of the consensus algorithm available to Substrate-based
//! chains in order to determine who is authorized to generate a block.
//!
//! Every block (with the exception of the genesis block) must contain, in its header, some data
//! that makes it possible to verify that it has been generated by a legitimate author.
//!
//! # Overview of AURA
//!
//! In the AURA algorithm, time is divided into non-overlapping **slots**. How long a slot is
//! never changes, and is determined by calling the `AuraApi_slot_duration` runtime entry point,
//! typically at initialization. The current slot number is equal to `unix_time / slot_duration`.
//! The slot number of a block can be found in its header.
//!
//! > **Note**: Slot durations values are usually around 3 to 20 seconds.
//!
//! A list of authorities (each authority being represented as a public key) is maintained by the
//! chain. The current authorities of a block can be found by calling the `AuraApi_authorities`
//! runtime entry point.
//! When this list is modified, a consensus log item is added to the header of the block, meaning
//! that this runtime entry point doesn't have to be called every single time.
//!
//! During a slot, only the authority whose public key is found at
//! `authorities[slot_number % authorities.len()]` is allowed to produce a block.
//!
//! Each block being produced must include an Aura seal in its header containing a signature of
//! the block header (with the exclusion of the seal itself) made using the public key in question.
//!

use crate::header;

use alloc::vec::Vec;
use core::{num::NonZero, time::Duration};

/// Configuration for [`verify_header`].
pub struct VerifyConfig<'a, TAuthList> {
    /// Header of the block to verify.
    pub header: header::HeaderRef<'a>,

    /// Number of bytes used to encode the block number in the header.
    pub block_number_bytes: usize,

    /// Header of the parent of the block to verify.
    ///
    /// [`verify_header`] assumes that this block has been successfully verified before.
    ///
    /// The hash of this header must be the one referenced in [`VerifyConfig::header`].
    pub parent_block_header: header::HeaderRef<'a>,

    /// Time elapsed since [the Unix Epoch](https://en.wikipedia.org/wiki/Unix_time) (i.e.
    /// 00:00:00 UTC on 1 January 1970), ignoring leap seconds.
    pub now_from_unix_epoch: Duration,

    /// Aura authorities that must validate the block.
    ///
    /// This list is either equal to the parent's list, or, if the parent changes the list of
    /// authorities, equal to that new modified list.
    pub current_authorities: TAuthList,

    /// Duration of a slot in milliseconds.
    /// Can be found by calling the `AuraApi_slot_duration` runtime function.
    pub slot_duration: NonZero<u64>,

    /// If `true`, allow the slot number to be equal to the parent's slot number.
    ///
    /// This is required for parachains which can produce multiple blocks per slot.
    pub allow_equal_slot_number: bool,
}

/// Information yielded back after successfully verifying a block.
#[derive(Debug)]
pub struct VerifySuccess {
    /// `Some` if the list of authorities is modified by this block. Contains the new list of
    /// authorities.
    pub authorities_change: Option<Vec<header::AuraAuthority>>,
}

/// Failure to verify a block.
#[derive(Debug, derive_more::Display, derive_more::Error)]
pub enum VerifyError {
    /// The seal (containing the signature of the authority) is missing from the header.
    MissingSeal,
    /// No pre-runtime digest in the block header.
    MissingPreRuntimeDigest,
    /// Parent block doesn't contain any Aura information.
    ParentIsntAuraConsensus,
    /// Slot number must be strictly increasing between a parent and its child.
    SlotNumberNotIncreasing,
    /// Slot number starts too far in the future.
    TooFarInFuture,
    /// Block header signature is invalid.
    BadSignature,
    /// Failed to parse Ed25519 public key.
    BadPublicKey,
    /// List of authorities is empty.
    EmptyAuthorities,
    /// Header contains more than one new list of authorities.
    MultipleAuthoritiesChangeDigestItems,
}

/// Verifies whether a block header provides a correct proof of the legitimacy of the authorship.
pub fn verify_header<'a>(
    mut config: VerifyConfig<'a, impl ExactSizeIterator<Item = header::AuraAuthorityRef<'a>>>,
) -> Result<VerifySuccess, VerifyError> {
    // TODO: handle OnDisabled

    // Gather the slot number from the header.
    let slot_number = config
        .header
        .digest
        .aura_pre_runtime()
        .ok_or(VerifyError::MissingPreRuntimeDigest)?
        .slot_number;

    // Make sure that the slot number is strictly increasing between a parent and its children.
    if config.parent_block_header.number != 0 {
        let parent_slot_number = match config.parent_block_header.digest.aura_pre_runtime() {
            Some(pr) => pr.slot_number,
            None => return Err(VerifyError::ParentIsntAuraConsensus),
        };

        if config.allow_equal_slot_number {
            if slot_number < parent_slot_number {
                return Err(VerifyError::SlotNumberNotIncreasing);
            }
        } else if slot_number <= parent_slot_number {
            return Err(VerifyError::SlotNumberNotIncreasing);
        }
    }

    // Check that the slot number isn't a slot in the future.
    // Since there might be a clock drift (either locally or on the authority that created the
    // block), a tolerance period is added.
    // If the local node is an authority itself, and the best block uses a slot number `N` seconds
    // in the future, then for the next `N` seconds the local node won't produce any block. As
    // such, a high tolerance level constitutes an attack vector.
    {
        const TOLERANCE: Duration = Duration::from_secs(30);
        let current_slot =
            (config.now_from_unix_epoch + TOLERANCE).as_secs() * 1000 / config.slot_duration.get();
        if slot_number > current_slot {
            return Err(VerifyError::TooFarInFuture);
        }
    };

    // Check whether there is an authority change in the block.
    // This information is used in case of success.
    let mut authorities_change = None;
    for digest_item in config.header.digest.logs() {
        if let header::DigestItemRef::AuraConsensus(
            header::AuraConsensusLogRef::AuthoritiesChange(new_list),
        ) = digest_item
        {
            if authorities_change.is_some() {
                return Err(VerifyError::MultipleAuthoritiesChangeDigestItems);
            }

            authorities_change = Some(new_list.map(Into::into).collect());
            // We continue looping even after finding a list, to make sure there is only one list.
        }
    }

    // The signature in the seal applies to the header from where the signature isn't present.
    // Extract the signature and build the hash that is expected to be signed.
    let (seal_bytes, pre_seal_hash) = {
        let mut unsealed_header = config.header;
        let seal_bytes: [u8; 64] = match unsealed_header.digest.pop_seal() {
            Some(header::Seal::Aura(seal)) => *seal,
            _ => return Err(VerifyError::MissingSeal),
        };
        (seal_bytes, unsealed_header.hash(config.block_number_bytes))
    };

    // Fetch the authority that has supposedly signed the block.
    if config.current_authorities.len() == 0 {
        // Checked beforehand in order to not do a modulo 0 operation.
        return Err(VerifyError::EmptyAuthorities);
    }
    // About overflows:
    // If `config.current_authorities.len()` doesn't fit in a `u64`, then
    // `slot_number % config.current_authorities.len()` is necessarily equal to `slot_number`.
    // Since we're using a modulo operation, `signing_authority` is always inferior to
    // `current_authorities.len()`, meaning that it always fits in a `usize`.
    let signing_authority = usize::try_from(
        slot_number % u64::try_from(config.current_authorities.len()).unwrap_or(u64::MAX),
    )
    .unwrap_or_else(|_| unreachable!());

    let authority_public_key = config
        .current_authorities
        .nth(signing_authority)
        .unwrap()
        .public_key;

    // Try sr25519 verification first, then fall back to ed25519.
    let sr25519_ok = schnorrkel::PublicKey::from_bytes(authority_public_key)
        .and_then(|pk| {
            pk.verify_simple(
                b"substrate",
                &pre_seal_hash,
                &schnorrkel::Signature::from_bytes(&seal_bytes)?,
            )
        })
        .is_ok();

    if !sr25519_ok {
        ed25519_zebra::VerificationKey::try_from(*authority_public_key)
            .map_err(|_| VerifyError::BadSignature)?
            .verify(&ed25519_zebra::Signature::from(seal_bytes), &pre_seal_hash)
            .map_err(|_| VerifyError::BadSignature)?;
    }

    // Success! 🚀
    Ok(VerifySuccess { authorities_change })
}
