quicproquo/crates/meshservice/src/message.rs

//! Core message types for the service layer.

use std::time::{SystemTime, UNIX_EPOCH};

use serde::{Deserialize, Serialize};

use crate::identity::ServiceIdentity;
use crate::verification::Verification;

/// Message types within a service.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[repr(u8)]
pub enum MessageType {
    /// Provider announces availability.
    Announce = 0x01,
    /// Consumer queries for matches.
    Query = 0x02,
    /// Response to a query.
    Response = 0x03,
    /// Consumer reserves a slot/item.
    Reserve = 0x04,
    /// Provider confirms/rejects reservation.
    Confirm = 0x05,
    /// Either party cancels.
    Cancel = 0x06,
    /// Provider updates an existing announce (partial).
    Update = 0x07,
    /// Provider revokes an announce.
    Revoke = 0x08,
}

impl TryFrom<u8> for MessageType {
    type Error = ();
    
    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0x01 => Ok(MessageType::Announce),
            0x02 => Ok(MessageType::Query),
            0x03 => Ok(MessageType::Response),
            0x04 => Ok(MessageType::Reserve),
            0x05 => Ok(MessageType::Confirm),
            0x06 => Ok(MessageType::Cancel),
            0x07 => Ok(MessageType::Update),
            0x08 => Ok(MessageType::Revoke),
            _ => Err(()),
        }
    }
}

/// A generic service message that can carry any application payload.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceMessage {
    /// Service identifier (which application).
    pub service_id: u32,
    /// Message type within service.
    pub message_type: MessageType,
    /// Protocol version for forward compatibility.
    pub version: u8,
    /// Unique message ID.
    pub id: [u8; 16],
    /// Sender's mesh address.
    pub sender_address: [u8; 16],
    /// Application-specific CBOR payload.
    pub payload: Vec<u8>,
    /// Ed25519 signature over signable fields.
    pub signature: Vec<u8>,
    /// Optional verifications from trusted parties.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub verifications: Vec<Verification>,
    /// Monotonically increasing per sender (dedup/supersede).
    pub sequence: u64,
    /// Time-to-live in hours.
    pub ttl_hours: u16,
    /// Unix timestamp of creation.
    pub timestamp: u64,
    /// Current hop count (incremented on re-broadcast).
    pub hop_count: u8,
    /// Maximum propagation hops.
    pub max_hops: u8,
}

/// Default TTL: 7 days.
const DEFAULT_TTL_HOURS: u16 = 168;
/// Default max hops.
const DEFAULT_MAX_HOPS: u8 = 8;

impl ServiceMessage {
    /// Create a new service message.
    pub fn new(
        identity: &ServiceIdentity,
        service_id: u32,
        message_type: MessageType,
        payload: Vec<u8>,
        sequence: u64,
    ) -> Self {
        Self::with_options(
            identity,
            service_id,
            message_type,
            payload,
            sequence,
            DEFAULT_TTL_HOURS,
            DEFAULT_MAX_HOPS,
        )
    }

    /// Create with custom TTL and max hops.
    pub fn with_options(
        identity: &ServiceIdentity,
        service_id: u32,
        message_type: MessageType,
        payload: Vec<u8>,
        sequence: u64,
        ttl_hours: u16,
        max_hops: u8,
    ) -> Self {
        use sha2::{Digest, Sha256};

        let sender_address = identity.address();
        
        // Generate unique ID from address + sequence
        let id_hash = Sha256::digest(
            [&sender_address[..], &sequence.to_le_bytes()].concat()
        );
        let mut id = [0u8; 16];
        id.copy_from_slice(&id_hash[..16]);

        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let mut msg = Self {
            service_id,
            message_type,
            version: 1,
            id,
            sender_address,
            payload,
            signature: Vec::new(),
            verifications: Vec::new(),
            sequence,
            ttl_hours,
            timestamp,
            hop_count: 0,
            max_hops,
        };

        let signable = msg.signable_bytes();
        msg.signature = identity.sign(&signable).to_vec();
        msg
    }

    /// Create an announce message.
    pub fn announce(
        identity: &ServiceIdentity,
        service_id: u32,
        payload: Vec<u8>,
        sequence: u64,
    ) -> Self {
        Self::new(identity, service_id, MessageType::Announce, payload, sequence)
    }

