qudag_protocol/
handshake.rs

1//! Quantum-resistant handshake procedures for QuDAG protocol.
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::time::{Duration, SystemTime, UNIX_EPOCH};
6use thiserror::Error;
7use tracing::{debug, error, info};
8use uuid::Uuid;
9
10use qudag_crypto::{
11    Ciphertext as KemCiphertext, KeyPair as KemKeyPair, MlDsaKeyPair, MlDsaPublicKey, MlKem768,
12    PublicKey as KemPublicKey, SecretKey, SharedSecret,
13};
14use rand;
15
16use crate::message::{HandshakeType, Message, MessageError, MessageType, ProtocolVersion};
17use crate::state::{ProtocolStateMachine, StateError};
18
19/// Handshake-related errors
20#[derive(Debug, Error)]
21pub enum HandshakeError {
22    /// Cryptographic operation failed
23    #[error("Cryptographic operation failed: {reason}")]
24    CryptoError { reason: String },
25
26    /// Invalid handshake message
27    #[error("Invalid handshake message: {reason}")]
28    InvalidMessage { reason: String },
29
30    /// Handshake timeout
31    #[error("Handshake timed out after {timeout:?}")]
32    Timeout { timeout: Duration },
33
34    /// Protocol version mismatch
35    #[error("Protocol version mismatch: expected {expected:?}, got {actual:?}")]
36    VersionMismatch {
37        expected: ProtocolVersion,
38        actual: ProtocolVersion,
39    },
40
41    /// Unsupported capabilities
42    #[error("Unsupported capabilities: {capabilities:?}")]
43    UnsupportedCapabilities { capabilities: Vec<String> },
44
45    /// Invalid peer credentials
46    #[error("Invalid peer credentials")]
47    InvalidCredentials,
48
49    /// State machine error
50    #[error("State machine error: {0}")]
51    StateMachine(#[from] StateError),
52
53    /// Message error
54    #[error("Message error: {0}")]
55    Message(#[from] MessageError),
56
57    /// Handshake already in progress
58    #[error("Handshake already in progress with session {session_id}")]
59    HandshakeInProgress { session_id: Uuid },
60
61    /// Replay attack detected
62    #[error("Replay attack detected: timestamp {timestamp} is too old")]
63    ReplayAttack { timestamp: u64 },
64}
65
66/// Handshake configuration
67#[derive(Debug, Clone)]
68pub struct HandshakeConfig {
69    /// Timeout for handshake completion
70    pub timeout: Duration,
71    /// Supported protocol versions
72    pub supported_versions: Vec<ProtocolVersion>,
73    /// Required capabilities
74    pub required_capabilities: Vec<String>,
75    /// Optional capabilities
76    pub optional_capabilities: Vec<String>,
77    /// Maximum timestamp skew allowed (to prevent replay attacks)
78    pub max_timestamp_skew: Duration,
79    /// Enable mutual authentication
80    pub mutual_auth: bool,
81}
82
83impl Default for HandshakeConfig {
84    fn default() -> Self {
85        Self {
86            timeout: Duration::from_secs(30),
87            supported_versions: vec![ProtocolVersion::CURRENT],
88            required_capabilities: vec![
89                "dag-consensus".to_string(),
90                "quantum-resistant-crypto".to_string(),
91            ],
92            optional_capabilities: vec![
93                "anonymous-routing".to_string(),
94                "dark-addressing".to_string(),
95            ],
96            max_timestamp_skew: Duration::from_secs(300), // 5 minutes
97            mutual_auth: true,
98        }
99    }
100}
101
102/// Handshake session state
103#[derive(Debug)]
104pub struct HandshakeSession {
105    /// Session identifier
106    pub session_id: Uuid,
107    /// Peer identifier (if known)
108    pub peer_id: Option<Vec<u8>>,
109    /// Handshake state
110    pub state: HandshakeSessionState,
111    /// Protocol version negotiated
112    pub negotiated_version: Option<ProtocolVersion>,
113    /// Peer capabilities
114    pub peer_capabilities: Vec<String>,
115    /// Our ephemeral keys for this session
116    pub our_keys: HandshakeKeys,
117    /// Peer's public keys
118    pub peer_keys: Option<PeerKeys>,
119    /// Shared secrets
120    pub shared_secrets: Option<SharedSecrets>,
121    /// Session start time
122    pub started_at: SystemTime,
123    /// Last activity
124    pub last_activity: SystemTime,
125    /// Nonce for this session
126    pub nonce: u64,
127}
128
129/// Handshake session state
130#[derive(Debug, Clone, PartialEq, Eq)]
131pub enum HandshakeSessionState {
132    /// Waiting to start handshake
133    Waiting,
134    /// Sent initial handshake message
135    InitSent,
136    /// Received initial handshake message
137    InitReceived,
138    /// Sent handshake response
139    ResponseSent,
140    /// Received handshake response
141    ResponseReceived,
142    /// Handshake completed successfully
143    Completed,
144    /// Handshake failed
145    Failed,
146}
147
148/// Handshake keys for a session
149#[derive(Debug)]
150pub struct HandshakeKeys {
151    /// ML-DSA keypair for authentication
152    pub signature_keypair: MlDsaKeyPair,
153    /// ML-KEM keypair for key exchange
154    pub kem_keypair: KemKeyPair,
155}
156
157/// Peer's public keys
158#[derive(Debug, Clone)]
159pub struct PeerKeys {
160    /// Peer's ML-DSA public key
161    pub signature_public_key: MlDsaPublicKey,
162    /// Peer's ML-KEM public key
163    pub kem_public_key: KemPublicKey,
164}
165
166/// Derived shared secrets
167#[derive(Debug, Clone)]
168pub struct SharedSecrets {
169    /// Shared secret from ML-KEM
170    pub kem_shared_secret: SharedSecret,
171    /// Derived encryption key
172    pub encryption_key: Vec<u8>,
173    /// Derived MAC key
174    pub mac_key: Vec<u8>,
175    /// Session identifier
176    pub session_key: Vec<u8>,
177}
178
179/// Handshake message payloads
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub enum HandshakeMessagePayload {
182    /// Initial handshake message
183    Init {
184        protocol_version: ProtocolVersion,
185        supported_versions: Vec<ProtocolVersion>,
186        capabilities: Vec<String>,
187        signature_public_key: Vec<u8>,
188        kem_public_key: Vec<u8>,
189        nonce: u64,
190        timestamp: u64,
191    },
192    /// Handshake response message
193    Response {
194        protocol_version: ProtocolVersion,
195        capabilities: Vec<String>,
196        signature_public_key: Vec<u8>,
197        kem_ciphertext: Vec<u8>, // Encapsulated shared secret
198        nonce: u64,
199        timestamp: u64,
200    },
201    /// Handshake completion message
202    Complete { session_id: Vec<u8>, timestamp: u64 },
203    /// Version negotiation message
204    VersionNegotiation {
205        supported_versions: Vec<ProtocolVersion>,
206        preferred_version: ProtocolVersion,
207    },
208}
209
210/// Handshake coordinator
211pub struct HandshakeCoordinator {
212    /// Handshake configuration
213    config: HandshakeConfig,
214    /// Active handshake sessions
215    sessions: HashMap<Uuid, HandshakeSession>,
216    /// Our long-term identity keys
217    #[allow(dead_code)]
218    identity_keys: HandshakeKeys,
219    /// Protocol state machine
220    state_machine: ProtocolStateMachine,
221}
222
223impl HandshakeCoordinator {
224    /// Create a new handshake coordinator
225    pub fn new(
226        config: HandshakeConfig,
227        identity_keys: HandshakeKeys,
228        state_machine: ProtocolStateMachine,
229    ) -> Self {
230        Self {
231            config,
232            sessions: HashMap::new(),
233            identity_keys,
234            state_machine,
235        }
236    }
237
238    /// Generate new handshake keys
239    pub fn generate_keys() -> Result<HandshakeKeys, HandshakeError> {
240        // Generate ML-DSA keypair
241        let signature_keypair = MlDsaKeyPair::generate(&mut rand::thread_rng()).map_err(|e| {
242            HandshakeError::CryptoError {
243                reason: format!("Failed to generate ML-DSA keypair: {:?}", e),
244            }
245        })?;
246
247        // Generate ML-KEM keypair
248        let (kem_public_key, kem_secret_key) =
249            MlKem768::keygen().map_err(|e| HandshakeError::CryptoError {
250                reason: format!("Failed to generate ML-KEM keypair: {:?}", e),
251            })?;
252        let kem_keypair = KemKeyPair {
253            public_key: kem_public_key.as_bytes().to_vec(),
254            secret_key: kem_secret_key.as_bytes().to_vec(),
255        };
256
257        Ok(HandshakeKeys {
258            signature_keypair,
259            kem_keypair,
260        })
261    }
262
263    /// Initiate handshake with a peer
264    pub fn initiate_handshake(
265        &mut self,
266        peer_id: Option<Vec<u8>>,
267    ) -> Result<(Uuid, Message), HandshakeError> {
268        // Generate ephemeral keys for this session
269        let session_keys = Self::generate_keys()?;
270        let session_id = Uuid::new_v4();
271        let nonce = rand::random::<u64>();
272
273        // Create handshake session
274        let session = HandshakeSession {
275            session_id,
276            peer_id,
277            state: HandshakeSessionState::Waiting,
278            negotiated_version: None,
279            peer_capabilities: Vec::new(),
280            our_keys: session_keys,
281            peer_keys: None,
282            shared_secrets: None,
283            started_at: SystemTime::now(),
284            last_activity: SystemTime::now(),
285            nonce,
286        };
287
288        // Create initial handshake message
289        let payload = HandshakeMessagePayload::Init {
290            protocol_version: ProtocolVersion::CURRENT,
291            supported_versions: self.config.supported_versions.clone(),
292            capabilities: [
293                self.config.required_capabilities.clone(),
294                self.config.optional_capabilities.clone(),
295            ]
296            .concat(),
297            signature_public_key: session.our_keys.signature_keypair.public_key().to_vec(),
298            kem_public_key: session.our_keys.kem_keypair.public_key().to_vec(),
299            nonce,
300            timestamp: SystemTime::now()
301                .duration_since(UNIX_EPOCH)
302                .unwrap()
303                .as_millis() as u64,
304        };
305
306        let payload_bytes =
307            bincode::serialize(&payload).map_err(|e| HandshakeError::InvalidMessage {
308                reason: format!("Failed to serialize handshake init: {:?}", e),
309            })?;
310
311        let mut message = Message::new(MessageType::Handshake(HandshakeType::Init), payload_bytes);
312
313        // Sign the message
314        message.sign(&session.our_keys.signature_keypair)?;
315
316        // Update session state
317        let mut session = session;
318        session.state = HandshakeSessionState::InitSent;
319        session.last_activity = SystemTime::now();
320
321        // Store session
322        self.sessions.insert(session_id, session);
323
324        info!("Initiated handshake session: {}", session_id);
325        Ok((session_id, message))
326    }
327
328    /// Process incoming handshake message
329    pub fn process_handshake_message(
330        &mut self,
331        message: &Message,
332        session_id: Option<Uuid>,
333    ) -> Result<Option<Message>, HandshakeError> {
334        // Validate message timestamp to prevent replay attacks
335        self.validate_timestamp(message)?;
336
337        match &message.msg_type {
338            MessageType::Handshake(HandshakeType::Init) => self.process_handshake_init(message),
339            MessageType::Handshake(HandshakeType::Response) => {
340                self.process_handshake_response(message, session_id)
341            }
342            MessageType::Handshake(HandshakeType::Complete) => {
343                self.process_handshake_complete(message, session_id)
344            }
345            MessageType::Handshake(HandshakeType::VersionNegotiation) => {
346                self.process_version_negotiation(message, session_id)
347            }
348            _ => Err(HandshakeError::InvalidMessage {
349                reason: "Not a handshake message".to_string(),
350            }),
351        }
352    }
353
354    /// Process handshake init message
355    fn process_handshake_init(
356        &mut self,
357        message: &Message,
358    ) -> Result<Option<Message>, HandshakeError> {
359        let payload: HandshakeMessagePayload =
360            bincode::deserialize(&message.payload).map_err(|e| HandshakeError::InvalidMessage {
361                reason: format!("Failed to deserialize handshake init: {:?}", e),
362            })?;
363
364        if let HandshakeMessagePayload::Init {
365            protocol_version,
366            supported_versions,
367            capabilities,
368            signature_public_key,
369            kem_public_key,
370            nonce,
371            timestamp: _,
372        } = payload
373        {
374            // Verify protocol version compatibility
375            let negotiated_version =
376                self.negotiate_version(&supported_versions, &protocol_version)?;
377
378            // Verify required capabilities
379            self.verify_capabilities(&capabilities)?;
380
381            // Parse peer keys
382            let peer_signature_key =
383                MlDsaPublicKey::from_bytes(&signature_public_key).map_err(|e| {
384                    HandshakeError::CryptoError {
385                        reason: format!("Invalid peer signature key: {:?}", e),
386                    }
387                })?;
388
389            let peer_kem_key = KemPublicKey::from_bytes(&kem_public_key).map_err(|e| {
390                HandshakeError::CryptoError {
391                    reason: format!("Invalid peer KEM key: {:?}", e),
392                }
393            })?;
394
395            // Verify message signature
396            if !message.verify(&peer_signature_key)? {
397                return Err(HandshakeError::InvalidCredentials);
398            }
399
400            // Generate ephemeral keys for this session
401            let session_keys = Self::generate_keys()?;
402            let session_id = Uuid::new_v4();
403            let our_nonce = rand::random::<u64>();
404
405            // Perform key exchange
406            let (kem_ciphertext, shared_secret) =
407                MlKem768::encapsulate(&peer_kem_key).map_err(|e| HandshakeError::CryptoError {
408                    reason: format!("KEM encapsulation failed: {:?}", e),
409                })?;
410
411            // Derive session keys
412            let shared_secrets = self.derive_session_keys(&shared_secret, nonce, our_nonce)?;
413
414            // Create handshake session
415            let session = HandshakeSession {
416                session_id,
417                peer_id: None, // Will be set later if needed
418                state: HandshakeSessionState::InitReceived,
419                negotiated_version: Some(negotiated_version.clone()),
420                peer_capabilities: capabilities,
421                our_keys: session_keys,
422                peer_keys: Some(PeerKeys {
423                    signature_public_key: peer_signature_key,
424                    kem_public_key: peer_kem_key,
425                }),
426                shared_secrets: Some(shared_secrets),
427                started_at: SystemTime::now(),
428                last_activity: SystemTime::now(),
429                nonce: our_nonce,
430            };
431
432            // Create response message
433            let response_payload = HandshakeMessagePayload::Response {
434                protocol_version: negotiated_version,
435                capabilities: [
436                    self.config.required_capabilities.clone(),
437                    self.config.optional_capabilities.clone(),
438                ]
439                .concat(),
440                signature_public_key: session.our_keys.signature_keypair.public_key().to_vec(),
441                kem_ciphertext: kem_ciphertext.as_bytes().to_vec(),
442                nonce: our_nonce,
443                timestamp: SystemTime::now()
444                    .duration_since(UNIX_EPOCH)
445                    .unwrap()
446                    .as_millis() as u64,
447            };
448
449            let response_bytes = bincode::serialize(&response_payload).map_err(|e| {
450                HandshakeError::InvalidMessage {
451                    reason: format!("Failed to serialize handshake response: {:?}", e),
452                }
453            })?;
454
455            let mut response_message = Message::new(
456                MessageType::Handshake(HandshakeType::Response),
457                response_bytes,
458            );
459
460            // Sign the response
461            response_message.sign(&session.our_keys.signature_keypair)?;
462
463            // Update session state
464            let mut session = session;
465            session.state = HandshakeSessionState::ResponseSent;
466            session.last_activity = SystemTime::now();
467
468            // Store session
469            self.sessions.insert(session_id, session);
470
471            // Update protocol state machine
472            self.state_machine
473                .process_message(message, Some(session_id))?;
474
475            info!(
476                "Processed handshake init, sending response for session: {}",
477                session_id
478            );
479            Ok(Some(response_message))
480        } else {
481            Err(HandshakeError::InvalidMessage {
482                reason: "Expected handshake init payload".to_string(),
483            })
484        }
485    }
486
487    /// Process handshake response message
488    fn process_handshake_response(
489        &mut self,
490        message: &Message,
491        session_id: Option<Uuid>,
492    ) -> Result<Option<Message>, HandshakeError> {
493        let session_id = session_id.ok_or(HandshakeError::InvalidMessage {
494            reason: "Session ID required for handshake response".to_string(),
495        })?;
496
497        let session = self
498            .sessions
499            .get_mut(&session_id)
500            .ok_or(HandshakeError::InvalidMessage {
501                reason: format!("Session not found: {}", session_id),
502            })?;
503
504        if session.state != HandshakeSessionState::InitSent {
505            return Err(HandshakeError::InvalidMessage {
506                reason: format!("Invalid session state for response: {:?}", session.state),
507            });
508        }
509
510        let payload: HandshakeMessagePayload =
511            bincode::deserialize(&message.payload).map_err(|e| HandshakeError::InvalidMessage {
512                reason: format!("Failed to deserialize handshake response: {:?}", e),
513            })?;
514
515        if let HandshakeMessagePayload::Response {
516            protocol_version,
517            capabilities,
518            signature_public_key,
519            kem_ciphertext,
520            nonce,
521            timestamp: _,
522        } = payload
523        {
524            // Parse peer keys
525            let peer_signature_key =
526                MlDsaPublicKey::from_bytes(&signature_public_key).map_err(|e| {
527                    HandshakeError::CryptoError {
528                        reason: format!("Invalid peer signature key: {:?}", e),
529                    }
530                })?;
531
532            // Verify message signature
533            if !message.verify(&peer_signature_key)? {
534                return Err(HandshakeError::InvalidCredentials);
535            }
536
537            // Verify protocol version
538            if !self.config.supported_versions.contains(&protocol_version) {
539                return Err(HandshakeError::VersionMismatch {
540                    expected: ProtocolVersion::CURRENT,
541                    actual: protocol_version,
542                });
543            }
544
545            // Verify capabilities first
546            {
547                // Check required capabilities
548                for required_cap in &self.config.required_capabilities {
549                    if !capabilities.contains(required_cap) {
550                        return Err(HandshakeError::UnsupportedCapabilities {
551                            capabilities: vec![required_cap.clone()],
552                        });
553                    }
554                }
555            }
556
557            // Decapsulate shared secret
558            let kem_ciphertext_bytes = KemCiphertext::from_bytes(&kem_ciphertext).map_err(|e| {
559                HandshakeError::CryptoError {
560                    reason: format!("Invalid KEM ciphertext: {:?}", e),
561                }
562            })?;
563
564            let secret_key = SecretKey::from_bytes(session.our_keys.kem_keypair.secret_key())
565                .map_err(|e| HandshakeError::CryptoError {
566                    reason: format!("Invalid secret key: {:?}", e),
567                })?;
568
569            let shared_secret =
570                MlKem768::decapsulate(&secret_key, &kem_ciphertext_bytes).map_err(|e| {
571                    HandshakeError::CryptoError {
572                        reason: format!("KEM decapsulation failed: {:?}", e),
573                    }
574                })?;
575
576            // Derive session keys
577            let session_nonce = session.nonce;
578            let shared_secrets = {
579                let secret_bytes = shared_secret.as_bytes();
580
581                // Combine nonces for key derivation
582                let combined_nonce = session_nonce ^ nonce;
583                let nonce_bytes = combined_nonce.to_be_bytes();
584
585                // Use HKDF for key derivation (simplified)
586                let mut key_material = Vec::new();
587                key_material.extend_from_slice(secret_bytes);
588                key_material.extend_from_slice(&nonce_bytes);
589
590                // Derive different keys using BLAKE3 with different contexts
591                let encryption_key = blake3::keyed_hash(
592                    blake3::hash(b"QuDAG-Encryption-Key").as_bytes(),
593                    &key_material,
594                )
595                .as_bytes()
596                .to_vec();
597
598                let mac_key =
599                    blake3::keyed_hash(blake3::hash(b"QuDAG-MAC-Key").as_bytes(), &key_material)
600                        .as_bytes()
601                        .to_vec();
602
603                let session_key = blake3::keyed_hash(
604                    blake3::hash(b"QuDAG-Session-Key").as_bytes(),
605                    &key_material,
606                )
607                .as_bytes()
608                .to_vec();
609
610                SharedSecrets {
611                    kem_shared_secret: shared_secret.clone(),
612                    encryption_key,
613                    mac_key,
614                    session_key,
615                }
616            };
617
618            // Update session
619            session.negotiated_version = Some(protocol_version);
620            session.peer_capabilities = capabilities;
621            session.peer_keys = Some(PeerKeys {
622                signature_public_key: peer_signature_key,
623                kem_public_key: KemPublicKey::from_bytes(session.our_keys.kem_keypair.public_key())
624                    .map_err(|e| HandshakeError::CryptoError {
625                        reason: format!("Invalid public key: {:?}", e),
626                    })?, // Placeholder
627            });
628            session.shared_secrets = Some(shared_secrets);
629            session.state = HandshakeSessionState::ResponseReceived;
630            session.last_activity = SystemTime::now();
631
632            // Create completion message
633            let complete_payload = HandshakeMessagePayload::Complete {
634                session_id: session_id.as_bytes().to_vec(),
635                timestamp: SystemTime::now()
636                    .duration_since(UNIX_EPOCH)
637                    .unwrap()
638                    .as_millis() as u64,
639            };
640
641            let complete_bytes = bincode::serialize(&complete_payload).map_err(|e| {
642                HandshakeError::InvalidMessage {
643                    reason: format!("Failed to serialize handshake complete: {:?}", e),
644                }
645            })?;
646
647            let mut complete_message = Message::new(
648                MessageType::Handshake(HandshakeType::Complete),
649                complete_bytes,
650            );
651
652            // Sign the completion message
653            complete_message.sign(&session.our_keys.signature_keypair)?;
654
655            // Update session state
656            session.state = HandshakeSessionState::Completed;
657
658            // Update protocol state machine
659            self.state_machine
660                .process_message(message, Some(session_id))?;
661
662            info!(
663                "Processed handshake response, sending completion for session: {}",
664                session_id
665            );
666            Ok(Some(complete_message))
667        } else {
668            Err(HandshakeError::InvalidMessage {
669                reason: "Expected handshake response payload".to_string(),
670            })
671        }
672    }
673
674    /// Process handshake complete message
675    fn process_handshake_complete(
676        &mut self,
677        message: &Message,
678        session_id: Option<Uuid>,
679    ) -> Result<Option<Message>, HandshakeError> {
680        let session_id = session_id.ok_or(HandshakeError::InvalidMessage {
681            reason: "Session ID required for handshake complete".to_string(),
682        })?;
683
684        let session = self
685            .sessions
686            .get_mut(&session_id)
687            .ok_or(HandshakeError::InvalidMessage {
688                reason: format!("Session not found: {}", session_id),
689            })?;
690
691        if session.state != HandshakeSessionState::ResponseSent {
692            return Err(HandshakeError::InvalidMessage {
693                reason: format!("Invalid session state for complete: {:?}", session.state),
694            });
695        }
696
697        // Verify message signature
698        if let Some(peer_keys) = &session.peer_keys {
699            if !message.verify(&peer_keys.signature_public_key)? {
700                return Err(HandshakeError::InvalidCredentials);
701            }
702        }
703
704        // Update session state
705        session.state = HandshakeSessionState::Completed;
706        session.last_activity = SystemTime::now();
707
708        // Update protocol state machine
709        self.state_machine
710            .process_message(message, Some(session_id))?;
711
712        info!(
713            "Handshake completed successfully for session: {}",
714            session_id
715        );
716        Ok(None) // No response needed
717    }
718
719    /// Process version negotiation message
720    fn process_version_negotiation(
721        &mut self,
722        message: &Message,
723        _session_id: Option<Uuid>,
724    ) -> Result<Option<Message>, HandshakeError> {
725        let payload: HandshakeMessagePayload =
726            bincode::deserialize(&message.payload).map_err(|e| HandshakeError::InvalidMessage {
727                reason: format!("Failed to deserialize version negotiation: {:?}", e),
728            })?;
729
730        if let HandshakeMessagePayload::VersionNegotiation {
731            supported_versions,
732            preferred_version,
733        } = payload
734        {
735            // Find compatible version
736            let compatible_version =
737                self.negotiate_version(&supported_versions, &preferred_version)?;
738
739            // Update our supported versions if needed
740            debug!("Negotiated protocol version: {:?}", compatible_version);
741
742            // No response needed for version negotiation
743            Ok(None)
744        } else {
745            Err(HandshakeError::InvalidMessage {
746                reason: "Expected version negotiation payload".to_string(),
747            })
748        }
749    }
750
751    /// Negotiate protocol version
752    fn negotiate_version(
753        &self,
754        peer_versions: &[ProtocolVersion],
755        peer_preferred: &ProtocolVersion,
756    ) -> Result<ProtocolVersion, HandshakeError> {
757        // Try peer's preferred version first
758        if self.config.supported_versions.contains(peer_preferred) {
759            return Ok(peer_preferred.clone());
760        }
761
762        // Find highest compatible version
763        for our_version in &self.config.supported_versions {
764            for peer_version in peer_versions {
765                if our_version.is_compatible(peer_version) {
766                    return Ok(our_version.clone());
767                }
768            }
769        }
770
771        Err(HandshakeError::VersionMismatch {
772            expected: ProtocolVersion::CURRENT,
773            actual: peer_preferred.clone(),
774        })
775    }
776
777    /// Verify peer capabilities
778    fn verify_capabilities(&self, peer_capabilities: &[String]) -> Result<(), HandshakeError> {
779        for required_cap in &self.config.required_capabilities {
780            if !peer_capabilities.contains(required_cap) {
781                return Err(HandshakeError::UnsupportedCapabilities {
782                    capabilities: vec![required_cap.clone()],
783                });
784            }
785        }
786        Ok(())
787    }
788
789    /// Derive session keys from shared secret
790    fn derive_session_keys(
791        &self,
792        shared_secret: &SharedSecret,
793        our_nonce: u64,
794        peer_nonce: u64,
795    ) -> Result<SharedSecrets, HandshakeError> {
796        let secret_bytes = shared_secret.as_bytes();
797
798        // Combine nonces for key derivation
799        let combined_nonce = our_nonce ^ peer_nonce;
800        let nonce_bytes = combined_nonce.to_be_bytes();
801
802        // Use HKDF for key derivation (simplified)
803        let mut key_material = Vec::new();
804        key_material.extend_from_slice(secret_bytes);
805        key_material.extend_from_slice(&nonce_bytes);
806
807        // Derive different keys using BLAKE3 with different contexts
808        let encryption_key = blake3::keyed_hash(
809            blake3::hash(b"QuDAG-Encryption-Key").as_bytes(),
810            &key_material,
811        )
812        .as_bytes()
813        .to_vec();
814
815        let mac_key = blake3::keyed_hash(blake3::hash(b"QuDAG-MAC-Key").as_bytes(), &key_material)
816            .as_bytes()
817            .to_vec();
818
819        let session_key =
820            blake3::keyed_hash(blake3::hash(b"QuDAG-Session-Key").as_bytes(), &key_material)
821                .as_bytes()
822                .to_vec();
823
824        Ok(SharedSecrets {
825            kem_shared_secret: shared_secret.clone(),
826            encryption_key,
827            mac_key,
828            session_key,
829        })
830    }
831
832    /// Validate message timestamp to prevent replay attacks
833    fn validate_timestamp(&self, message: &Message) -> Result<(), HandshakeError> {
834        let now = SystemTime::now()
835            .duration_since(UNIX_EPOCH)
836            .unwrap()
837            .as_millis() as u64;
838
839        let message_time = message.timestamp;
840        let max_skew = self.config.max_timestamp_skew.as_millis() as u64;
841
842        if now > message_time + max_skew || message_time > now + max_skew {
843            return Err(HandshakeError::ReplayAttack {
844                timestamp: message_time,
845            });
846        }
847
848        Ok(())
849    }
850
851    /// Get handshake session
852    pub fn get_session(&self, session_id: &Uuid) -> Option<&HandshakeSession> {
853        self.sessions.get(session_id)
854    }
855
856    /// Remove completed or failed sessions
857    pub fn cleanup_sessions(&mut self) {
858        let now = SystemTime::now();
859        let timeout = self.config.timeout;
860
861        self.sessions.retain(|_, session| {
862            let elapsed = now
863                .duration_since(session.started_at)
864                .unwrap_or(Duration::ZERO);
865
866            match session.state {
867                HandshakeSessionState::Completed | HandshakeSessionState::Failed => false,
868                _ => elapsed < timeout,
869            }
870        });
871    }
872
873    /// Get all active sessions
874    pub fn get_active_sessions(&self) -> Vec<&HandshakeSession> {
875        self.sessions
876            .values()
877            .filter(|s| {
878                !matches!(
879                    s.state,
880                    HandshakeSessionState::Completed | HandshakeSessionState::Failed
881                )
882            })
883            .collect()
884    }
885
886    /// Check if handshake is completed for a session
887    pub fn is_handshake_completed(&self, session_id: &Uuid) -> bool {
888        self.sessions
889            .get(session_id)
890            .map(|s| s.state == HandshakeSessionState::Completed)
891            .unwrap_or(false)
892    }
893
894    /// Get shared secrets for a completed session
895    pub fn get_shared_secrets(&self, session_id: &Uuid) -> Option<&SharedSecrets> {
896        self.sessions
897            .get(session_id)
898            .and_then(|s| s.shared_secrets.as_ref())
899    }
900}
qudag_protocol/handshake.rs

qudag_protocol/
handshake.rs
