Skip to main content

fips_core/noise/
handshake.rs

1use super::{
2    CipherState, EPOCH_ENCRYPTED_SIZE, EPOCH_SIZE, HANDSHAKE_MSG1_SIZE, HANDSHAKE_MSG2_SIZE,
3    HandshakeProgress, HandshakeRole, NoiseError, NoisePattern, NoiseSession, PROTOCOL_NAME_IK,
4    PROTOCOL_NAME_XK, PUBKEY_SIZE, XK_HANDSHAKE_MSG1_SIZE, XK_HANDSHAKE_MSG2_SIZE,
5    XK_HANDSHAKE_MSG3_SIZE,
6};
7use hkdf::Hkdf;
8use rand::Rng;
9use secp256k1::{Keypair, PublicKey, Secp256k1, SecretKey, ecdh::shared_secret_point};
10use sha2::{Digest, Sha256};
11use std::fmt;
12
13/// Symmetric state during handshake.
14///
15/// Maintains the chaining key (ck), handshake hash (h), and current cipher.
16#[derive(Clone)]
17struct SymmetricState {
18    /// Chaining key for key derivation.
19    ck: [u8; 32],
20    /// Handshake hash for transcript binding.
21    h: [u8; 32],
22    /// Current cipher state for encrypting handshake payloads.
23    cipher: CipherState,
24}
25
26impl SymmetricState {
27    /// Initialize with protocol name.
28    fn initialize(protocol_name: &[u8]) -> Self {
29        // If protocol name <= 32 bytes, pad with zeros
30        // If > 32 bytes, hash it
31        let h = if protocol_name.len() <= 32 {
32            let mut h = [0u8; 32];
33            h[..protocol_name.len()].copy_from_slice(protocol_name);
34            h
35        } else {
36            let mut hasher = Sha256::new();
37            hasher.update(protocol_name);
38            hasher.finalize().into()
39        };
40
41        Self {
42            ck: h,
43            h,
44            cipher: CipherState::empty(),
45        }
46    }
47
48    /// Mix data into the handshake hash.
49    fn mix_hash(&mut self, data: &[u8]) {
50        let mut hasher = Sha256::new();
51        hasher.update(self.h);
52        hasher.update(data);
53        self.h = hasher.finalize().into();
54    }
55
56    /// Mix key material into the chaining key.
57    fn mix_key(&mut self, input_key_material: &[u8]) {
58        let hk = Hkdf::<Sha256>::new(Some(&self.ck), input_key_material);
59        let mut output = [0u8; 64];
60        hk.expand(&[], &mut output)
61            .expect("64 bytes is valid output length");
62
63        self.ck.copy_from_slice(&output[..32]);
64
65        // Initialize cipher with derived key for handshake encryption
66        let mut key = [0u8; 32];
67        key.copy_from_slice(&output[32..64]);
68        self.cipher.initialize_key(key);
69    }
70
71    /// Encrypt and mix into hash.
72    fn encrypt_and_hash(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseError> {
73        let ciphertext = self.cipher.encrypt(plaintext)?;
74        self.mix_hash(&ciphertext);
75        Ok(ciphertext)
76    }
77
78    /// Decrypt and mix ciphertext into hash.
79    fn decrypt_and_hash(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseError> {
80        let plaintext = self.cipher.decrypt(ciphertext)?;
81        self.mix_hash(ciphertext);
82        Ok(plaintext)
83    }
84
85    /// Split into two cipher states for transport.
86    fn split(&self) -> (CipherState, CipherState) {
87        let hk = Hkdf::<Sha256>::new(Some(&self.ck), &[]);
88        let mut output = [0u8; 64];
89        hk.expand(&[], &mut output)
90            .expect("64 bytes is valid output length");
91
92        let mut k1 = [0u8; 32];
93        let mut k2 = [0u8; 32];
94        k1.copy_from_slice(&output[..32]);
95        k2.copy_from_slice(&output[32..64]);
96
97        (CipherState::new(k1), CipherState::new(k2))
98    }
99
100    /// Get the handshake hash (for channel binding).
101    fn handshake_hash(&self) -> [u8; 32] {
102        self.h
103    }
104}
105
106/// Handshake state for Noise IK and XK patterns.
107pub struct HandshakeState {
108    /// Which Noise pattern is being used.
109    pattern: NoisePattern,
110    /// Our role in the handshake.
111    role: HandshakeRole,
112    /// Current progress.
113    progress: HandshakeProgress,
114    /// Symmetric state.
115    symmetric: SymmetricState,
116    /// Our static keypair.
117    static_keypair: Keypair,
118    /// Our ephemeral keypair (generated at handshake start).
119    ephemeral_keypair: Option<Keypair>,
120    /// Remote static public key.
121    /// For IK initiator: known before handshake (from config).
122    /// For IK responder: learned from message 1.
123    /// For XK initiator: known before handshake (from config).
124    /// For XK responder: learned from message 3.
125    remote_static: Option<PublicKey>,
126    /// Remote ephemeral public key (learned during handshake).
127    remote_ephemeral: Option<PublicKey>,
128    /// Secp256k1 context.
129    secp: Secp256k1<secp256k1::All>,
130    /// Our startup epoch for restart detection.
131    local_epoch: Option<[u8; 8]>,
132    /// Remote peer's startup epoch (learned during handshake).
133    remote_epoch: Option<[u8; 8]>,
134}
135
136impl HandshakeState {
137    /// Normalize a compressed public key to even parity for pre-message hashing.
138    ///
139    /// Nostr npubs encode x-only keys (no parity). The Noise IK pre-message
140    /// mixes the responder's static key into the hash before any messages.
141    /// Both sides must mix identical bytes. Since the initiator may only have
142    /// the x-only key (from an npub), we normalize to even parity (0x02 prefix)
143    /// so the hash chain matches regardless of the key's actual parity.
144    ///
145    /// This does NOT affect ECDH operations (which use x-coordinate-only output)
146    /// or the keys sent in handshake messages (which use actual parity).
147    fn normalize_for_premessage(pubkey: &PublicKey) -> [u8; PUBKEY_SIZE] {
148        let mut bytes = pubkey.serialize();
149        bytes[0] = 0x02; // Force even parity
150        bytes
151    }
152
153    /// Create a new IK handshake as initiator.
154    ///
155    /// The initiator knows the responder's static key and will send first.
156    /// Used by FMP (link layer).
157    pub fn new_initiator(static_keypair: Keypair, remote_static: PublicKey) -> Self {
158        let secp = Secp256k1::new();
159        let mut state = Self {
160            pattern: NoisePattern::Ik,
161            role: HandshakeRole::Initiator,
162            progress: HandshakeProgress::Initial,
163            symmetric: SymmetricState::initialize(PROTOCOL_NAME_IK),
164            static_keypair,
165            ephemeral_keypair: None,
166            remote_static: Some(remote_static),
167            remote_ephemeral: None,
168            secp,
169            local_epoch: None,
170            remote_epoch: None,
171        };
172
173        // Mix in pre-message: <- s (responder's static is known)
174        // Normalize to even parity so initiator and responder hash chains match
175        // even when the initiator only has the x-only key (from npub).
176        let normalized = Self::normalize_for_premessage(&remote_static);
177        state.symmetric.mix_hash(&normalized);
178
179        state
180    }
181
182    /// Create a new IK handshake as responder.
183    ///
184    /// The responder does NOT know the initiator's static key - it will be
185    /// learned from message 1. Used by FMP (link layer).
186    pub fn new_responder(static_keypair: Keypair) -> Self {
187        let secp = Secp256k1::new();
188        let mut state = Self {
189            pattern: NoisePattern::Ik,
190            role: HandshakeRole::Responder,
191            progress: HandshakeProgress::Initial,
192            symmetric: SymmetricState::initialize(PROTOCOL_NAME_IK),
193            static_keypair,
194            ephemeral_keypair: None,
195            remote_static: None, // Will learn from message 1
196            remote_ephemeral: None,
197            secp,
198            local_epoch: None,
199            remote_epoch: None,
200        };
201
202        // Mix in pre-message: <- s (our static, since we're responder)
203        // Normalize to even parity to match initiator's hash chain.
204        let normalized = Self::normalize_for_premessage(&state.static_keypair.public_key());
205        state.symmetric.mix_hash(&normalized);
206
207        state
208    }
209
210    /// Create a new XK handshake as initiator.
211    ///
212    /// The initiator knows the responder's static key. XK defers the
213    /// initiator's static key reveal to msg3. Used by FSP (session layer).
214    pub fn new_xk_initiator(static_keypair: Keypair, remote_static: PublicKey) -> Self {
215        let secp = Secp256k1::new();
216        let mut state = Self {
217            pattern: NoisePattern::Xk,
218            role: HandshakeRole::Initiator,
219            progress: HandshakeProgress::Initial,
220            symmetric: SymmetricState::initialize(PROTOCOL_NAME_XK),
221            static_keypair,
222            ephemeral_keypair: None,
223            remote_static: Some(remote_static),
224            remote_ephemeral: None,
225            secp,
226            local_epoch: None,
227            remote_epoch: None,
228        };
229
230        // Mix in pre-message: <- s (responder's static is known)
231        let normalized = Self::normalize_for_premessage(&remote_static);
232        state.symmetric.mix_hash(&normalized);
233
234        state
235    }
236
237    /// Create a new XK handshake as responder.
238    ///
239    /// The responder does NOT know the initiator's static key - it will be
240    /// learned from message 3. Used by FSP (session layer).
241    pub fn new_xk_responder(static_keypair: Keypair) -> Self {
242        let secp = Secp256k1::new();
243        let mut state = Self {
244            pattern: NoisePattern::Xk,
245            role: HandshakeRole::Responder,
246            progress: HandshakeProgress::Initial,
247            symmetric: SymmetricState::initialize(PROTOCOL_NAME_XK),
248            static_keypair,
249            ephemeral_keypair: None,
250            remote_static: None, // Will learn from message 3
251            remote_ephemeral: None,
252            secp,
253            local_epoch: None,
254            remote_epoch: None,
255        };
256
257        // Mix in pre-message: <- s (our static, since we're responder)
258        let normalized = Self::normalize_for_premessage(&state.static_keypair.public_key());
259        state.symmetric.mix_hash(&normalized);
260
261        state
262    }
263
264    /// Get our role.
265    pub fn role(&self) -> HandshakeRole {
266        self.role
267    }
268
269    /// Get current progress.
270    pub fn progress(&self) -> HandshakeProgress {
271        self.progress
272    }
273
274    /// Check if handshake is complete.
275    pub fn is_complete(&self) -> bool {
276        self.progress == HandshakeProgress::Complete
277    }
278
279    /// Get the remote static key (available after message 1 for responder).
280    pub fn remote_static(&self) -> Option<&PublicKey> {
281        self.remote_static.as_ref()
282    }
283
284    /// Set the local startup epoch for restart detection.
285    pub fn set_local_epoch(&mut self, epoch: [u8; 8]) {
286        self.local_epoch = Some(epoch);
287    }
288
289    /// Get the remote peer's startup epoch (available after processing their message).
290    pub fn remote_epoch(&self) -> Option<[u8; 8]> {
291        self.remote_epoch
292    }
293
294    /// Generate ephemeral keypair.
295    fn generate_ephemeral(&mut self) {
296        let mut rng = rand::rng();
297        let mut secret_bytes = [0u8; 32];
298        rng.fill_bytes(&mut secret_bytes);
299
300        let secret_key =
301            SecretKey::from_slice(&secret_bytes).expect("32 random bytes is valid secret key");
302        self.ephemeral_keypair = Some(Keypair::from_secret_key(&self.secp, &secret_key));
303    }
304
305    /// Perform ECDH between our secret and their public key.
306    ///
307    /// Uses x-only hashing (SHA-256 of just the x-coordinate) to produce
308    /// a parity-independent shared secret. This is necessary because Nostr
309    /// npubs encode x-only keys without parity information, so the initiator
310    /// may have the wrong parity for the responder's static key. Since P and
311    /// -P produce ECDH result points with the same x-coordinate, hashing
312    /// only x ensures both sides derive the same shared secret.
313    fn ecdh(&self, our_secret: &SecretKey, their_public: &PublicKey) -> [u8; 32] {
314        // Get raw (x, y) coordinates (64 bytes) without any hashing
315        let point = shared_secret_point(their_public, our_secret);
316        // Hash only the x-coordinate (first 32 bytes), ignoring y/parity
317        let mut hasher = Sha256::new();
318        hasher.update(&point[..32]);
319        let hash = hasher.finalize();
320        let mut result = [0u8; 32];
321        result.copy_from_slice(&hash);
322        result
323    }
324
325    /// Write message 1 (initiator only).
326    ///
327    /// Message 1 contains:
328    /// - e: ephemeral public key (33 bytes)
329    /// - encrypted s: our static public key encrypted (33 + 16 = 49 bytes)
330    /// - encrypted epoch: startup epoch for restart detection (8 + 16 = 24 bytes)
331    ///
332    /// Total: 106 bytes
333    pub fn write_message_1(&mut self) -> Result<Vec<u8>, NoiseError> {
334        if self.role != HandshakeRole::Initiator {
335            return Err(NoiseError::WrongState {
336                expected: "initiator".to_string(),
337                got: "responder".to_string(),
338            });
339        }
340        if self.progress != HandshakeProgress::Initial {
341            return Err(NoiseError::WrongState {
342                expected: HandshakeProgress::Initial.to_string(),
343                got: self.progress.to_string(),
344            });
345        }
346
347        let remote_static = self
348            .remote_static
349            .expect("initiator must have remote static");
350        let epoch = self
351            .local_epoch
352            .expect("local epoch must be set before write_message_1");
353
354        // Generate ephemeral keypair
355        self.generate_ephemeral();
356        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
357        let e_pub = ephemeral.public_key().serialize();
358
359        let mut message = Vec::with_capacity(HANDSHAKE_MSG1_SIZE);
360
361        // -> e: send ephemeral, mix into hash
362        message.extend_from_slice(&e_pub);
363        self.symmetric.mix_hash(&e_pub);
364
365        // -> es: DH(e, rs), mix into key
366        let es = self.ecdh(&ephemeral.secret_key(), &remote_static);
367        self.symmetric.mix_key(&es);
368
369        // -> s: encrypt our static and send
370        let our_static = self.static_keypair.public_key().serialize();
371        let encrypted_static = self.symmetric.encrypt_and_hash(&our_static)?;
372        message.extend_from_slice(&encrypted_static);
373
374        // -> ss: DH(s, rs), mix into key
375        let ss = self.ecdh(&self.static_keypair.secret_key(), &remote_static);
376        self.symmetric.mix_key(&ss);
377
378        // -> epoch: encrypt startup epoch for restart detection
379        let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?;
380        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
381        message.extend_from_slice(&encrypted_epoch);
382
383        self.progress = HandshakeProgress::Message1Done;
384
385        Ok(message)
386    }
387
388    /// Read message 1 (responder only).
389    ///
390    /// Processes the initiator's first message and learns their identity and epoch.
391    pub fn read_message_1(&mut self, message: &[u8]) -> Result<(), NoiseError> {
392        if self.role != HandshakeRole::Responder {
393            return Err(NoiseError::WrongState {
394                expected: "responder".to_string(),
395                got: "initiator".to_string(),
396            });
397        }
398        if self.progress != HandshakeProgress::Initial {
399            return Err(NoiseError::WrongState {
400                expected: HandshakeProgress::Initial.to_string(),
401                got: self.progress.to_string(),
402            });
403        }
404        if message.len() != HANDSHAKE_MSG1_SIZE {
405            return Err(NoiseError::MessageTooShort {
406                expected: HANDSHAKE_MSG1_SIZE,
407                got: message.len(),
408            });
409        }
410
411        // -> e: parse remote ephemeral, mix into hash
412        let re = PublicKey::from_slice(&message[..PUBKEY_SIZE])
413            .map_err(|_| NoiseError::InvalidPublicKey)?;
414        self.remote_ephemeral = Some(re);
415        self.symmetric.mix_hash(&message[..PUBKEY_SIZE]);
416
417        // -> es: DH(s, re), mix into key
418        // (responder uses their static with initiator's ephemeral)
419        let es = self.ecdh(&self.static_keypair.secret_key(), &re);
420        self.symmetric.mix_key(&es);
421
422        // -> s: decrypt initiator's static
423        let encrypted_static_end = PUBKEY_SIZE + PUBKEY_SIZE + super::TAG_SIZE;
424        let encrypted_static = &message[PUBKEY_SIZE..encrypted_static_end];
425        let decrypted_static = self.symmetric.decrypt_and_hash(encrypted_static)?;
426        let rs =
427            PublicKey::from_slice(&decrypted_static).map_err(|_| NoiseError::InvalidPublicKey)?;
428        self.remote_static = Some(rs);
429
430        // -> ss: DH(s, rs), mix into key
431        let ss = self.ecdh(&self.static_keypair.secret_key(), &rs);
432        self.symmetric.mix_key(&ss);
433
434        // -> epoch: decrypt initiator's startup epoch
435        let encrypted_epoch = &message[encrypted_static_end..];
436        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
437        let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?;
438        debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE);
439        let mut epoch = [0u8; EPOCH_SIZE];
440        epoch.copy_from_slice(&decrypted_epoch);
441        self.remote_epoch = Some(epoch);
442
443        self.progress = HandshakeProgress::Message1Done;
444
445        Ok(())
446    }
447
448    /// Write message 2 (responder only).
449    ///
450    /// Message 2 contains:
451    /// - e: ephemeral public key (33 bytes)
452    /// - encrypted epoch: startup epoch for restart detection (8 + 16 = 24 bytes)
453    ///
454    /// Total: 57 bytes
455    pub fn write_message_2(&mut self) -> Result<Vec<u8>, NoiseError> {
456        if self.role != HandshakeRole::Responder {
457            return Err(NoiseError::WrongState {
458                expected: "responder".to_string(),
459                got: "initiator".to_string(),
460            });
461        }
462        if self.progress != HandshakeProgress::Message1Done {
463            return Err(NoiseError::WrongState {
464                expected: HandshakeProgress::Message1Done.to_string(),
465                got: self.progress.to_string(),
466            });
467        }
468
469        let re = self.remote_ephemeral.expect("should have remote ephemeral");
470        let epoch = self
471            .local_epoch
472            .expect("local epoch must be set before write_message_2");
473
474        // Generate ephemeral keypair
475        self.generate_ephemeral();
476        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
477        let e_pub = ephemeral.public_key().serialize();
478
479        let mut message = Vec::with_capacity(HANDSHAKE_MSG2_SIZE);
480
481        // <- e: send ephemeral, mix into hash
482        message.extend_from_slice(&e_pub);
483        self.symmetric.mix_hash(&e_pub);
484
485        // <- ee: DH(e, re), mix into key
486        let ee = self.ecdh(&ephemeral.secret_key(), &re);
487        self.symmetric.mix_key(&ee);
488
489        // <- se: DH(s, re), mix into key
490        let se = self.ecdh(&self.static_keypair.secret_key(), &re);
491        self.symmetric.mix_key(&se);
492
493        // <- epoch: encrypt startup epoch for restart detection
494        let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?;
495        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
496        message.extend_from_slice(&encrypted_epoch);
497
498        self.progress = HandshakeProgress::Complete;
499
500        Ok(message)
501    }
502
503    /// Read message 2 (initiator only).
504    ///
505    /// Processes the responder's message and completes the handshake.
506    pub fn read_message_2(&mut self, message: &[u8]) -> Result<(), NoiseError> {
507        if self.role != HandshakeRole::Initiator {
508            return Err(NoiseError::WrongState {
509                expected: "initiator".to_string(),
510                got: "responder".to_string(),
511            });
512        }
513        if self.progress != HandshakeProgress::Message1Done {
514            return Err(NoiseError::WrongState {
515                expected: HandshakeProgress::Message1Done.to_string(),
516                got: self.progress.to_string(),
517            });
518        }
519        if message.len() != HANDSHAKE_MSG2_SIZE {
520            return Err(NoiseError::MessageTooShort {
521                expected: HANDSHAKE_MSG2_SIZE,
522                got: message.len(),
523            });
524        }
525
526        // <- e: parse remote ephemeral, mix into hash
527        let e_pub = &message[..PUBKEY_SIZE];
528        let re = PublicKey::from_slice(e_pub).map_err(|_| NoiseError::InvalidPublicKey)?;
529        self.remote_ephemeral = Some(re);
530        self.symmetric.mix_hash(e_pub);
531
532        // <- ee: DH(e, re), mix into key
533        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
534        let ee = self.ecdh(&ephemeral.secret_key(), &re);
535        self.symmetric.mix_key(&ee);
536
537        // <- se: DH(e, rs), mix into key
538        // (initiator uses their ephemeral with responder's static)
539        let rs = self.remote_static.expect("initiator has remote static");
540        let se = self.ecdh(&ephemeral.secret_key(), &rs);
541        self.symmetric.mix_key(&se);
542
543        // <- epoch: decrypt responder's startup epoch
544        let encrypted_epoch = &message[PUBKEY_SIZE..];
545        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
546        let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?;
547        debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE);
548        let mut epoch = [0u8; EPOCH_SIZE];
549        epoch.copy_from_slice(&decrypted_epoch);
550        self.remote_epoch = Some(epoch);
551
552        self.progress = HandshakeProgress::Complete;
553
554        Ok(())
555    }
556
557    // ========================================================================
558    // XK Pattern Methods (Session Layer)
559    // ========================================================================
560
561    /// Write XK message 1 (initiator only).
562    ///
563    /// XK msg1: `-> e, es`
564    /// - e: ephemeral public key (33 bytes)
565    /// - es: DH(e_priv, rs_pub), mix_key
566    ///
567    /// Total: 33 bytes (ephemeral only — no static, no epoch)
568    pub fn write_xk_message_1(&mut self) -> Result<Vec<u8>, NoiseError> {
569        if self.role != HandshakeRole::Initiator {
570            return Err(NoiseError::WrongState {
571                expected: "initiator".to_string(),
572                got: "responder".to_string(),
573            });
574        }
575        if self.progress != HandshakeProgress::Initial {
576            return Err(NoiseError::WrongState {
577                expected: HandshakeProgress::Initial.to_string(),
578                got: self.progress.to_string(),
579            });
580        }
581
582        let remote_static = self
583            .remote_static
584            .expect("initiator must have remote static");
585
586        // Generate ephemeral keypair
587        self.generate_ephemeral();
588        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
589        let e_pub = ephemeral.public_key().serialize();
590
591        let mut message = Vec::with_capacity(XK_HANDSHAKE_MSG1_SIZE);
592
593        // -> e: send ephemeral, mix into hash
594        message.extend_from_slice(&e_pub);
595        self.symmetric.mix_hash(&e_pub);
596
597        // -> es: DH(e, rs), mix into key
598        let es = self.ecdh(&ephemeral.secret_key(), &remote_static);
599        self.symmetric.mix_key(&es);
600
601        self.progress = HandshakeProgress::Message1Done;
602
603        Ok(message)
604    }
605
606    /// Read XK message 1 (responder only).
607    ///
608    /// Processes the initiator's first message. Does NOT learn initiator's
609    /// identity (that comes in msg3).
610    pub fn read_xk_message_1(&mut self, message: &[u8]) -> Result<(), NoiseError> {
611        if self.role != HandshakeRole::Responder {
612            return Err(NoiseError::WrongState {
613                expected: "responder".to_string(),
614                got: "initiator".to_string(),
615            });
616        }
617        if self.progress != HandshakeProgress::Initial {
618            return Err(NoiseError::WrongState {
619                expected: HandshakeProgress::Initial.to_string(),
620                got: self.progress.to_string(),
621            });
622        }
623        if message.len() != XK_HANDSHAKE_MSG1_SIZE {
624            return Err(NoiseError::MessageTooShort {
625                expected: XK_HANDSHAKE_MSG1_SIZE,
626                got: message.len(),
627            });
628        }
629
630        // -> e: parse remote ephemeral, mix into hash
631        let re = PublicKey::from_slice(&message[..PUBKEY_SIZE])
632            .map_err(|_| NoiseError::InvalidPublicKey)?;
633        self.remote_ephemeral = Some(re);
634        self.symmetric.mix_hash(&message[..PUBKEY_SIZE]);
635
636        // -> es: DH(s, re), mix into key
637        // (responder uses their static with initiator's ephemeral)
638        let es = self.ecdh(&self.static_keypair.secret_key(), &re);
639        self.symmetric.mix_key(&es);
640
641        self.progress = HandshakeProgress::Message1Done;
642
643        Ok(())
644    }
645
646    /// Write XK message 2 (responder only).
647    ///
648    /// XK msg2: `<- e, ee` + encrypted epoch
649    /// - e: ephemeral public key (33 bytes)
650    /// - ee: DH(e_priv, re_pub), mix_key
651    /// - encrypted epoch (24 bytes)
652    ///
653    /// Total: 57 bytes
654    pub fn write_xk_message_2(&mut self) -> Result<Vec<u8>, NoiseError> {
655        if self.role != HandshakeRole::Responder {
656            return Err(NoiseError::WrongState {
657                expected: "responder".to_string(),
658                got: "initiator".to_string(),
659            });
660        }
661        if self.progress != HandshakeProgress::Message1Done {
662            return Err(NoiseError::WrongState {
663                expected: HandshakeProgress::Message1Done.to_string(),
664                got: self.progress.to_string(),
665            });
666        }
667
668        let re = self.remote_ephemeral.expect("should have remote ephemeral");
669        let epoch = self
670            .local_epoch
671            .expect("local epoch must be set before write_xk_message_2");
672
673        // Generate ephemeral keypair
674        self.generate_ephemeral();
675        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
676        let e_pub = ephemeral.public_key().serialize();
677
678        let mut message = Vec::with_capacity(XK_HANDSHAKE_MSG2_SIZE);
679
680        // <- e: send ephemeral, mix into hash
681        message.extend_from_slice(&e_pub);
682        self.symmetric.mix_hash(&e_pub);
683
684        // <- ee: DH(e, re), mix into key
685        let ee = self.ecdh(&ephemeral.secret_key(), &re);
686        self.symmetric.mix_key(&ee);
687
688        // <- epoch: encrypt startup epoch for restart detection
689        let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?;
690        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
691        message.extend_from_slice(&encrypted_epoch);
692
693        self.progress = HandshakeProgress::Message2Done;
694
695        Ok(message)
696    }
697
698    /// Read XK message 2 (initiator only).
699    ///
700    /// Processes the responder's message and extracts the responder's epoch.
701    /// Does NOT complete the handshake — msg3 still needed.
702    pub fn read_xk_message_2(&mut self, message: &[u8]) -> Result<(), NoiseError> {
703        if self.role != HandshakeRole::Initiator {
704            return Err(NoiseError::WrongState {
705                expected: "initiator".to_string(),
706                got: "responder".to_string(),
707            });
708        }
709        if self.progress != HandshakeProgress::Message1Done {
710            return Err(NoiseError::WrongState {
711                expected: HandshakeProgress::Message1Done.to_string(),
712                got: self.progress.to_string(),
713            });
714        }
715        if message.len() != XK_HANDSHAKE_MSG2_SIZE {
716            return Err(NoiseError::MessageTooShort {
717                expected: XK_HANDSHAKE_MSG2_SIZE,
718                got: message.len(),
719            });
720        }
721
722        // <- e: parse remote ephemeral, mix into hash
723        let e_pub = &message[..PUBKEY_SIZE];
724        let re = PublicKey::from_slice(e_pub).map_err(|_| NoiseError::InvalidPublicKey)?;
725        self.remote_ephemeral = Some(re);
726        self.symmetric.mix_hash(e_pub);
727
728        // <- ee: DH(e, re), mix into key
729        let ephemeral = self.ephemeral_keypair.as_ref().unwrap();
730        let ee = self.ecdh(&ephemeral.secret_key(), &re);
731        self.symmetric.mix_key(&ee);
732
733        // <- epoch: decrypt responder's startup epoch
734        let encrypted_epoch = &message[PUBKEY_SIZE..];
735        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
736        let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?;
737        debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE);
738        let mut epoch = [0u8; EPOCH_SIZE];
739        epoch.copy_from_slice(&decrypted_epoch);
740        self.remote_epoch = Some(epoch);
741
742        self.progress = HandshakeProgress::Message2Done;
743
744        Ok(())
745    }
746
747    /// Write XK message 3 (initiator only).
748    ///
749    /// XK msg3: `-> s, se` + encrypted epoch
750    /// - s: encrypt_and_hash(s_pub) — encrypted static (49 bytes)
751    /// - se: DH(s_priv, re_pub), mix_key
752    /// - encrypted epoch (24 bytes)
753    ///
754    /// Total: 73 bytes
755    pub fn write_xk_message_3(&mut self) -> Result<Vec<u8>, NoiseError> {
756        if self.role != HandshakeRole::Initiator {
757            return Err(NoiseError::WrongState {
758                expected: "initiator".to_string(),
759                got: "responder".to_string(),
760            });
761        }
762        if self.progress != HandshakeProgress::Message2Done {
763            return Err(NoiseError::WrongState {
764                expected: HandshakeProgress::Message2Done.to_string(),
765                got: self.progress.to_string(),
766            });
767        }
768
769        let re = self
770            .remote_ephemeral
771            .expect("should have remote ephemeral after msg2");
772        let epoch = self
773            .local_epoch
774            .expect("local epoch must be set before write_xk_message_3");
775
776        let mut message = Vec::with_capacity(XK_HANDSHAKE_MSG3_SIZE);
777
778        // -> s: encrypt our static and send
779        let our_static = self.static_keypair.public_key().serialize();
780        let encrypted_static = self.symmetric.encrypt_and_hash(&our_static)?;
781        message.extend_from_slice(&encrypted_static);
782
783        // -> se: DH(s, re), mix into key
784        let se = self.ecdh(&self.static_keypair.secret_key(), &re);
785        self.symmetric.mix_key(&se);
786
787        // -> epoch: encrypt startup epoch for restart detection
788        let encrypted_epoch = self.symmetric.encrypt_and_hash(&epoch)?;
789        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
790        message.extend_from_slice(&encrypted_epoch);
791
792        self.progress = HandshakeProgress::Complete;
793
794        Ok(message)
795    }
796
797    /// Read XK message 3 (responder only).
798    ///
799    /// Processes the initiator's encrypted static key and epoch.
800    /// After this, the responder learns the initiator's identity.
801    pub fn read_xk_message_3(&mut self, message: &[u8]) -> Result<(), NoiseError> {
802        let previous_symmetric = self.symmetric.clone();
803        let previous_progress = self.progress;
804        let previous_remote_static = self.remote_static;
805        let previous_remote_epoch = self.remote_epoch;
806        let result = self.read_xk_message_3_in_place(message);
807        if result.is_err() {
808            self.symmetric = previous_symmetric;
809            self.progress = previous_progress;
810            self.remote_static = previous_remote_static;
811            self.remote_epoch = previous_remote_epoch;
812        }
813        result
814    }
815
816    fn read_xk_message_3_in_place(&mut self, message: &[u8]) -> Result<(), NoiseError> {
817        if self.role != HandshakeRole::Responder {
818            return Err(NoiseError::WrongState {
819                expected: "responder".to_string(),
820                got: "initiator".to_string(),
821            });
822        }
823        if self.progress != HandshakeProgress::Message2Done {
824            return Err(NoiseError::WrongState {
825                expected: HandshakeProgress::Message2Done.to_string(),
826                got: self.progress.to_string(),
827            });
828        }
829        if message.len() != XK_HANDSHAKE_MSG3_SIZE {
830            return Err(NoiseError::MessageTooShort {
831                expected: XK_HANDSHAKE_MSG3_SIZE,
832                got: message.len(),
833            });
834        }
835
836        // -> s: decrypt initiator's static
837        let encrypted_static_end = PUBKEY_SIZE + super::TAG_SIZE;
838        let encrypted_static = &message[..encrypted_static_end];
839        let decrypted_static = self.symmetric.decrypt_and_hash(encrypted_static)?;
840        let rs =
841            PublicKey::from_slice(&decrypted_static).map_err(|_| NoiseError::InvalidPublicKey)?;
842        self.remote_static = Some(rs);
843
844        // -> se: DH(e, rs), mix into key
845        // (responder uses their ephemeral with initiator's now-known static)
846        let ephemeral = self
847            .ephemeral_keypair
848            .as_ref()
849            .expect("should have ephemeral after msg2");
850        let se = self.ecdh(&ephemeral.secret_key(), &rs);
851        self.symmetric.mix_key(&se);
852
853        // -> epoch: decrypt initiator's startup epoch
854        let encrypted_epoch = &message[encrypted_static_end..];
855        debug_assert_eq!(encrypted_epoch.len(), EPOCH_ENCRYPTED_SIZE);
856        let decrypted_epoch = self.symmetric.decrypt_and_hash(encrypted_epoch)?;
857        debug_assert_eq!(decrypted_epoch.len(), EPOCH_SIZE);
858        let mut epoch = [0u8; EPOCH_SIZE];
859        epoch.copy_from_slice(&decrypted_epoch);
860        self.remote_epoch = Some(epoch);
861
862        self.progress = HandshakeProgress::Complete;
863
864        Ok(())
865    }
866
867    /// Complete the handshake and return a NoiseSession.
868    ///
869    /// Must be called after the handshake is complete.
870    pub fn into_session(self) -> Result<NoiseSession, NoiseError> {
871        if !self.is_complete() {
872            return Err(NoiseError::HandshakeNotComplete);
873        }
874
875        let (c1, c2) = self.symmetric.split();
876        let handshake_hash = self.symmetric.handshake_hash();
877        let remote_static = self
878            .remote_static
879            .expect("remote static must be known after handshake");
880        let remote_epoch = self
881            .remote_epoch
882            .expect("remote epoch must be known after handshake");
883
884        // Initiator sends with c1, receives with c2
885        // Responder sends with c2, receives with c1
886        let (send_cipher, recv_cipher) = match self.role {
887            HandshakeRole::Initiator => (c1, c2),
888            HandshakeRole::Responder => (c2, c1),
889        };
890
891        Ok(NoiseSession::from_handshake(
892            self.role,
893            send_cipher,
894            recv_cipher,
895            handshake_hash,
896            remote_static,
897            remote_epoch,
898        ))
899    }
900
901    /// Get the handshake hash (for channel binding, available after complete).
902    pub fn handshake_hash(&self) -> [u8; 32] {
903        self.symmetric.handshake_hash()
904    }
905}
906
907impl fmt::Debug for HandshakeState {
908    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
909        f.debug_struct("HandshakeState")
910            .field("pattern", &self.pattern)
911            .field("role", &self.role)
912            .field("progress", &self.progress)
913            .field("has_ephemeral", &self.ephemeral_keypair.is_some())
914            .field("has_remote_static", &self.remote_static.is_some())
915            .field("has_remote_ephemeral", &self.remote_ephemeral.is_some())
916            .field("has_local_epoch", &self.local_epoch.is_some())
917            .field("has_remote_epoch", &self.remote_epoch.is_some())
918            .finish()
919    }
920}