zakura-network 2.0.0

use super::{config::*, error::*, events::HeaderSyncRequestId, validation::*, *};

/// Zakura stream kind reserved for native header sync.
pub const ZAKURA_STREAM_HEADER_SYNC: u16 = 5;
/// Version of the native header-sync stream.
///
/// Version 4 carries one tree-aux root for each non-empty range header.
/// Version 6 extends each tree-aux root with the full set of per-block ZIP-221 history-leaf
/// inputs a recipient needs to rebuild the leaf and verify the roots against its own header
/// commitments during header sync, without the block body: the Ironwood note-commitment
/// root, the three per-pool shielded transaction counts (Sapling/Orchard/Ironwood), and the
/// block's ZIP-244 `auth_data_root` (the co-input to its NU5+ header commitment).
/// Version 7 prefixes `GetHeaders`/`Headers` with a non-zero request ID so a response is
/// correlated with the exact request that solicited it, instead of by FIFO arrival order.
/// Each of these is a breaking wire change: a peer that speaks an earlier version cannot
/// exchange header-sync ranges with version 7 and does not negotiate header sync at all.
pub const ZAKURA_HEADER_SYNC_STREAM_VERSION: u16 = 7;

/// Peer status advertisement.
pub const MSG_HS_STATUS: u8 = 1;
/// Request a contiguous range of headers by height.
pub const MSG_HS_GET_HEADERS: u8 = 2;
/// Respond with a contiguous run of headers.
pub const MSG_HS_HEADERS: u8 = 3;
/// Flood a newly seen tip block, including its full body.
pub const MSG_HS_NEW_BLOCK: u8 = 4;

/// Maximum encoded header-sync message bytes.
pub const MAX_HS_MESSAGE_BYTES: usize = 2 * 1024 * 1024;
/// Default number of headers advertised per response.
pub const DEFAULT_HS_RANGE: u32 = 1000;
/// Maximum number of headers ever honored by header sync.
pub const MAX_HS_RANGE: u32 = 4000;
/// Default number of in-flight header requests advertised per peer.
pub const DEFAULT_HS_MAX_INFLIGHT: u16 = 10;

pub(super) const HEADER_SYNC_MESSAGE_TYPE_BYTES: usize = 1;
pub(super) const HEADER_SYNC_REQUEST_ID_BYTES: usize = 8;
pub(super) const HEADER_SYNC_COUNT_BYTES: usize = 4;
pub(super) const HEADER_SYNC_HAS_ROOTS_BYTES: usize = 1;
pub(super) const HEADER_SYNC_BODY_SIZE_BYTES: usize = 4;
/// Encoded [`BlockCommitmentRoots`]: height + Sapling root + Orchard root + Ironwood root
/// + three `u64` shielded tx-counts (Sapling/Orchard/Ironwood) + auth-data root.
pub(super) const HEADER_SYNC_BLOCK_COMMITMENT_ROOTS_BYTES: usize =
    4 + 32 + 32 + 32 + 8 + 8 + 8 + 32;
pub(super) const COMMON_HEADER_BYTES: usize = 1_487;
pub(super) const REGTEST_HEADER_BYTES: usize = 177;
pub(super) const LOCAL_MAX_HS_INFLIGHT_PER_PEER: u16 = 16;
pub(super) const HEADER_SYNC_MAX_RESIDENT_BATCHES: u32 = 16;
pub(super) const HEADER_SYNC_RETRY_AVOIDANCE: Duration = Duration::from_millis(50);
pub(super) const STATUS_PUBLICATION_RETRY_DELAY: Duration = Duration::from_millis(50);
pub(super) const HEADER_SYNC_SEEN_HASH_CAPACITY: usize = 4096;
pub(super) const DEFAULT_HS_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
pub(super) const EMPTY_HEADERS_RETRY_DELAY: Duration = Duration::from_secs(1);
pub(super) const DEFAULT_HS_STATUS_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
// v1 semantic meters intentionally use strict spacing for unsolicited status refreshes and
// distinct unseen full-block floods. Cheap duplicate `NewBlock` echoes are deduped before this
// meter, so they do not consume tokens or cause honest peers to be scored.
pub(super) const DEFAULT_HS_INBOUND_STATUS_MIN_INTERVAL: Duration = Duration::from_secs(5);
pub(super) const DEFAULT_HS_INBOUND_NEW_BLOCK_MIN_INTERVAL: Duration = Duration::from_secs(5);

const _: () = assert!(MAX_HS_MESSAGE_BYTES < LOCAL_MAX_MESSAGE_BYTES as usize);
const _: () = assert!(
    HEADER_SYNC_MESSAGE_TYPE_BYTES
        + HEADER_SYNC_REQUEST_ID_BYTES
        + HEADER_SYNC_COUNT_BYTES
        + (COMMON_HEADER_BYTES
            + HEADER_SYNC_BODY_SIZE_BYTES
            + HEADER_SYNC_BLOCK_COMMITMENT_ROOTS_BYTES)
            * (DEFAULT_HS_RANGE as usize)
        < MAX_HS_MESSAGE_BYTES
);