    /// Create a query message.
    pub fn query(
        identity: &ServiceIdentity,
        service_id: u32,
        payload: Vec<u8>,
    ) -> Self {
        // Queries use random sequence (not monotonic)
        let sequence = rand::random();
        Self::with_options(
            identity,
            service_id,
            MessageType::Query,
            payload,
            sequence,
            1, // 1 hour TTL for queries
            DEFAULT_MAX_HOPS,
        )
    }

    /// Create a response message.
    pub fn response(
        identity: &ServiceIdentity,
        service_id: u32,
        query_id: [u8; 16],
        payload: Vec<u8>,
    ) -> Self {
        let mut msg = Self::new(
            identity,
            service_id,
            MessageType::Response,
            payload,
            rand::random(),
        );
        // Response ID matches query ID for correlation
        msg.id = query_id;
        msg
    }

    /// Assemble bytes for signing/verification.
    /// Excludes signature, hop_count, verifications (mutable fields).
    fn signable_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(256);
        buf.extend_from_slice(&self.service_id.to_le_bytes());
        buf.push(self.message_type as u8);
        buf.push(self.version);
        buf.extend_from_slice(&self.id);
        buf.extend_from_slice(&self.sender_address);
        buf.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
        buf.extend_from_slice(&self.payload);
        buf.extend_from_slice(&self.sequence.to_le_bytes());
        buf.extend_from_slice(&self.ttl_hours.to_le_bytes());
        buf.extend_from_slice(&self.timestamp.to_le_bytes());
        buf.push(self.max_hops);
        buf
    }

    /// Verify the signature using the sender's public key.
    pub fn verify(&self, sender_public_key: &[u8; 32]) -> bool {
        use crate::identity::compute_address;

        // Verify address matches key
        if compute_address(sender_public_key) != self.sender_address {
            return false;
        }

        let sig: [u8; 64] = match self.signature.as_slice().try_into() {
            Ok(s) => s,
            Err(_) => return false,
        };

        let signable = self.signable_bytes();
        ServiceIdentity::verify(sender_public_key, &signable, &sig)
    }

    /// Check if the message has expired.
    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let ttl_secs = u64::from(self.ttl_hours) * 3600;
        now.saturating_sub(self.timestamp) > ttl_secs
    }

    /// Check if the message can still propagate.
    pub fn can_propagate(&self) -> bool {
        self.hop_count < self.max_hops && !self.is_expired()
    }

    /// Create a forwarded copy with incremented hop count.
    pub fn forwarded(&self) -> Self {
        let mut copy = self.clone();
        copy.hop_count = copy.hop_count.saturating_add(1);
        copy
    }

    /// Get the highest verification level attached.
    pub fn verification_level(&self) -> u8 {
        self.verifications
            .iter()
            .map(|v| v.level)
            .max()
            .unwrap_or(0)
    }

    /// Add a verification to the message.
    pub fn add_verification(&mut self, verification: Verification) {
        self.verifications.push(verification);
    }
}

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

    #[test]
    fn create_and_verify() {
        let id = ServiceIdentity::generate();
        let msg = ServiceMessage::announce(
            &id,
            crate::service_ids::FAPP,
            b"test payload".to_vec(),
            1,
        );
        
        assert!(msg.verify(&id.public_key()));
        assert!(!msg.is_expired());
        assert!(msg.can_propagate());
        assert_eq!(msg.hop_count, 0);
    }

    #[test]
    fn forwarded_increments_hop() {
        let id = ServiceIdentity::generate();
        let msg = ServiceMessage::announce(&id, 1, vec![], 1);
        let fwd = msg.forwarded();
        
        assert_eq!(fwd.hop_count, 1);
        assert!(fwd.verify(&id.public_key())); // Still valid
    }

    #[test]
    fn tampered_fails_verify() {
        let id = ServiceIdentity::generate();
        let mut msg = ServiceMessage::announce(&id, 1, b"original".to_vec(), 1);
        msg.payload = b"tampered".to_vec();
        
        assert!(!msg.verify(&id.public_key()));
    }

    #[test]
    fn query_has_short_ttl() {
        let id = ServiceIdentity::generate();
        let msg = ServiceMessage::query(&id, 1, vec![]);
        
        assert_eq!(msg.ttl_hours, 1);
    }
}