synor/crates/synor-network/src/message.rs

//! Network message types for Synor.

use borsh::{BorshDeserialize, BorshSerialize};
use serde::{Deserialize, Serialize};
use synor_types::{BlockHeader, BlockId, Hash256, Transaction, TransactionId};

/// Message types for P2P communication.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub enum Message {
    /// Block announcement.
    BlockAnnouncement(BlockAnnouncement),
    /// Transaction announcement.
    TransactionAnnouncement(TransactionAnnouncement),
    /// Header announcement.
    HeaderAnnouncement(HeaderAnnouncement),
    /// Ping message.
    Ping(u64),
    /// Pong response.
    Pong(u64),
}

impl Message {
    /// Returns the message type as a string.
    pub fn message_type(&self) -> &'static str {
        match self {
            Message::BlockAnnouncement(_) => "block_announcement",
            Message::TransactionAnnouncement(_) => "tx_announcement",
            Message::HeaderAnnouncement(_) => "header_announcement",
            Message::Ping(_) => "ping",
            Message::Pong(_) => "pong",
        }
    }

    /// Serializes the message to bytes.
    pub fn to_bytes(&self) -> Result<Vec<u8>, borsh::io::Error> {
        borsh::to_vec(self)
    }

    /// Deserializes a message from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, borsh::io::Error> {
        borsh::from_slice(bytes)
    }
}

/// Announcement of a new block.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
pub struct BlockAnnouncement {
    /// Block hash.
    pub hash: BlockId,
    /// Block header.
    pub header: BlockHeader,
    /// Number of transactions in the block.
    pub tx_count: u32,
    /// Blue score of the block.
    pub blue_score: u64,
}

impl BlockAnnouncement {
    /// Creates a new block announcement.
    pub fn new(header: BlockHeader, tx_count: u32, blue_score: u64) -> Self {
        let hash = header.block_id();
        BlockAnnouncement {
            hash,
            header,
            tx_count,
            blue_score,
        }
    }
}

/// Announcement of a new transaction.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
pub struct TransactionAnnouncement {
    /// Transaction ID.
    pub txid: TransactionId,
    /// Transaction (optional, may be just the ID for inv messages).
    pub transaction: Option<Transaction>,
    /// Fee rate in sompi per byte.
    pub fee_rate: Option<u64>,
}

impl TransactionAnnouncement {
    /// Creates a new transaction announcement with full transaction.
    pub fn with_transaction(tx: Transaction, fee_rate: u64) -> Self {
        let txid = tx.txid();
        TransactionAnnouncement {
            txid,
            transaction: Some(tx),
            fee_rate: Some(fee_rate),
        }
    }

    /// Creates a new transaction announcement with just the ID.
    pub fn id_only(txid: TransactionId) -> Self {
        TransactionAnnouncement {
            txid,
            transaction: None,
            fee_rate: None,
        }
    }
}

/// Announcement of a new header.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
pub struct HeaderAnnouncement {
    /// Block hash.
    pub hash: BlockId,
    /// Block header.
    pub header: BlockHeader,
    /// Blue score.
    pub blue_score: u64,
    /// DAA score.
    pub daa_score: u64,
}

impl HeaderAnnouncement {
    /// Creates a new header announcement.
    pub fn new(header: BlockHeader, blue_score: u64, daa_score: u64) -> Self {
        let hash = header.block_id();
        HeaderAnnouncement {
            hash,
            header,
            blue_score,
            daa_score,
        }
    }
}

/// Inventory item types.
#[derive(Clone, Copy, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub enum InvType {
    /// Block.
    Block,
    /// Transaction.
    Transaction,
    /// Block header.
    Header,
}

/// Inventory item.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct InvItem {
    /// Item type.
    pub inv_type: InvType,
    /// Item hash.
    pub hash: Hash256,
}

impl InvItem {
    /// Creates a block inventory item.
    pub fn block(hash: Hash256) -> Self {
        InvItem {
            inv_type: InvType::Block,
            hash,
        }
    }

    /// Creates a transaction inventory item.
    pub fn transaction(hash: Hash256) -> Self {
        InvItem {
            inv_type: InvType::Transaction,
            hash,
        }
    }

    /// Creates a header inventory item.
    pub fn header(hash: Hash256) -> Self {
        InvItem {
            inv_type: InvType::Header,
            hash,
        }
    }
}

/// Inventory message for announcing available data.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct Inventory {
    /// List of inventory items.
    pub items: Vec<InvItem>,
}

impl Inventory {
    /// Creates a new inventory message.
    pub fn new(items: Vec<InvItem>) -> Self {
        Inventory { items }
    }