/// Native versioned header-sync message.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum HeaderSyncMessage {
    /// Peer tip, anchor, and served-range advertisement.
    Status(HeaderSyncStatus),
    /// Request `count` headers starting at `start_height`.
    GetHeaders {
        /// First requested height.
        start_height: block::Height,
        /// Requested header count.
        count: u32,
        /// Whether the requester wants all-or-nothing tree-aux roots.
        /// A sender who is syncing in vct mode will always request these.
        /// A sender who is syncing in non-checkpoint mode does not need these but still requests them.
        /// A sender who is syncing above the last checkpoint height does not request these.
        want_tree_aux_roots: bool,
    },
    /// A bounded contiguous header run with one advisory body-size hint per header.
    ///
    /// A `0` size means "unknown"; the hint is not consensus data. Tree-aux roots
    /// are peer-sourced execution hints and are verified by state before use.
    Headers {
        /// Headers in ascending height order.
        headers: Vec<Arc<block::Header>>,
        /// Advisory serialized body sizes, parallel to `headers`.
        body_sizes: Vec<u32>,
        /// Per-height commitment roots, parallel to `headers`.
        tree_aux_roots: Vec<BlockCommitmentRoots>,
    },
    /// Full block tip-flood payload.
    NewBlock(Arc<block::Block>),
}

impl HeaderSyncMessage {
    /// Returns this message's header-sync discriminator.
    pub fn message_type(&self) -> u8 {
        match self {
            Self::Status(_) => MSG_HS_STATUS,
            Self::GetHeaders { .. } => MSG_HS_GET_HEADERS,
            Self::Headers { .. } => MSG_HS_HEADERS,
            Self::NewBlock(_) => MSG_HS_NEW_BLOCK,
        }
    }

    /// Encode this message as `[u8 message_type][request id][bounded fields...]`.
    ///
    /// `request_id` is required by `GetHeaders`/`Headers` and unused by the other
    /// message types, which are not correlated.
    pub fn encode(
        &self,
        request_id: Option<HeaderSyncRequestId>,
    ) -> Result<Vec<u8>, HeaderSyncWireError> {
        let mut bytes = Vec::new();
        bytes.write_u8(self.message_type())?;
        match self {
            Self::Status(status) => status.encode_to(&mut bytes)?,
            Self::GetHeaders {
                start_height,
                count,
                want_tree_aux_roots,
            } => {
                validate_get_headers_count(*count)?;
                bytes.write_u64::<LittleEndian>(
                    request_id
                        .ok_or(HeaderSyncWireError::MissingRequestId {
                            message: "GetHeaders",
                        })?
                        .get(),
                )?;
                write_height(&mut bytes, *start_height)?;
                bytes.write_u32::<LittleEndian>(*count)?;
                bytes.write_u8(u8::from(*want_tree_aux_roots))?;
            }
            Self::Headers {
                headers,
                body_sizes,
                tree_aux_roots,
            } => {
                validate_headers_len(headers.len(), usize_from_u32(MAX_HS_RANGE, "headers cap")?)?;
                validate_body_sizes_len(headers.len(), body_sizes.len())?;
                validate_tree_aux_roots_len(headers.len(), tree_aux_roots.len())?;
                bytes.write_u64::<LittleEndian>(
                    request_id
                        .ok_or(HeaderSyncWireError::MissingRequestId { message: "Headers" })?
                        .get(),
                )?;
                bytes.write_u32::<LittleEndian>(u32_from_usize(headers.len(), "headers count")?)?;
                bytes.write_u8(u8::from(!tree_aux_roots.is_empty()))?;
                for (header, body_size) in headers.iter().zip(body_sizes) {
                    header.zcash_serialize(&mut bytes)?;
                    bytes.write_u32::<LittleEndian>(*body_size)?;
                }
                for roots in tree_aux_roots {
                    roots.zcash_serialize(&mut bytes)?;
                }
            }
            Self::NewBlock(block) => {
                block.zcash_serialize(&mut bytes)?;
            }
        }
        if bytes.len() > MAX_HS_MESSAGE_BYTES {
            return Err(HeaderSyncWireError::OversizedPayload {
                actual: bytes.len(),
                max: MAX_HS_MESSAGE_BYTES,
            });
        }
        Ok(bytes)
    }

