1  //! Handshake state machine for Dead Drop protocol.
  2  //!
  3  //! This module provides a high-level wrapper around the Noise XX handshake,
  4  //! adding application-specific functionality like contact verification and
  5  //! state tracking for the BLE communication flow.
  6  //!
  7  //! # Overview
  8  //!
  9  //! The handshake establishes a secure channel between two devices:
 10  //!
 11  //! 1. **Discovery**: Devices detect each other via BLE rotating IDs
 12  //! 2. **Connection**: BLE GATT connection is established
 13  //! 3. **Handshake**: Noise XX handshake authenticates both parties
 14  //! 4. **Verification**: Remote identity is verified against contacts
 15  //! 5. **Exchange**: Secure channel is used for message exchange
 16  //!
 17  //! # Handshake Flow
 18  //!
 19  //! ```text
 20  //! Initiator                              Responder
 21  //!     |                                      |
 22  //!     |──────────── Message 1 ──────────────>|  (ephemeral key)
 23  //!     |                                      |
 24  //!     |<─────────── Message 2 ───────────────|  (ephemeral + static, encrypted)
 25  //!     |                                      |
 26  //!     |──────────── Message 3 ──────────────>|  (static, encrypted)
 27  //!     |                                      |
 28  //!     |        [Handshake Complete]          |
 29  //!     |    [Secure Channel Established]      |
 30  //!     |                                      |
 31  //! ```
 32  //!
 33  //! # States
 34  //!
 35  //! | State | Description | Next Action |
 36  //! |-------|-------------|-------------|
 37  //! | `AwaitingOutgoing` | Need to send a message | Call `generate()` |
 38  //! | `AwaitingSend` | Have message to send | Send via BLE |
 39  //! | `AwaitingIncoming` | Waiting for peer message | Wait for BLE data |
 40  //! | `Complete` | Handshake finished | Call `finalize()` |
 41  //! | `Failed` | Handshake failed | Handle error |
 42  //!
 43  //! # Example
 44  //!
 45  //! ```
 46  //! use dead_drop_core::crypto::keys::ExchangeKeyPair;
 47  //! use dead_drop_core::protocol::handshake::Handshake;
 48  //!
 49  //! let my_keys = ExchangeKeyPair::generate();
 50  //!
 51  //! // As initiator
 52  //! let mut handshake = Handshake::new_initiator(&my_keys).unwrap();
 53  //!
 54  //! // Generate first message
 55  //! let msg1 = handshake.generate().unwrap();
 56  //! // ... send msg1 over BLE ...
 57  //!
 58  //! // After receiving response
 59  //! // let response = handshake.receive(&msg2_from_peer).unwrap();
 60  //! ```
 61  
 62  use crate::crypto::keys::ExchangeKeyPair;
 63  use crate::crypto::noise::{HandshakeStatus, NoiseHandshake, NoiseTransport, Role};
 64  use crate::error::{DeadDropError, Result};
 65  
 66  // =============================================================================
 67  // HANDSHAKE STATE
 68  // =============================================================================
 69  
 70  /// State of the handshake protocol.
 71  ///
 72  /// Tracks the current state of the handshake for the application layer.
 73  /// This provides a more granular view than the underlying Noise states.
 74  #[derive(Debug, Clone, PartialEq, Eq)]
 75  pub enum HandshakeState {
 76      /// Initial state, not yet started.
 77      Initial,
 78  
 79      /// Initiator: Sent message 1, waiting for message 2.
 80      /// Pattern: → e
 81      SentEphemeral,
 82  
 83      /// Responder: Received message 1, ready to send message 2.
 84      /// Pattern: received e, ready to send e,ee,s,es
 85      ReceivedEphemeral,
 86  
 87      /// Responder: Sent message 2, waiting for message 3.
 88      /// Pattern: ← e, ee, s, es
 89      SentEphemeralAndStatic,
 90  
 91      /// Initiator: Received message 2, ready to send message 3.
 92      /// Pattern: received e,ee,s,es, ready to send s,se
 93      ReceivedEphemeralAndStatic,
 94  
 95      /// Initiator: Sent message 3, handshake complete.
 96      /// Pattern: → s, se
 97      SentStatic,
 98  
 99      /// Responder: Received message 3, handshake complete.
100      ReceivedStatic,
101  
102      /// Handshake completed successfully.
103      Complete,
104  
105      /// Handshake failed.
106      Failed(String),
107  }
108  
109  impl HandshakeState {
110      /// Check if the handshake is complete.
111      pub fn is_complete(&self) -> bool {
112          matches!(self, HandshakeState::Complete)
113      }
114  
115      /// Check if the handshake has failed.
116      pub fn is_failed(&self) -> bool {
117          matches!(self, HandshakeState::Failed(_))
118      }
119  
120      /// Check if we need to send a message.
121      pub fn needs_to_send(&self) -> bool {
122          matches!(
123              self,
124              HandshakeState::Initial
125                  | HandshakeState::ReceivedEphemeral
126                  | HandshakeState::ReceivedEphemeralAndStatic
127          )
128      }
129  
130      /// Check if we are waiting to receive a message.
131      pub fn needs_to_receive(&self) -> bool {
132          matches!(
133              self,
134              HandshakeState::SentEphemeral | HandshakeState::SentEphemeralAndStatic
135          )
136      }
137  }
138  
139  // =============================================================================
140  // HANDSHAKE RESULT
141  // =============================================================================
142  
143  /// Result of processing a handshake operation.
144  ///
145  /// Contains the outcome of a `receive()` or `generate()` call,
146  /// including any data to send and the new state.
147  #[derive(Debug)]
148  pub struct HandshakeResult {
149      /// Data to send to the peer (if any).
150      pub data_to_send: Option<Vec<u8>>,
151      /// Whether the handshake is now complete.
152      pub is_complete: bool,
153      /// The remote peer's static public key (if known).
154      pub remote_static: Option<[u8; 32]>,
155  }
156  
157  // =============================================================================
158  // HANDSHAKE
159  // =============================================================================
160  
161  /// High-level handshake handler for Dead Drop protocol.
162  ///
163  /// Wraps the Noise XX handshake with application-specific state tracking
164  /// and contact verification support.
165  ///
166  /// # Lifecycle
167  ///
168  /// 1. Create with `new_initiator()` or `new_responder()`
169  /// 2. Call `generate()` when `state().needs_to_send()`
170  /// 3. Call `receive()` when data arrives from peer
171  /// 4. Repeat until `state().is_complete()`
172  /// 5. Call `finalize()` to get the secure transport
173  ///
174  /// # Thread Safety
175  ///
176  /// This type is `Send` but not `Sync`. For concurrent access,
177  /// wrap in `Arc<Mutex<Handshake>>`.
178  pub struct Handshake {
179      /// The underlying Noise handshake.
180      noise: NoiseHandshake,
181      /// Application-level state tracking.
182      state: HandshakeState,
183      /// The role in the handshake.
184      role: Role,
185      /// Remote peer's static key (once known).
186      remote_static: Option<[u8; 32]>,
187      /// Count of messages sent.
188      messages_sent: u8,
189      /// Count of messages received.
190      messages_received: u8,
191  }
192  
193  impl Handshake {
194      /// Create a new handshake as the initiator.
195      ///
196      /// The initiator is the device that detected the other device's
197      /// rotating ID and is initiating the connection.
198      ///
199      /// # Arguments
200      ///
201      /// * `local_exchange` - Our exchange key pair for the handshake
202      ///
203      /// # Example
204      ///
205      /// ```
206      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
207      /// use dead_drop_core::protocol::handshake::Handshake;
208      ///
209      /// let keys = ExchangeKeyPair::generate();
210      /// let handshake = Handshake::new_initiator(&keys).unwrap();
211      ///
212      /// assert!(handshake.state().needs_to_send());
213      /// ```
214      pub fn new_initiator(local_exchange: &ExchangeKeyPair) -> Result<Self> {
215          let noise = NoiseHandshake::new_initiator(local_exchange)?;
216  
217          Ok(Self {
218              noise,
219              state: HandshakeState::Initial,
220              role: Role::Initiator,
221              remote_static: None,
222              messages_sent: 0,
223              messages_received: 0,
224          })
225      }
226  
227      /// Create a new handshake as the responder.
228      ///
229      /// The responder is the device that was discovered and is responding
230      /// to an incoming connection.
231      ///
232      /// # Arguments
233      ///
234      /// * `local_exchange` - Our exchange key pair for the handshake
235      ///
236      /// # Example
237      ///
238      /// ```
239      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
240      /// use dead_drop_core::protocol::handshake::Handshake;
241      ///
242      /// let keys = ExchangeKeyPair::generate();
243      /// let handshake = Handshake::new_responder(&keys).unwrap();
244      ///
245      /// assert!(handshake.state().needs_to_receive());
246      /// ```
247      pub fn new_responder(local_exchange: &ExchangeKeyPair) -> Result<Self> {
248          let noise = NoiseHandshake::new_responder(local_exchange)?;
249  
250          Ok(Self {
251              noise,
252              state: HandshakeState::SentEphemeral, // Waiting for first message
253              role: Role::Responder,
254              remote_static: None,
255              messages_sent: 0,
256              messages_received: 0,
257          })
258      }
259  
260      /// Get the current handshake state.
261      pub fn state(&self) -> &HandshakeState {
262          &self.state
263      }
264  
265      /// Get the handshake role.
266      pub fn role(&self) -> Role {
267          self.role
268      }
269  
270      /// Check if the handshake is complete.
271      pub fn is_complete(&self) -> bool {
272          self.state.is_complete()
273      }
274  
275      /// Check if the handshake has failed.
276      pub fn is_failed(&self) -> bool {
277          self.state.is_failed()
278      }
279  
280      /// Get the remote peer's static public key.
281      ///
282      /// Only available after the peer has sent their static key
283      /// (message 2 for initiator, message 3 for responder).
284      pub fn remote_static(&self) -> Option<&[u8; 32]> {
285          self.remote_static.as_ref()
286      }
287  
288      /// Get the number of messages sent.
289      pub fn messages_sent(&self) -> u8 {
290          self.messages_sent
291      }
292  
293      /// Get the number of messages received.
294      pub fn messages_received(&self) -> u8 {
295          self.messages_received
296      }
297  
298      /// Generate the next outgoing handshake message.
299      ///
300      /// Call this when `state().needs_to_send()` returns true.
301      ///
302      /// # Returns
303      ///
304      /// The handshake message to send to the peer.
305      ///
306      /// # Errors
307      ///
308      /// - `InvalidHandshakeState` if not in a sending state
309      /// - `Handshake` if message generation fails
310      ///
311      /// # Example
312      ///
313      /// ```
314      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
315      /// use dead_drop_core::protocol::handshake::Handshake;
316      ///
317      /// let keys = ExchangeKeyPair::generate();
318      /// let mut handshake = Handshake::new_initiator(&keys).unwrap();
319      ///
320      /// let msg = handshake.generate().unwrap();
321      /// assert!(!msg.is_empty());
322      /// ```
323      pub fn generate(&mut self) -> Result<Vec<u8>> {
324          if !self.state.needs_to_send() && self.state != HandshakeState::Initial {
325              // Special case: responder starts in "waiting" state
326              if self.role == Role::Responder && self.noise.status() == HandshakeStatus::NeedToReceive
327              {
328                  return Err(DeadDropError::InvalidHandshakeState(
329                      "Responder must receive first message before generating".to_string(),
330                  ));
331              }
332          }
333  
334          // Generate the message
335          let message = self.noise.write_message().map_err(|e| {
336              self.state = HandshakeState::Failed(e.to_string());
337              e
338          })?;
339  
340          self.messages_sent += 1;
341  
342          // Update state based on role and message count
343          self.update_state_after_send();
344  
345          // Check for remote static key
346          if let Some(rs) = self.noise.remote_static() {
347              self.remote_static = Some(*rs);
348          }
349  
350          Ok(message)
351      }
352  
353      /// Process an incoming handshake message from the peer.
354      ///
355      /// Call this when handshake data is received over BLE.
356      ///
357      /// # Arguments
358      ///
359      /// * `data` - The handshake message from the peer
360      ///
361      /// # Returns
362      ///
363      /// A `HandshakeResult` indicating what to do next.
364      ///
365      /// # Errors
366      ///
367      /// - `InvalidHandshakeState` if not expecting incoming data
368      /// - `Handshake` if message processing fails
369      ///
370      /// # Example
371      ///
372      /// ```
373      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
374      /// use dead_drop_core::protocol::handshake::Handshake;
375      ///
376      /// let alice_keys = ExchangeKeyPair::generate();
377      /// let bob_keys = ExchangeKeyPair::generate();
378      ///
379      /// let mut alice = Handshake::new_initiator(&alice_keys).unwrap();
380      /// let mut bob = Handshake::new_responder(&bob_keys).unwrap();
381      ///
382      /// // Alice sends first message
383      /// let msg1 = alice.generate().unwrap();
384      ///
385      /// // Bob receives and processes
386      /// let result = bob.receive(&msg1).unwrap();
387      /// assert!(result.data_to_send.is_none()); // Bob needs to call generate()
388      /// ```
389      pub fn receive(&mut self, data: &[u8]) -> Result<HandshakeResult> {
390          // For responder, the first receive is valid even in "initial" equivalent state
391          let is_responder_first = self.role == Role::Responder && self.messages_received == 0;
392  
393          if !self.state.needs_to_receive() && !is_responder_first && !self.state.needs_to_send() {
394              // Special case for responder waiting state
395              if !(self.role == Role::Responder
396                  && self.noise.status() == HandshakeStatus::NeedToReceive)
397              {
398                  return Err(DeadDropError::InvalidHandshakeState(format!(
399                      "Cannot receive in state {:?}",
400                      self.state
401                  )));
402              }
403          }
404  
405          // Process the message
406          let _payload = self.noise.read_message(data).map_err(|e| {
407              self.state = HandshakeState::Failed(e.to_string());
408              e
409          })?;
410  
411          self.messages_received += 1;
412  
413          // Check for remote static key
414          if let Some(rs) = self.noise.remote_static() {
415              self.remote_static = Some(*rs);
416          }
417  
418          // Update state based on role and message count
419          self.update_state_after_receive();
420  
421          Ok(HandshakeResult {
422              data_to_send: None, // Caller should call generate() if needed
423              is_complete: self.is_complete(),
424              remote_static: self.remote_static,
425          })
426      }
427  
428      /// Update state after sending a message.
429      fn update_state_after_send(&mut self) {
430          match self.role {
431              Role::Initiator => {
432                  match self.messages_sent {
433                      1 => self.state = HandshakeState::SentEphemeral,
434                      2 => {
435                          self.state = HandshakeState::SentStatic;
436                          if self.noise.is_complete() {
437                              self.state = HandshakeState::Complete;
438                          }
439                      }
440                      _ => {}
441                  }
442              }
443              Role::Responder => {
444                  if self.messages_sent == 1 {
445                      self.state = HandshakeState::SentEphemeralAndStatic;
446                  }
447              }
448          }
449  
450          // Check if noise thinks we're complete
451          if self.noise.is_complete() {
452              self.state = HandshakeState::Complete;
453          }
454      }
455  
456      /// Update state after receiving a message.
457      fn update_state_after_receive(&mut self) {
458          match self.role {
459              Role::Initiator => {
460                  if self.messages_received == 1 {
461                      self.state = HandshakeState::ReceivedEphemeralAndStatic;
462                  }
463              }
464              Role::Responder => {
465                  match self.messages_received {
466                      1 => self.state = HandshakeState::ReceivedEphemeral,
467                      2 => {
468                          self.state = HandshakeState::ReceivedStatic;
469                          if self.noise.is_complete() {
470                              self.state = HandshakeState::Complete;
471                          }
472                      }
473                      _ => {}
474                  }
475              }
476          }
477  
478          // Check if noise thinks we're complete
479          if self.noise.is_complete() {
480              self.state = HandshakeState::Complete;
481          }
482      }
483  
484      /// Finalize the handshake and get the secure transport.
485      ///
486      /// Consumes the handshake and returns a `NoiseTransport` for
487      /// encrypted communication.
488      ///
489      /// # Errors
490      ///
491      /// - `InvalidHandshakeState` if the handshake is not complete
492      ///
493      /// # Example
494      ///
495      /// ```
496      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
497      /// use dead_drop_core::crypto::noise::perform_handshake;
498      ///
499      /// let alice = ExchangeKeyPair::generate();
500      /// let bob = ExchangeKeyPair::generate();
501      ///
502      /// // Use the helper for complete handshake
503      /// let (alice_transport, bob_transport) = perform_handshake(&alice, &bob).unwrap();
504      ///
505      /// // Now both can send encrypted messages
506      /// ```
507      pub fn finalize(self) -> Result<NoiseTransport> {
508          if !self.is_complete() {
509              return Err(DeadDropError::InvalidHandshakeState(format!(
510                  "Cannot finalize in state {:?}",
511                  self.state
512              )));
513          }
514  
515          self.noise.into_transport()
516      }
517  
518      /// Verify that the remote peer is a known contact.
519      ///
520      /// Call this after receiving the peer's static key to verify they
521      /// are in your contact list.
522      ///
523      /// # Arguments
524      ///
525      /// * `known_contacts` - List of known contact exchange public keys
526      ///
527      /// # Returns
528      ///
529      /// `true` if the remote static key matches one of the known contacts.
530      ///
531      /// # Example
532      ///
533      /// ```
534      /// use dead_drop_core::crypto::keys::ExchangeKeyPair;
535      /// use dead_drop_core::protocol::handshake::Handshake;
536      ///
537      /// let my_keys = ExchangeKeyPair::generate();
538      /// let contact_keys = ExchangeKeyPair::generate();
539      ///
540      /// // After handshake completes...
541      /// // let is_known = handshake.verify_contact(&[contact_keys.public_bytes()]);
542      /// ```
543      pub fn verify_contact(&self, known_contacts: &[[u8; 32]]) -> bool {
544          match &self.remote_static {
545              Some(remote) => known_contacts.iter().any(|contact| contact == remote),
546              None => false,
547          }
548      }
549  }
550  
551  impl std::fmt::Debug for Handshake {
552      fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
553          f.debug_struct("Handshake")
554              .field("state", &self.state)
555              .field("role", &self.role)
556              .field("messages_sent", &self.messages_sent)
557              .field("messages_received", &self.messages_received)
558              .field("remote_static", &self.remote_static.map(|k| hex::encode(&k[..8])))
559              .finish()
560      }
561  }
562  
563  // =============================================================================
564  // CONVENIENCE FUNCTIONS
565  // =============================================================================
566  
567  /// Perform a complete handshake between two parties.
568  ///
569  /// This is a convenience function for testing. In production, handshake
570  /// messages are exchanged over BLE.
571  ///
572  /// # Arguments
573  ///
574  /// * `initiator_keys` - Initiator's exchange key pair
575  /// * `responder_keys` - Responder's exchange key pair
576  ///
577  /// # Returns
578  ///
579  /// A tuple of (initiator_transport, responder_transport).
580  pub fn perform_handshake(
581      initiator_keys: &ExchangeKeyPair,
582      responder_keys: &ExchangeKeyPair,
583  ) -> Result<(NoiseTransport, NoiseTransport)> {
584      let mut initiator = Handshake::new_initiator(initiator_keys)?;
585      let mut responder = Handshake::new_responder(responder_keys)?;
586  
587      // Message 1: Initiator → Responder
588      let msg1 = initiator.generate()?;
589      responder.receive(&msg1)?;
590  
591      // Message 2: Responder → Initiator
592      let msg2 = responder.generate()?;
593      initiator.receive(&msg2)?;
594  
595      // Message 3: Initiator → Responder
596      let msg3 = initiator.generate()?;
597      responder.receive(&msg3)?;
598  
599      // Both should be complete
600      assert!(initiator.is_complete(), "Initiator not complete");
601      assert!(responder.is_complete(), "Responder not complete");
602  
603      // Finalize into transports
604      let initiator_transport = initiator.finalize()?;
605      let responder_transport = responder.finalize()?;
606  
607      Ok((initiator_transport, responder_transport))
608  }
609  
610  #[cfg(test)]
611  mod tests {
612      use super::*;
613  
614      // ==================== Creation Tests ====================
615  
616      #[test]
617      fn test_create_initiator() {
618          let keys = ExchangeKeyPair::generate();
619          let handshake = Handshake::new_initiator(&keys).unwrap();
620  
621          assert_eq!(handshake.role(), Role::Initiator);
622          assert!(handshake.state().needs_to_send() || *handshake.state() == HandshakeState::Initial);
623          assert!(!handshake.is_complete());
624          assert!(!handshake.is_failed());
625          assert_eq!(handshake.messages_sent(), 0);
626          assert_eq!(handshake.messages_received(), 0);
627      }
628  
629      #[test]
630      fn test_create_responder() {
631          let keys = ExchangeKeyPair::generate();
632          let handshake = Handshake::new_responder(&keys).unwrap();
633  
634          assert_eq!(handshake.role(), Role::Responder);
635          // Responder waits for first message
636          assert!(!handshake.is_complete());
637          assert!(!handshake.is_failed());
638      }
639  
640      // ==================== Complete Handshake Tests ====================
641  
642      #[test]
643      fn test_complete_handshake_flow() {
644          let alice_keys = ExchangeKeyPair::generate();
645          let bob_keys = ExchangeKeyPair::generate();
646  
647          let mut alice = Handshake::new_initiator(&alice_keys).unwrap();
648          let mut bob = Handshake::new_responder(&bob_keys).unwrap();
649  
650          // Message 1: Alice → Bob (ephemeral)
651          let msg1 = alice.generate().unwrap();
652          assert_eq!(alice.messages_sent(), 1);
653          assert_eq!(*alice.state(), HandshakeState::SentEphemeral);
654  
655          bob.receive(&msg1).unwrap();
656          assert_eq!(bob.messages_received(), 1);
657          assert_eq!(*bob.state(), HandshakeState::ReceivedEphemeral);
658  
659          // Message 2: Bob → Alice (ephemeral + static)
660          let msg2 = bob.generate().unwrap();
661          assert_eq!(bob.messages_sent(), 1);
662          assert_eq!(*bob.state(), HandshakeState::SentEphemeralAndStatic);
663  
664          alice.receive(&msg2).unwrap();
665          assert_eq!(alice.messages_received(), 1);
666          assert_eq!(*alice.state(), HandshakeState::ReceivedEphemeralAndStatic);
667  
668          // Alice should now have Bob's static key
669          assert!(alice.remote_static().is_some());
670          assert_eq!(alice.remote_static().unwrap(), &bob_keys.public_bytes());
671  
672          // Message 3: Alice → Bob (static)
673          let msg3 = alice.generate().unwrap();
674          assert_eq!(alice.messages_sent(), 2);
675          assert!(alice.is_complete());
676  
677          bob.receive(&msg3).unwrap();
678          assert_eq!(bob.messages_received(), 2);
679          assert!(bob.is_complete());
680  
681          // Bob should now have Alice's static key
682          assert!(bob.remote_static().is_some());
683          assert_eq!(bob.remote_static().unwrap(), &alice_keys.public_bytes());
684      }
685  
686      #[test]
687      fn test_perform_handshake_helper() {
688          let alice = ExchangeKeyPair::generate();
689          let bob = ExchangeKeyPair::generate();
690  
691          let result = perform_handshake(&alice, &bob);
692          assert!(result.is_ok());
693  
694          let (alice_transport, bob_transport) = result.unwrap();
695  
696          // Verify they have each other's static keys
697          assert_eq!(alice_transport.remote_static().unwrap(), &bob.public_bytes());
698          assert_eq!(bob_transport.remote_static().unwrap(), &alice.public_bytes());
699      }
700  
701      // ==================== Transport Tests ====================
702  
703      #[test]
704      fn test_finalize_to_transport() {
705          let alice = ExchangeKeyPair::generate();
706          let bob = ExchangeKeyPair::generate();
707  
708          let (mut alice_transport, mut bob_transport) = perform_handshake(&alice, &bob).unwrap();
709  
710          // Test encryption/decryption
711          let plaintext = b"Hello from Alice!";
712          let ciphertext = alice_transport.encrypt(plaintext).unwrap();
713          let decrypted = bob_transport.decrypt(&ciphertext).unwrap();
714  
715          assert_eq!(plaintext.as_slice(), decrypted.as_slice());
716  
717          // Bidirectional
718          let reply = b"Hello from Bob!";
719          let ciphertext2 = bob_transport.encrypt(reply).unwrap();
720          let decrypted2 = alice_transport.decrypt(&ciphertext2).unwrap();
721  
722          assert_eq!(reply.as_slice(), decrypted2.as_slice());
723      }
724  
725      // ==================== Error Cases ====================
726  
727      #[test]
728      fn test_cannot_generate_when_not_ready() {
729          let keys = ExchangeKeyPair::generate();
730          let mut responder = Handshake::new_responder(&keys).unwrap();
731  
732          // Responder can't generate before receiving
733          let result = responder.generate();
734          assert!(result.is_err());
735      }
736  
737      #[test]
738      fn test_cannot_finalize_incomplete() {
739          let keys = ExchangeKeyPair::generate();
740          let handshake = Handshake::new_initiator(&keys).unwrap();
741  
742          let result = handshake.finalize();
743          assert!(result.is_err());
744          assert!(matches!(
745              result.unwrap_err(),
746              DeadDropError::InvalidHandshakeState(_)
747          ));
748      }
749  
750      #[test]
751      fn test_invalid_handshake_message() {
752          let alice_keys = ExchangeKeyPair::generate();
753          let bob_keys = ExchangeKeyPair::generate();
754  
755          let mut alice = Handshake::new_initiator(&alice_keys).unwrap();
756          let mut bob = Handshake::new_responder(&bob_keys).unwrap();
757  
758          // Get a valid first message
759          let msg1 = alice.generate().unwrap();
760          bob.receive(&msg1).unwrap();
761  
762          // Now Bob generates message 2
763          let msg2 = bob.generate().unwrap();
764  
765          // Tamper with the encrypted portion of message 2
766          let mut garbage = msg2.clone();
767          for byte in garbage.iter_mut().skip(32) {
768              // Skip the ephemeral key, tamper with encrypted payload
769              *byte ^= 0xFF;
770          }
771  
772          // Alice should reject the tampered message
773          let result = alice.receive(&garbage);
774          assert!(result.is_err(), "Should reject tampered handshake message");
775          assert!(alice.is_failed());
776      }
777  
778      // ==================== Contact Verification ====================
779  
780      #[test]
781      fn test_verify_contact_success() {
782          let alice = ExchangeKeyPair::generate();
783          let bob = ExchangeKeyPair::generate();
784  
785          let (alice_hs, bob_hs) = {
786              let mut alice_hs = Handshake::new_initiator(&alice).unwrap();
787              let mut bob_hs = Handshake::new_responder(&bob).unwrap();
788  
789              let msg1 = alice_hs.generate().unwrap();
790              bob_hs.receive(&msg1).unwrap();
791  
792              let msg2 = bob_hs.generate().unwrap();
793              alice_hs.receive(&msg2).unwrap();
794  
795              let msg3 = alice_hs.generate().unwrap();
796              bob_hs.receive(&msg3).unwrap();
797  
798              (alice_hs, bob_hs)
799          };
800  
801          // Alice verifies Bob is a known contact
802          let known_contacts = [bob.public_bytes()];
803          assert!(alice_hs.verify_contact(&known_contacts));
804  
805          // Bob verifies Alice is a known contact
806          let known_contacts = [alice.public_bytes()];
807          assert!(bob_hs.verify_contact(&known_contacts));
808      }
809  
810      #[test]
811      fn test_verify_contact_failure() {
812          let alice = ExchangeKeyPair::generate();
813          let bob = ExchangeKeyPair::generate();
814          let charlie = ExchangeKeyPair::generate();
815  
816          let (alice_hs, _) = {
817              let mut alice_hs = Handshake::new_initiator(&alice).unwrap();
818              let mut bob_hs = Handshake::new_responder(&bob).unwrap();
819  
820              let msg1 = alice_hs.generate().unwrap();
821              bob_hs.receive(&msg1).unwrap();
822  
823              let msg2 = bob_hs.generate().unwrap();
824              alice_hs.receive(&msg2).unwrap();
825  
826              let msg3 = alice_hs.generate().unwrap();
827              bob_hs.receive(&msg3).unwrap();
828  
829              (alice_hs, bob_hs)
830          };
831  
832          // Alice's contact list only has Charlie, not Bob
833          let known_contacts = [charlie.public_bytes()];
834          assert!(!alice_hs.verify_contact(&known_contacts));
835  
836          // Empty contact list
837          assert!(!alice_hs.verify_contact(&[]));
838      }
839  
840      #[test]
841      fn test_verify_contact_before_key_exchange() {
842          let keys = ExchangeKeyPair::generate();
843          let handshake = Handshake::new_initiator(&keys).unwrap();
844  
845          // Remote static not yet known
846          assert!(handshake.remote_static().is_none());
847          assert!(!handshake.verify_contact(&[keys.public_bytes()]));
848      }
849  
850      // ==================== State Tests ====================
851  
852      #[test]
853      fn test_handshake_state_needs_to_send() {
854          assert!(HandshakeState::Initial.needs_to_send());
855          assert!(HandshakeState::ReceivedEphemeral.needs_to_send());
856          assert!(HandshakeState::ReceivedEphemeralAndStatic.needs_to_send());
857          assert!(!HandshakeState::SentEphemeral.needs_to_send());
858          assert!(!HandshakeState::Complete.needs_to_send());
859      }
860  
861      #[test]
862      fn test_handshake_state_needs_to_receive() {
863          assert!(HandshakeState::SentEphemeral.needs_to_receive());
864          assert!(HandshakeState::SentEphemeralAndStatic.needs_to_receive());
865          assert!(!HandshakeState::Initial.needs_to_receive());
866          assert!(!HandshakeState::Complete.needs_to_receive());
867      }
868  
869      #[test]
870      fn test_handshake_state_is_terminal() {
871          assert!(HandshakeState::Complete.is_complete());
872          assert!(!HandshakeState::Initial.is_complete());
873  
874          assert!(HandshakeState::Failed("test".to_string()).is_failed());
875          assert!(!HandshakeState::Complete.is_failed());
876      }
877  
878      // ==================== Debug Tests ====================
879  
880      #[test]
881      fn test_handshake_debug() {
882          let keys = ExchangeKeyPair::generate();
883          let handshake = Handshake::new_initiator(&keys).unwrap();
884  
885          let debug_str = format!("{:?}", handshake);
886          assert!(debug_str.contains("Handshake"));
887          assert!(debug_str.contains("state"));
888          assert!(debug_str.contains("role"));
889      }
890  }