    /// Creates an empty inventory.
    pub fn empty() -> Self {
        Inventory { items: Vec::new() }
    }

    /// Adds an item to the inventory.
    pub fn add(&mut self, item: InvItem) {
        self.items.push(item);
    }

    /// Returns the number of items.
    pub fn len(&self) -> usize {
        self.items.len()
    }

    /// Returns true if empty.
    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }
}

/// Get data request message.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct GetData {
    /// Requested items.
    pub items: Vec<InvItem>,
}

impl GetData {
    /// Creates a new get data request.
    pub fn new(items: Vec<InvItem>) -> Self {
        GetData { items }
    }

    /// Creates a request for blocks.
    pub fn blocks(hashes: Vec<Hash256>) -> Self {
        GetData {
            items: hashes.into_iter().map(InvItem::block).collect(),
        }
    }

    /// Creates a request for transactions.
    pub fn transactions(hashes: Vec<Hash256>) -> Self {
        GetData {
            items: hashes.into_iter().map(InvItem::transaction).collect(),
        }
    }

    /// Creates a request for headers.
    pub fn headers(hashes: Vec<Hash256>) -> Self {
        GetData {
            items: hashes.into_iter().map(InvItem::header).collect(),
        }
    }
}

/// Get headers request.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct GetHeaders {
    /// Start hash (get headers after this).
    pub start_hash: Hash256,
    /// Stop hash (optional, get up to this hash).
    pub stop_hash: Option<Hash256>,
    /// Maximum number of headers to return.
    pub max_headers: u32,
}

impl GetHeaders {
    /// Creates a new get headers request.
    pub fn new(start_hash: Hash256, stop_hash: Option<Hash256>, max_headers: u32) -> Self {
        GetHeaders {
            start_hash,
            stop_hash,
            max_headers,
        }
    }
}

/// Headers response.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct Headers {
    /// List of headers.
    pub headers: Vec<BlockHeader>,
}

impl Headers {
    /// Creates a new headers response.
    pub fn new(headers: Vec<BlockHeader>) -> Self {
        Headers { headers }
    }

    /// Creates an empty headers response.
    pub fn empty() -> Self {
        Headers {
            headers: Vec::new(),
        }
    }
}

/// Status message for handshake.
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
pub struct StatusMessage {
    /// Protocol version.
    pub version: u32,
    /// Chain ID magic bytes.
    pub chain_id: [u8; 4],
    /// Genesis block hash.
    pub genesis_hash: Hash256,
    /// Best block hash (current tip).
    pub best_hash: Hash256,
    /// Best block blue score.
    pub best_blue_score: u64,
    /// Best block DAA score.
    pub best_daa_score: u64,
    /// Timestamp.
    pub timestamp: u64,
    /// User agent.
    pub user_agent: String,
}

impl StatusMessage {
    /// Protocol version constant.
    pub const PROTOCOL_VERSION: u32 = 1;

    /// Creates a new status message.
    pub fn new(
        chain_id: [u8; 4],
        genesis_hash: Hash256,
        best_hash: Hash256,
        best_blue_score: u64,
        best_daa_score: u64,
        user_agent: String,
    ) -> Self {
        StatusMessage {
            version: Self::PROTOCOL_VERSION,
            chain_id,
            genesis_hash,
            best_hash,
            best_blue_score,
            best_daa_score,
            timestamp: chrono::Utc::now().timestamp() as u64,
            user_agent,
        }
    }

    /// Checks if this status is compatible with ours.
    pub fn is_compatible(&self, our_chain_id: &[u8; 4], our_genesis: &Hash256) -> bool {
        self.chain_id == *our_chain_id && self.genesis_hash == *our_genesis
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_message_serialization() {
        let msg = Message::Ping(12345);
        let bytes = msg.to_bytes().unwrap();
        let decoded = Message::from_bytes(&bytes).unwrap();

        match decoded {
            Message::Ping(n) => assert_eq!(n, 12345),
            _ => panic!("Wrong message type"),
        }
    }

    #[test]
    fn test_inventory() {
        let mut inv = Inventory::empty();
        inv.add(InvItem::block(Hash256::from([1u8; 32])));
        inv.add(InvItem::transaction(Hash256::from([2u8; 32])));

        assert_eq!(inv.len(), 2);
    }

    #[test]
    fn test_get_data() {
        let hashes = vec![Hash256::from([1u8; 32]), Hash256::from([2u8; 32])];
        let get_data = GetData::blocks(hashes);

        assert_eq!(get_data.items.len(), 2);
        assert_eq!(get_data.items[0].inv_type, InvType::Block);
    }
}