    /// Decode a header-sync message and its request ID using the request bounds in `context`.
    ///
    /// The request ID is `None` for the message types that do not carry one.
    pub fn decode(
        bytes: &[u8],
        context: HeaderSyncDecodeContext,
    ) -> Result<(Self, Option<HeaderSyncRequestId>), HeaderSyncWireError> {
        if bytes.len() > MAX_HS_MESSAGE_BYTES {
            return Err(HeaderSyncWireError::OversizedPayload {
                actual: bytes.len(),
                max: MAX_HS_MESSAGE_BYTES,
            });
        }
        let mut reader = Cursor::new(bytes);
        let message_type = reader.read_u8()?;
        let mut request_id = None;
        let message = match message_type {
            MSG_HS_STATUS => Self::Status(HeaderSyncStatus::decode_from(&mut reader)?),
            MSG_HS_GET_HEADERS => {
                request_id = HeaderSyncRequestId::new(reader.read_u64::<LittleEndian>()?);
                if request_id.is_none() {
                    return Err(HeaderSyncWireError::MissingRequestId {
                        message: "GetHeaders",
                    });
                }
                let start_height = read_height(&mut reader)?;
                let count = reader.read_u32::<LittleEndian>()?;
                let want_tree_aux_roots = read_bool_marker(&mut reader, "want_tree_aux_roots")?;
                validate_get_headers_count(count)?;
                Self::GetHeaders {
                    start_height,
                    count,
                    want_tree_aux_roots,
                }
            }
            MSG_HS_HEADERS => {
                request_id = HeaderSyncRequestId::new(reader.read_u64::<LittleEndian>()?);
                if request_id.is_none() {
                    return Err(HeaderSyncWireError::MissingRequestId { message: "Headers" });
                }
                let count = usize_from_u32(reader.read_u32::<LittleEndian>()?, "headers count")?;
                let has_roots = read_bool_marker(&mut reader, "has_roots")?;
                let Some(max_headers) = context.headers_response_limit()? else {
                    return Err(HeaderSyncWireError::UnsolicitedHeaders);
                };
                if has_roots && !context.wants_tree_aux_roots() {
                    return Err(HeaderSyncWireError::UnrequestedTreeAuxRoots);
                }
                validate_headers_len(count, max_headers)?;
                let mut headers = Vec::with_capacity(count);
                let mut body_sizes = Vec::with_capacity(count);
                let mut tree_aux_roots = if has_roots {
                    Vec::with_capacity(count)
                } else {
                    Vec::new()
                };
                for _ in 0..count {
                    headers.push(Arc::new(block::Header::zcash_deserialize(&mut reader)?));
                    body_sizes.push(reader.read_u32::<LittleEndian>()?);
                }
                if has_roots {
                    for _ in 0..count {
                        tree_aux_roots.push(BlockCommitmentRoots::zcash_deserialize(&mut reader)?);
                    }
                }
                validate_tree_aux_roots_len(count, tree_aux_roots.len())?;
                if let Some(requested) = context.requested {
                    validate_tree_aux_root_heights(requested.start_height, &tree_aux_roots)?;
                }
                Self::Headers {
                    headers,
                    body_sizes,
                    tree_aux_roots,
                }
            }
            MSG_HS_NEW_BLOCK => {
                Self::NewBlock(Arc::new(block::Block::zcash_deserialize(&mut reader)?))
            }
            value => return Err(HeaderSyncWireError::UnknownMessageType(value)),
        };
        reject_trailing(bytes, &reader)?;
        Ok((message, request_id))
    }

    /// Return a `Headers` request ID without decoding the full header payload.
    pub fn peek_headers_request_id(
        bytes: &[u8],
    ) -> Result<HeaderSyncRequestId, HeaderSyncWireError> {
        let mut reader = Cursor::new(bytes);
        let message_type = reader.read_u8()?;
        if message_type != MSG_HS_HEADERS {
            return Err(HeaderSyncWireError::UnknownMessageType(message_type));
        }
        HeaderSyncRequestId::new(reader.read_u64::<LittleEndian>()?)
            .ok_or(HeaderSyncWireError::MissingRequestId { message: "Headers" })
    }

    /// Convert this message into a bounded Zakura frame.
    pub fn encode_frame(
        &self,
        request_id: Option<HeaderSyncRequestId>,
    ) -> Result<Frame, HeaderSyncWireError> {
        Ok(Frame {
            message_type: u16::from(self.message_type()),
            flags: 0,
            payload: self.encode(request_id)?,
        })
    }

    /// Decode this message and its request ID from a Zakura frame, after checking flags
    /// and type agreement.
    pub fn decode_frame(
        frame: Frame,
        context: HeaderSyncDecodeContext,
    ) -> Result<(Self, Option<HeaderSyncRequestId>), HeaderSyncWireError> {
        if frame.flags != 0 {
            return Err(HeaderSyncWireError::UnsupportedFlags(frame.flags));
        }
        let (message, request_id) = Self::decode(&frame.payload, context)?;
        let frame_message_type = u8::try_from(frame.message_type)
            .map_err(|_| HeaderSyncWireError::UnknownFrameMessageType(frame.message_type))?;
        if frame_message_type != message.message_type() {
            return Err(HeaderSyncWireError::MismatchedFrameMessageType {
                frame: frame.message_type,
                payload: message.message_type(),
            });
        }
        Ok((message, request_id))
    }
}