qudag_protocol/
message.rs

1//! Protocol message implementation.
2
3use qudag_crypto::{Ciphertext, MlDsaKeyPair, MlDsaPublicKey, MlKem768, PublicKey, SecretKey};
4use qudag_dag::vertex::VertexId;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use thiserror::Error;
8use uuid::Uuid;
9
10/// Errors that can occur during message operations.
11#[derive(Debug, Error)]
12pub enum MessageError {
13    /// Invalid message format
14    #[error("Invalid message format")]
15    InvalidFormat,
16
17    /// Message too large
18    #[error("Message too large: {0} bytes")]
19    MessageTooLarge(usize),
20
21    /// Invalid signature
22    #[error("Invalid signature")]
23    InvalidSignature,
24
25    /// Missing signature when required
26    #[error("Missing signature")]
27    MissingSignature,
28
29    /// Signing failed
30    #[error("Message signing failed")]
31    SigningFailed,
32
33    /// Verification failed
34    #[error("Signature verification failed")]
35    VerificationFailed,
36
37    /// Encryption failed
38    #[error("Encryption failed")]
39    EncryptionFailed,
40
41    /// Decryption failed
42    #[error("Decryption failed")]
43    DecryptionFailed,
44
45    /// Serialization failed
46    #[error("Message serialization failed")]
47    SerializationFailed,
48
49    /// Deserialization failed
50    #[error("Message deserialization failed")]
51    DeserializationFailed,
52
53    /// Message expired
54    #[error("Message has expired")]
55    MessageExpired,
56
57    /// Invalid timestamp
58    #[error("Invalid message timestamp")]
59    InvalidTimestamp,
60
61    /// Incompatible protocol version
62    #[error("Incompatible protocol version: {0:?}")]
63    IncompatibleVersion(ProtocolVersion),
64}
65
66/// Protocol version information
67#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
68pub struct ProtocolVersion {
69    /// Major version
70    pub major: u16,
71    /// Minor version
72    pub minor: u16,
73    /// Patch version
74    pub patch: u16,
75    /// Protocol features supported
76    pub features: Vec<String>,
77}
78
79impl ProtocolVersion {
80    /// Current protocol version
81    pub const CURRENT: ProtocolVersion = ProtocolVersion {
82        major: 1,
83        minor: 0,
84        patch: 0,
85        features: vec![],
86    };
87
88    /// Check if this version is compatible with another
89    pub fn is_compatible(&self, other: &ProtocolVersion) -> bool {
90        self.major == other.major && self.minor <= other.minor
91    }
92}
93
94/// Message type enumeration.
95#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
96pub enum MessageType {
97    /// Protocol handshake messages
98    Handshake(HandshakeType),
99
100    /// DAG consensus messages
101    Consensus(ConsensusMessageType),
102
103    /// Network routing messages
104    Routing(RoutingMessageType),
105
106    /// Anonymous communication messages
107    Anonymous(AnonymousMessageType),
108
109    /// Protocol control messages
110    Control(ControlMessageType),
111
112    /// State synchronization messages
113    Sync(SyncMessageType),
114
115    /// Generic data messages
116    Data(Vec<u8>),
117}
118
119/// Handshake message types
120#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
121pub enum HandshakeType {
122    /// Initial handshake request
123    Init,
124    /// Handshake response with key exchange
125    Response,
126    /// Handshake completion confirmation
127    Complete,
128    /// Protocol version negotiation
129    VersionNegotiation,
130}
131
132/// Consensus message types for DAG
133#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
134pub enum ConsensusMessageType {
135    /// New vertex proposal
136    VertexProposal,
137    /// Vertex vote
138    Vote,
139    /// Finality announcement
140    Finality,
141    /// Query for missing vertices
142    Query,
143}
144
145/// Routing message types
146#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
147pub enum RoutingMessageType {
148    /// Onion routing message
149    OnionRouted,
150    /// Direct peer message
151    Direct,
152    /// Anonymous broadcast
153    Broadcast,
154}
155
156/// Anonymous communication message types
157#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
158pub enum AnonymousMessageType {
159    /// Anonymous data payload
160    Data,
161    /// Mix network message
162    Mix,
163    /// Cover traffic
164    Cover,
165}
166
167/// Control message types
168#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
169pub enum ControlMessageType {
170    /// Ping for connectivity
171    Ping,
172    /// Pong response
173    Pong,
174    /// Disconnect notification
175    Disconnect,
176    /// Keep-alive message
177    KeepAlive,
178}
179
180/// Synchronization message types
181#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
182pub enum SyncMessageType {
183    /// Request state sync
184    StateRequest,
185    /// State sync response
186    StateResponse,
187    /// Delta sync
188    DeltaSync,
189    /// Checkpoint sync
190    CheckpointSync,
191}
192
193/// Protocol message structure.
194#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct Message {
196    /// Unique message identifier
197    pub id: Uuid,
198
199    /// Protocol version
200    pub version: ProtocolVersion,
201
202    /// Message type
203    pub msg_type: MessageType,
204
205    /// Message payload
206    pub payload: Vec<u8>,
207
208    /// Message timestamp (Unix timestamp in milliseconds)
209    pub timestamp: u64,
210
211    /// Message signature (ML-DSA)
212    pub signature: Option<Vec<u8>>,
213
214    /// Message headers for metadata
215    pub headers: HashMap<String, String>,
216
217    /// Sender's public key hash for verification
218    pub sender_key_hash: Option<Vec<u8>>,
219
220    /// Message sequence number for ordering
221    pub sequence: u64,
222
223    /// Time-to-live for message expiration
224    pub ttl: Option<u64>,
225}
226
227impl Message {
228    /// Create a new message
229    pub fn new(msg_type: MessageType, payload: Vec<u8>) -> Self {
230        Self {
231            id: Uuid::new_v4(),
232            version: ProtocolVersion::CURRENT,
233            msg_type,
234            payload,
235            timestamp: std::time::SystemTime::now()
236                .duration_since(std::time::UNIX_EPOCH)
237                .unwrap()
238                .as_millis() as u64,
239            signature: None,
240            headers: HashMap::new(),
241            sender_key_hash: None,
242            sequence: 0,
243            ttl: None,
244        }
245    }
246
247    /// Create a new message with version
248    pub fn new_with_version(
249        version: ProtocolVersion,
250        msg_type: MessageType,
251        payload: Vec<u8>,
252    ) -> Self {
253        Self {
254            id: Uuid::new_v4(),
255            version,
256            msg_type,
257            payload,
258            timestamp: std::time::SystemTime::now()
259                .duration_since(std::time::UNIX_EPOCH)
260                .unwrap()
261                .as_millis() as u64,
262            signature: None,
263            headers: HashMap::new(),
264            sender_key_hash: None,
265            sequence: 0,
266            ttl: None,
267        }
268    }
269
270    /// Set message sequence number
271    pub fn with_sequence(mut self, sequence: u64) -> Self {
272        self.sequence = sequence;
273        self
274    }
275
276    /// Set message TTL
277    pub fn with_ttl(mut self, ttl: u64) -> Self {
278        self.ttl = Some(ttl);
279        self
280    }
281
282    /// Add header to message
283    pub fn with_header(mut self, key: String, value: String) -> Self {
284        self.headers.insert(key, value);
285        self
286    }
287
288    /// Check if message has expired
289    pub fn is_expired(&self) -> bool {
290        if let Some(ttl) = self.ttl {
291            let now = std::time::SystemTime::now()
292                .duration_since(std::time::UNIX_EPOCH)
293                .unwrap()
294                .as_millis() as u64;
295            now > self.timestamp + ttl
296        } else {
297            false
298        }
299    }
300
301    /// Get message for signing (excludes signature field)
302    fn get_signable_data(&self) -> Result<Vec<u8>, MessageError> {
303        let mut msg_copy = self.clone();
304        msg_copy.signature = None;
305
306        bincode::serialize(&msg_copy).map_err(|_| MessageError::SerializationFailed)
307    }
308
309    /// Sign message with ML-DSA
310    pub fn sign(&mut self, keypair: &MlDsaKeyPair) -> Result<(), MessageError> {
311        let signable_data = self.get_signable_data()?;
312
313        // Sign using the keypair directly
314        let signature = keypair
315            .sign(&signable_data, &mut rand::thread_rng())
316            .map_err(|_| MessageError::SigningFailed)?;
317
318        self.signature = Some(signature);
319
320        // Set sender key hash for verification
321        let public_key_bytes = keypair.public_key();
322        self.sender_key_hash = Some(blake3::hash(public_key_bytes).as_bytes().to_vec());
323
324        Ok(())
325    }
326
327    /// Verify message signature with ML-DSA
328    pub fn verify(&self, public_key: &MlDsaPublicKey) -> Result<bool, MessageError> {
329        let signature = self
330            .signature
331            .as_ref()
332            .ok_or(MessageError::MissingSignature)?;
333
334        // Verify sender key hash matches
335        if let Some(sender_hash) = &self.sender_key_hash {
336            let public_key_bytes = public_key.as_bytes();
337            let expected_hash = blake3::hash(public_key_bytes).as_bytes().to_vec();
338            if sender_hash != &expected_hash {
339                return Ok(false);
340            }
341        }
342
343        let signable_data = self.get_signable_data()?;
344
345        // Verify using the public key directly
346        public_key
347            .verify(&signable_data, signature)
348            .map_err(|_| MessageError::VerificationFailed)
349            .map(|_| true)
350    }
351
352    /// Serialize message to bytes
353    pub fn to_bytes(&self) -> Result<Vec<u8>, MessageError> {
354        bincode::serialize(self).map_err(|_| MessageError::SerializationFailed)
355    }
356
357    /// Deserialize message from bytes
358    pub fn from_bytes(data: &[u8]) -> Result<Self, MessageError> {
359        bincode::deserialize(data).map_err(|_| MessageError::DeserializationFailed)
360    }
361
362    /// Validate message structure and content
363    pub fn validate(&self) -> Result<(), MessageError> {
364        // Check if message has expired
365        if self.is_expired() {
366            return Err(MessageError::MessageExpired);
367        }
368
369        // Check payload size limits (max 1MB)
370        if self.payload.len() > 1024 * 1024 {
371            return Err(MessageError::MessageTooLarge(self.payload.len()));
372        }
373
374        // Validate timestamp is reasonable (not too far in future)
375        let now = std::time::SystemTime::now()
376            .duration_since(std::time::UNIX_EPOCH)
377            .unwrap()
378            .as_millis() as u64;
379
380        // Allow up to 5 minutes in the future
381        if self.timestamp > now + (5 * 60 * 1000) {
382            return Err(MessageError::InvalidTimestamp);
383        }
384
385        Ok(())
386    }
387}
388
389/// Encrypted message container
390#[derive(Debug, Clone, Serialize, Deserialize)]
391pub struct EncryptedMessage {
392    /// Encrypted message data
393    pub ciphertext: Vec<u8>,
394    /// Encapsulated key (ML-KEM)
395    pub encapsulation: Vec<u8>,
396    /// Message headers (unencrypted)
397    pub headers: HashMap<String, String>,
398    /// Timestamp
399    pub timestamp: u64,
400}
401
402impl EncryptedMessage {
403    /// Encrypt a message using ML-KEM + AES-GCM
404    pub fn encrypt(
405        message: &Message,
406        recipient_public_key: &PublicKey,
407    ) -> Result<Self, MessageError> {
408        // Serialize the message
409        let message_bytes = message.to_bytes()?;
410
411        // Use ML-KEM for key encapsulation
412        let (ciphertext, shared_secret) = MlKem768::encapsulate(recipient_public_key)
413            .map_err(|_| MessageError::EncryptionFailed)?;
414
415        // Use shared secret as AES key (first 32 bytes)
416        let _aes_key = &shared_secret.as_bytes()[..32];
417
418        // Encrypt message with AES-GCM (simplified - in real implementation use proper AEAD)
419        let encrypted_data = message_bytes; // TODO: Implement actual AES-GCM encryption
420
421        Ok(Self {
422            ciphertext: encrypted_data,
423            encapsulation: ciphertext.as_bytes().to_vec(),
424            headers: message.headers.clone(),
425            timestamp: message.timestamp,
426        })
427    }
428
429    /// Decrypt a message using ML-KEM + AES-GCM
430    pub fn decrypt(&self, recipient_secret_key: &SecretKey) -> Result<Message, MessageError> {
431        // Decapsulate the shared secret
432        let encapsulation = Ciphertext::from_bytes(&self.encapsulation)
433            .map_err(|_| MessageError::DecryptionFailed)?;
434        let shared_secret = MlKem768::decapsulate(recipient_secret_key, &encapsulation)
435            .map_err(|_| MessageError::DecryptionFailed)?;
436
437        // Use shared secret as AES key
438        let _aes_key = &shared_secret.as_bytes()[..32];
439
440        // Decrypt message with AES-GCM (simplified)
441        let message_bytes = &self.ciphertext; // TODO: Implement actual AES-GCM decryption
442
443        Message::from_bytes(message_bytes)
444    }
445}
446
447/// Message factory for creating different types of protocol messages
448pub struct MessageFactory;
449
450impl MessageFactory {
451    /// Create a handshake init message
452    pub fn create_handshake_init(
453        protocol_version: ProtocolVersion,
454        public_key: &MlDsaPublicKey,
455        kem_public_key: &PublicKey,
456    ) -> Result<Message, MessageError> {
457        let payload = HandshakePayload {
458            protocol_version: protocol_version.clone(),
459            public_key: public_key.as_bytes().to_vec(),
460            kem_public_key: kem_public_key.as_bytes().to_vec(),
461            capabilities: vec!["anonymous-routing".to_string(), "dag-consensus".to_string()],
462            timestamp: std::time::SystemTime::now()
463                .duration_since(std::time::UNIX_EPOCH)
464                .unwrap()
465                .as_millis() as u64,
466        };
467
468        let payload_bytes =
469            bincode::serialize(&payload).map_err(|_| MessageError::SerializationFailed)?;
470
471        Ok(Message::new_with_version(
472            protocol_version,
473            MessageType::Handshake(HandshakeType::Init),
474            payload_bytes,
475        ))
476    }
477
478    /// Create a consensus vertex proposal message
479    pub fn create_vertex_proposal(
480        vertex_id: VertexId,
481        vertex_data: Vec<u8>,
482        parent_vertices: Vec<VertexId>,
483    ) -> Result<Message, MessageError> {
484        let payload = ConsensusPayload::VertexProposal {
485            vertex_id,
486            vertex_data,
487            parent_vertices,
488            timestamp: std::time::SystemTime::now()
489                .duration_since(std::time::UNIX_EPOCH)
490                .unwrap()
491                .as_millis() as u64,
492        };
493
494        let payload_bytes =
495            bincode::serialize(&payload).map_err(|_| MessageError::SerializationFailed)?;
496
497        Ok(Message::new(
498            MessageType::Consensus(ConsensusMessageType::VertexProposal),
499            payload_bytes,
500        ))
501    }
502
503    /// Create a ping message
504    pub fn create_ping() -> Result<Message, MessageError> {
505        let payload = ControlPayload::Ping {
506            timestamp: std::time::SystemTime::now()
507                .duration_since(std::time::UNIX_EPOCH)
508                .unwrap()
509                .as_millis() as u64,
510            nonce: rand::random::<u64>(),
511        };
512
513        let payload_bytes =
514            bincode::serialize(&payload).map_err(|_| MessageError::SerializationFailed)?;
515
516        Ok(Message::new(
517            MessageType::Control(ControlMessageType::Ping),
518            payload_bytes,
519        )
520        .with_ttl(30000)) // 30 second TTL
521    }
522}
523
524/// Handshake payload structure
525#[derive(Debug, Clone, Serialize, Deserialize)]
526pub struct HandshakePayload {
527    pub protocol_version: ProtocolVersion,
528    pub public_key: Vec<u8>,
529    pub kem_public_key: Vec<u8>,
530    pub capabilities: Vec<String>,
531    pub timestamp: u64,
532}
533
534/// Consensus message payloads
535#[derive(Debug, Clone, Serialize, Deserialize)]
536pub enum ConsensusPayload {
537    VertexProposal {
538        vertex_id: VertexId,
539        vertex_data: Vec<u8>,
540        parent_vertices: Vec<VertexId>,
541        timestamp: u64,
542    },
543    Vote {
544        vertex_id: VertexId,
545        vote: bool,
546        timestamp: u64,
547    },
548    Finality {
549        vertex_ids: Vec<VertexId>,
550        timestamp: u64,
551    },
552    Query {
553        requested_vertices: Vec<VertexId>,
554        timestamp: u64,
555    },
556}
557
558/// Control message payloads
559#[derive(Debug, Clone, Serialize, Deserialize)]
560pub enum ControlPayload {
561    Ping { timestamp: u64, nonce: u64 },
562    Pong { timestamp: u64, nonce: u64 },
563    Disconnect { reason: String, timestamp: u64 },
564    KeepAlive { timestamp: u64 },
565}
qudag_protocol/message.rs

qudag_protocol/
message.rs
