wlan_sae/
lib.rs

1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4mod boringssl;
5mod ecc;
6mod frame;
7pub mod hmac_utils;
8mod state;
9
10use anyhow::{Error, bail};
11use boringssl::{Bignum, EcGroupId};
12use fidl_fuchsia_wlan_ieee80211 as fidl_ieee80211;
13pub use frame::{AntiCloggingTokenMsg, CommitMsg, ConfirmMsg};
14use hmac_utils::{HmacUtils, HmacUtilsImpl};
15use ieee80211::{MacAddr, Ssid};
16use log::warn;
17use mundane::hash::Sha256;
18use num::FromPrimitive;
19use wlan_common::ie::rsn::akm::{self, Akm};
20
21/// Maximum number of incorrect frames sent before SAE fails.
22const MAX_RETRIES_PER_EXCHANGE: u16 = 3;
23
24/// A shared key computed by an SAE handshake.
25#[derive(Clone, PartialEq, Debug)]
26pub struct Key {
27    pub pmk: Vec<u8>,
28    pub pmkid: Vec<u8>,
29}
30
31/// IEEE Std 802.11-2020 9.4.2.241
32/// Method used to generate the PWE from a password.
33#[derive(Debug, Clone, PartialEq)]
34pub enum PweMethod {
35    /// IEEE Std 802.11-2020, 12.4.4.2.2/12.4.4.3.2
36    /// Generate the PWE using the looping hunt-and-peck method.
37    Loop = 0,
38
39    /// IEEE Std 802.11-2020, 12.4.4.2.3/12.4.4.3.3
40    /// Generate the PWE using the direct hashing method, hash-to-curve or hash-to-element.
41    Direct = 1,
42}
43
44/// Types of timeout that are used by SAE handshakes. Duration and scheduling of these timeouts
45/// is left to the user of this library.
46#[derive(Debug, Clone, PartialEq)]
47pub enum Timeout {
48    /// Timeout before the most recent message(s) should be resent.
49    Retransmission,
50    /// Timeout before the PMK produced by a successful handshake is considered invalid.
51    KeyExpiration,
52}
53
54#[derive(Debug)]
55pub enum RejectReason {
56    /// We experienced a failure that was unrelated to data received from the peer. This likely
57    /// means we are not in a good state.
58    InternalError(Error),
59    /// Data received from the peer failed validation, and we cannot generate a PMK.
60    AuthFailed,
61    /// The peer has failed to respond or sent incorrect responses too many times.
62    TooManyRetries,
63    /// The SAE PMKSA has expired, reauthenticate.
64    KeyExpiration,
65}
66
67impl From<Error> for RejectReason {
68    fn from(e: Error) -> Self {
69        Self::InternalError(e)
70    }
71}
72
73#[derive(Debug)]
74pub struct AuthFrameRx<'a> {
75    pub seq: u16,
76    pub status_code: fidl_ieee80211::StatusCode,
77    pub body: &'a [u8],
78}
79
80#[derive(Debug, Clone, Eq, PartialEq)]
81pub struct AuthFrameTx {
82    pub seq: u16,
83    pub status_code: fidl_ieee80211::StatusCode,
84    pub body: Vec<u8>,
85}
86
87/// An update generated to progress an SAE handshake. These updates should generally be converted
88/// into a frame and sent to the SAE peer.
89#[derive(Debug)]
90pub enum SaeUpdate {
91    /// Send an auth frame to the peer.
92    SendFrame(AuthFrameTx),
93    /// Indicates the handshake is complete. The handshake should *not* be deleted at this point.
94    Success(Key),
95    /// Indicates that the handshake has failed and must be aborted or restarted.
96    Reject(RejectReason),
97    /// Request the user of the library to set or reset a timeout. If this timeout expires, it
98    /// should be passed to SaeHandshake::handle_timeout.
99    ResetTimeout(Timeout),
100    /// Request the user of the library to cancel a timeout that was previously set.
101    CancelTimeout(Timeout),
102}
103
104pub type SaeUpdateSink = Vec<SaeUpdate>;
105
106/// IEEE 802.11-2016 12.4: Simultaneous Authentication of Equals (SAE)
107///
108/// An SAE handshake with a peer is a symmetric handshake that may be used in place of open
109/// authentication as the AKM. A full handshake consists of both peers sending a Commit and Confirm
110/// frame, at which point they have both derived a shared key that is unique to those peers and that
111/// session.
112///
113/// Structs implementing this trait are responsible for handling both a successful SAE handshake,
114/// various failure modes, and edge cases such as retries and timeouts.
115///
116/// None of the functions in this trait return errors. Instead, non-fatal errors are logged, and
117/// fatal errors push an SaeUpdate::Reject to the update sink. Once an SaeUpdate::Reject is pushed,
118/// all further operations are no-ops.
119pub trait SaeHandshake: Send {
120    /// Initiate SAE by sending the first commit message. If the peer STA sends the first commit
121    /// message, handle_commit should be called first and initiate_sae should never be called.
122    fn initiate_sae(&mut self, sink: &mut SaeUpdateSink);
123
124    fn handle_commit(&mut self, sink: &mut SaeUpdateSink, commit_msg: &CommitMsg<'_>);
125    fn handle_confirm(&mut self, sink: &mut SaeUpdateSink, confirm_msg: &ConfirmMsg<'_>);
126    fn handle_anti_clogging_token(
127        &mut self,
128        sink: &mut SaeUpdateSink,
129        act_msg: &AntiCloggingTokenMsg<'_>,
130    );
131    fn handle_timeout(&mut self, sink: &mut SaeUpdateSink, timeout: Timeout);
132
133    fn handle_frame(&mut self, sink: &mut SaeUpdateSink, frame: &AuthFrameRx<'_>) {
134        match frame::parse(frame) {
135            Ok(parse) => match parse {
136                frame::ParseSuccess::Commit(commit) => self.handle_commit(sink, &commit),
137                frame::ParseSuccess::Confirm(confirm) => self.handle_confirm(sink, &confirm),
138                frame::ParseSuccess::AntiCloggingToken(act_msg) => {
139                    self.handle_anti_clogging_token(sink, &act_msg)
140                }
141            },
142            Err(e) => warn!("Failed to parse SAE auth frame: {}", e),
143        }
144    }
145}
146
147/// Creates a new SAE handshake for the given group ID and authentication parameters.
148pub fn new_sae_handshake(
149    group_id: u16,
150    akm: Akm,
151    pwe_method: PweMethod,
152    ssid: Ssid,
153    password: Vec<u8>,
154    password_id: Option<Vec<u8>>,
155    mac: MacAddr,
156    peer_mac: MacAddr,
157) -> Result<Box<dyn SaeHandshake>, Error> {
158    match akm.suite_type {
159        akm::SAE | akm::FT_SAE => (),
160        _ => bail!("Cannot construct SAE handshake with AKM {:?}", akm),
161    };
162    let (hmac, group_constructor) = match EcGroupId::from_u16(group_id) {
163        Some(EcGroupId::P256) => {
164            // IEEE 802.11-2020 12.4.2
165            // Group 19 has a 256-bit prime length, thus we use SHA256.
166            let hmac = Box::new(HmacUtilsImpl::<Sha256>::new());
167            let group_constructor = Box::new(|| {
168                ecc::Group::new(EcGroupId::P256).map(|group| {
169                    Box::new(group)
170                        as Box<
171                            dyn internal::FiniteCyclicGroup<
172                                    Element = <ecc::Group as internal::FiniteCyclicGroup>::Element,
173                                >,
174                        >
175                })
176            });
177            (hmac, group_constructor)
178        }
179        _ => bail!("Unsupported SAE group id: {}", group_id),
180    };
181    Ok(Box::new(state::SaeHandshakeImpl::new(
182        group_constructor,
183        internal::SaeParameters {
184            hmac,
185            pwe_method,
186            ssid,
187            password,
188            password_id,
189            sta_a_mac: mac,
190            sta_b_mac: peer_mac,
191        },
192    )?))
193}
194
195/// Creates a new SAE handshake in response to a first message from a peer, using the FCG indiated
196/// by the peer if possible. In a successful handshake, this will immediately push a Commit and
197/// Confirm to the given update sink.
198pub fn join_sae_handshake(
199    sink: &mut SaeUpdateSink,
200    first_frame: &AuthFrameRx<'_>,
201    akm: Akm,
202    ssid: Ssid,
203    password: Vec<u8>,
204    mac: MacAddr,
205    peer_mac: MacAddr,
206) -> Result<Box<dyn SaeHandshake>, Error> {
207    let parsed_frame = frame::parse(first_frame)?;
208    match parsed_frame {
209        frame::ParseSuccess::Commit(commit) => {
210            let mut handshake = new_sae_handshake(
211                commit.group_id,
212                akm,
213                PweMethod::Loop,
214                ssid,
215                password,
216                None,
217                mac,
218                peer_mac,
219            )?;
220            handshake.handle_commit(sink, &commit);
221            Ok(handshake)
222        }
223        _ => bail!("Recieved incorrect first frame of SAE handshake"),
224    }
225}
226
227// Internal mod for structs with mod-public visibility.
228mod internal {
229    use super::*;
230
231    /// IEEE 802.11-2016 12.4.4
232    /// SAE may use many different finite cyclic groups (FCGs) to compute the various values used
233    /// during the handshake. This trait allows our SAE implementation to seamlessly handle
234    /// different classes of FCG. IEEE 802.11-2016 defines support for both elliptic curve groups
235    /// and finite field cryptography groups.
236    ///
237    /// All functions provided by this trait will only return an Error when something internal has
238    /// gone wrong.
239    pub trait FiniteCyclicGroup {
240        /// Different classes of FCG have different Element types, but scalars can always be
241        /// represented by a Bignum.
242        type Element;
243
244        fn group_id(&self) -> u16;
245
246        /// IEEE 802.11-2016 12.4.3
247        /// Generates a new password element, a secret value shared by the two peers in SAE.
248        fn generate_pwe(&self, params: &SaeParameters) -> Result<Self::Element, Error>;
249
250        /// IEEE 12.4.4.1
251        /// These three operators are used to manipulate FCG elements for the purposes of the
252        /// Diffie-Hellman key exchange used by SAE.
253        fn scalar_op(
254            &self,
255            scalar: &Bignum,
256            element: &Self::Element,
257        ) -> Result<Self::Element, Error>;
258        fn elem_op(
259            &self,
260            element1: &Self::Element,
261            element2: &Self::Element,
262        ) -> Result<Self::Element, Error>;
263        fn inverse_op(&self, element: Self::Element) -> Result<Self::Element, Error>;
264
265        /// Returns the prime order of the FCG.
266        fn order(&self) -> Result<Bignum, Error>;
267        /// IEEE 802.11-2016 12.4.5.4
268        /// Maps the given secret element to the shared secret value. Returns None if this is the
269        /// identity element for this FCG, indicating that we have in invalid secret element.
270        fn map_to_secret_value(&self, element: &Self::Element) -> Result<Option<Vec<u8>>, Error>;
271        /// IEEE 802.11-2016 12.4.2: The FCG Element must convert into an octet string such
272        /// that it may be included in the confirmation hash when completing SAE.
273        fn element_to_octets(&self, element: &Self::Element) -> Result<Vec<u8>, Error>;
274        /// Convert octets into an element. Returns None if the given octet string does not
275        /// contain a valid element for this group.
276        fn element_from_octets(&self, octets: &[u8]) -> Result<Option<Self::Element>, Error>;
277
278        /// Return the expected size of scalar and element values when serialized into a frame.
279        fn scalar_size(&self) -> Result<usize, Error> {
280            self.order().map(|order| order.len())
281        }
282    }
283
284    pub struct SaeParameters {
285        pub hmac: Box<dyn HmacUtils + Send>,
286        pub pwe_method: PweMethod,
287        // IEEE Std 802.11-2020 12.4.4.2.3/12.4.4.3.3: The SSID is needed to generate a password
288        // seed.
289        pub ssid: Ssid,
290        // IEEE Std 802.11-2020 12.4.3
291        pub password: Vec<u8>,
292        pub password_id: Option<Vec<u8>>,
293        // IEEE Std 802.11-2016 12.4.4.2.2: The MacAddrs are needed to generate a password seed.
294        pub sta_a_mac: MacAddr,
295        pub sta_b_mac: MacAddr,
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    #![allow(unused_variables)] // Allow unused variables in tests.
302    use super::*;
303    use assert_matches::assert_matches;
304    use std::convert::TryFrom;
305    use std::sync::LazyLock;
306    use wlan_common::ie::rsn::akm::{AKM_PSK, AKM_SAE};
307
308    // IEEE 802.11-2016 Annex J.10 SAE test vector
309    const TEST_SSID: &'static str = "SSID not in 802.11-2016";
310    const TEST_PWD: &'static str = "thisisreallysecret";
311
312    pub(crate) static TEST_STA_A: LazyLock<MacAddr> =
313        LazyLock::new(|| MacAddr::from([0x7b, 0x88, 0x56, 0x20, 0x2d, 0x8d]));
314    pub(crate) static TEST_STA_B: LazyLock<MacAddr> =
315        LazyLock::new(|| MacAddr::from([0xe2, 0x47, 0x1c, 0x0a, 0x5a, 0xcb]));
316
317    #[test]
318    fn bad_akm() {
319        let akm = AKM_PSK;
320        let res = new_sae_handshake(
321            19,
322            akm,
323            PweMethod::Loop,
324            Ssid::try_from(TEST_SSID).unwrap(),
325            Vec::from(TEST_PWD),
326            None, // Not required for PweMethod::Loop
327            *TEST_STA_A,
328            *TEST_STA_B,
329        );
330        assert!(res.is_err());
331        assert!(
332            format!("{}", res.err().unwrap())
333                .contains("Cannot construct SAE handshake with AKM 00-0F-AC:2")
334        );
335    }
336
337    #[test]
338    fn bad_fcg() {
339        let akm = AKM_SAE;
340        let res = new_sae_handshake(
341            200,
342            akm,
343            PweMethod::Loop,
344            Ssid::try_from(TEST_SSID).unwrap(),
345            Vec::from(TEST_PWD),
346            None, // Not required for PweMethod::Loop
347            *TEST_STA_A,
348            *TEST_STA_B,
349        );
350        assert!(res.is_err());
351        assert!(format!("{}", res.err().unwrap()).contains("Unsupported SAE group id: 200"));
352    }
353
354    struct TestHandshake {
355        sta1: Box<dyn SaeHandshake>,
356        sta2: Box<dyn SaeHandshake>,
357    }
358
359    // Helper structs for differentiating Commit/Confirm messages once they've been converted into
360    // generic auth frames.
361    #[derive(Clone, Eq, PartialEq, Debug)]
362    struct CommitTx(AuthFrameTx);
363    #[derive(Clone, Eq, PartialEq, Debug)]
364    struct ConfirmTx(AuthFrameTx);
365    struct CommitRx<'a>(AuthFrameRx<'a>);
366    struct ConfirmRx<'a>(AuthFrameRx<'a>);
367
368    fn to_rx(frame: &AuthFrameTx) -> AuthFrameRx<'_> {
369        AuthFrameRx { seq: frame.seq, status_code: frame.status_code, body: &frame.body[..] }
370    }
371
372    impl CommitTx {
373        fn to_rx(&self) -> CommitRx<'_> {
374            CommitRx(to_rx(&self.0))
375        }
376    }
377
378    impl ConfirmTx {
379        fn to_rx(&self) -> ConfirmRx<'_> {
380            ConfirmRx(to_rx(&self.0))
381        }
382    }
383
384    impl<'a> CommitRx<'a> {
385        fn msg(&'a self) -> CommitMsg<'a> {
386            assert_matches!(frame::parse(&self.0),
387                Ok(frame::ParseSuccess::Commit(commit)) => commit)
388        }
389    }
390
391    impl<'a> ConfirmRx<'a> {
392        fn msg(&'a self) -> ConfirmMsg<'a> {
393            assert_matches!(frame::parse(&self.0),
394                Ok(frame::ParseSuccess::Confirm(confirm)) => confirm)
395        }
396    }
397
398    fn expect_commit(sink: &mut Vec<SaeUpdate>) -> CommitTx {
399        let commit = assert_matches!(sink.remove(0), SaeUpdate::SendFrame(frame) => frame);
400        assert_matches!(frame::parse(&to_rx(&commit)), Ok(frame::ParseSuccess::Commit(msg)));
401        CommitTx(commit)
402    }
403
404    fn expect_confirm(sink: &mut Vec<SaeUpdate>) -> ConfirmTx {
405        let confirm = assert_matches!(sink.remove(0), SaeUpdate::SendFrame(frame) => frame);
406        assert_matches!(frame::parse(&to_rx(&confirm)), Ok(frame::ParseSuccess::Confirm(msg)));
407        ConfirmTx(confirm)
408    }
409
410    fn expect_reset_timeout(sink: &mut Vec<SaeUpdate>, timeout: Timeout) {
411        assert_matches!(sink.remove(0), SaeUpdate::ResetTimeout(timeout));
412    }
413
414    fn expect_cancel_timeout(sink: &mut Vec<SaeUpdate>, timeout: Timeout) {
415        assert_matches!(sink.remove(0), SaeUpdate::CancelTimeout(timeout));
416    }
417
418    // Test helper to advance through successful steps of an SAE handshake.
419    impl TestHandshake {
420        fn new() -> Self {
421            let akm = AKM_SAE;
422            let sta1 = new_sae_handshake(
423                19,
424                akm.clone(),
425                PweMethod::Loop,
426                Ssid::try_from(TEST_SSID).unwrap(),
427                Vec::from(TEST_PWD),
428                None, // Not required for PweMethod::Loop
429                *TEST_STA_A,
430                *TEST_STA_B,
431            )
432            .unwrap();
433            let sta2 = new_sae_handshake(
434                19,
435                akm,
436                PweMethod::Loop,
437                Ssid::try_from(TEST_SSID).unwrap(),
438                Vec::from(TEST_PWD),
439                None, // Not required for PweMethod::Loop
440                *TEST_STA_B,
441                *TEST_STA_A,
442            )
443            .unwrap();
444            Self { sta1, sta2 }
445        }
446
447        fn sta1_init(&mut self) -> CommitTx {
448            let mut sink = vec![];
449            self.sta1.initiate_sae(&mut sink);
450            assert_eq!(sink.len(), 2);
451            let commit = expect_commit(&mut sink);
452            expect_reset_timeout(&mut sink, Timeout::Retransmission);
453            commit
454        }
455
456        fn sta2_handle_commit(&mut self, commit1: CommitRx<'_>) -> (CommitTx, ConfirmTx) {
457            let mut sink = vec![];
458            self.sta2.handle_commit(&mut sink, &commit1.msg());
459            assert_eq!(sink.len(), 3);
460            let commit2 = expect_commit(&mut sink);
461            let confirm2 = expect_confirm(&mut sink);
462            expect_reset_timeout(&mut sink, Timeout::Retransmission);
463            (commit2, confirm2)
464        }
465
466        fn sta1_handle_commit(&mut self, commit2: CommitRx<'_>) -> ConfirmTx {
467            let mut sink = vec![];
468            self.sta1.handle_commit(&mut sink, &commit2.msg());
469            assert_eq!(sink.len(), 2);
470            let confirm1 = expect_confirm(&mut sink);
471            expect_reset_timeout(&mut sink, Timeout::Retransmission);
472            confirm1
473        }
474
475        fn sta1_handle_confirm(&mut self, confirm2: ConfirmRx<'_>) -> Key {
476            Self::__internal_handle_confirm(&mut self.sta1, confirm2.msg())
477        }
478
479        fn sta2_handle_confirm(&mut self, confirm1: ConfirmRx<'_>) -> Key {
480            Self::__internal_handle_confirm(&mut self.sta2, confirm1.msg())
481        }
482
483        fn __internal_handle_confirm(
484            sta: &mut Box<dyn SaeHandshake>,
485            confirm: ConfirmMsg<'_>,
486        ) -> Key {
487            let mut sink = vec![];
488            sta.handle_confirm(&mut sink, &confirm);
489            assert_eq!(sink.len(), 3);
490            expect_cancel_timeout(&mut sink, Timeout::Retransmission);
491            expect_reset_timeout(&mut sink, Timeout::KeyExpiration);
492            assert_matches!(sink.remove(0), SaeUpdate::Success(key) => key)
493        }
494    }
495
496    #[test]
497    fn sae_handshake_success() {
498        let mut handshake = TestHandshake::new();
499        let commit1 = handshake.sta1_init();
500        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
501        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
502        let key1 = handshake.sta1_handle_confirm(confirm2.to_rx());
503        let key2 = handshake.sta2_handle_confirm(confirm1.to_rx());
504        assert_eq!(key1, key2);
505    }
506
507    #[test]
508    fn password_mismatch() {
509        let akm = AKM_SAE;
510        let sta1 = new_sae_handshake(
511            19,
512            akm.clone(),
513            PweMethod::Loop,
514            Ssid::try_from(TEST_SSID).unwrap(),
515            Vec::from(TEST_PWD),
516            None, // Not required for PweMethod::Loop
517            *TEST_STA_A,
518            *TEST_STA_B,
519        )
520        .unwrap();
521        let sta2 = new_sae_handshake(
522            19,
523            akm,
524            PweMethod::Loop,
525            Ssid::try_from(TEST_SSID).unwrap(),
526            Vec::from("other_pwd"),
527            None, // Not required for PweMethod::Loop
528            *TEST_STA_B,
529            *TEST_STA_A,
530        )
531        .unwrap();
532        let mut handshake = TestHandshake { sta1, sta2 };
533
534        let commit1 = handshake.sta1_init();
535        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
536        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
537
538        let mut sink1 = vec![];
539        handshake.sta1.handle_confirm(&mut sink1, &confirm2.to_rx().msg());
540        let mut sink2 = vec![];
541        handshake.sta2.handle_confirm(&mut sink2, &confirm1.to_rx().msg());
542        // The confirm is dropped both ways.
543        assert_eq!(sink1.len(), 0);
544        assert_eq!(sink2.len(), 0);
545    }
546
547    #[test]
548    fn retry_commit_on_unexpected_confirm() {
549        let mut handshake = TestHandshake::new();
550
551        let commit1 = handshake.sta1_init();
552        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
553        let mut sink = vec![];
554        handshake.sta1.handle_confirm(&mut sink, &confirm2.to_rx().msg());
555        assert_eq!(sink.len(), 2);
556        let commit1_retry = expect_commit(&mut sink);
557        assert_matches!(sink.remove(0), SaeUpdate::ResetTimeout(Timeout::Retransmission));
558
559        // We retransmit the same commit in response to a faulty confirm.
560        assert_eq!(commit1, commit1_retry);
561    }
562
563    #[test]
564    fn retry_commit_on_anti_clogging_token() {
565        let mut handshake = TestHandshake::new();
566
567        let commit1 = handshake.sta1_init();
568
569        // Simulate an anti-clogging token sent to sta1.
570        let mut sink = vec![];
571        let anti_clogging_token = "anticloggingtokentext";
572        let act_msg = AntiCloggingTokenMsg {
573            group_id: 19,
574            anti_clogging_token: anti_clogging_token.as_bytes(),
575        };
576        handshake.sta1.handle_anti_clogging_token(&mut sink, &act_msg);
577        let commit1_retry = expect_commit(&mut sink);
578        assert_eq!(
579            commit1_retry.clone().to_rx().msg().anti_clogging_token,
580            Some(anti_clogging_token.as_bytes())
581        );
582
583        // Finish the handshake.
584        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1_retry.to_rx());
585        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
586        let key1 = handshake.sta1_handle_confirm(confirm2.to_rx());
587        let key2 = handshake.sta2_handle_confirm(confirm1.to_rx());
588        assert_eq!(key1, key2);
589    }
590
591    #[test]
592    fn ignore_wrong_confirm() {
593        let mut handshake = TestHandshake::new();
594
595        let commit1 = handshake.sta1_init();
596        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
597        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
598
599        let mut sink = vec![];
600        let confirm2_wrong = ConfirmTx(frame::write_confirm(1, &[1; 32][..]));
601        handshake.sta1.handle_confirm(&mut sink, &confirm2_wrong.to_rx().msg());
602        assert_eq!(sink.len(), 0); // Ignored.
603
604        // STA1 should still be able to handle a subsequent correct confirm.
605        handshake.sta1_handle_confirm(confirm2.to_rx());
606    }
607
608    #[test]
609    fn handle_resent_commit() {
610        let mut handshake = TestHandshake::new();
611        let commit1 = handshake.sta1_init();
612        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
613        let (commit2_retry, confirm2_retry) = handshake.sta2_handle_commit(commit1.to_rx());
614
615        // The resent commit message should be unchanged, but the resent confirm should increment
616        // sc and produce a different value.
617        assert_eq!(commit2, commit2_retry);
618        assert_eq!(confirm2.to_rx().msg().send_confirm, 1);
619        assert_eq!(confirm2_retry.to_rx().msg().send_confirm, 2);
620        assert!(confirm2.to_rx().msg().confirm != confirm2_retry.to_rx().msg().confirm);
621
622        // Now complete the handshake.
623        let confirm1 = handshake.sta1_handle_commit(commit2_retry.to_rx());
624        let key1 = handshake.sta1_handle_confirm(confirm2_retry.to_rx());
625        let key2 = handshake.sta2_handle_confirm(confirm1.to_rx());
626        assert_eq!(key1, key2);
627    }
628
629    #[test]
630    fn completed_handshake_handles_resent_confirm() {
631        let mut handshake = TestHandshake::new();
632        let commit1 = handshake.sta1_init();
633        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
634        let (commit2_retry, confirm2_retry) = handshake.sta2_handle_commit(commit1.to_rx());
635        // Send STA1 the second confirm message first.
636        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
637        let key1 = handshake.sta1_handle_confirm(confirm2.clone().to_rx());
638
639        // STA1 should respond to the second confirm with its own confirm.
640        let mut sink = vec![];
641        handshake.sta1.handle_confirm(&mut sink, &confirm2_retry.to_rx().msg());
642        assert_eq!(sink.len(), 1);
643        let confirm1_retry = expect_confirm(&mut sink);
644        assert!(confirm1.to_rx().msg().confirm != confirm1_retry.to_rx().msg().confirm);
645        assert_eq!(confirm1_retry.to_rx().msg().send_confirm, u16::max_value());
646
647        // STA2 should complete the handshake with the resent confirm.
648        let key2 = handshake.sta2_handle_confirm(confirm1_retry.to_rx());
649        assert_eq!(key1, key2);
650
651        // STA1 should silently drop either of our confirm frames now.
652        handshake.sta1.handle_confirm(&mut sink, &confirm2_retry.to_rx().msg());
653        assert!(sink.is_empty());
654        handshake.sta1.handle_confirm(&mut sink, &confirm2.to_rx().msg());
655        assert!(sink.is_empty());
656
657        // STA1 should also silently drop an incorrect confirm, even if send_confirm is incremented.
658        let confirm2_wrong = ConfirmMsg { send_confirm: 10, confirm: &[0xab; 32][..] };
659        handshake.sta1.handle_confirm(&mut sink, &confirm2_wrong);
660        assert!(sink.is_empty());
661    }
662
663    #[test]
664    fn completed_handshake_ignores_commit() {
665        let mut handshake = TestHandshake::new();
666        let commit1 = handshake.sta1_init();
667        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
668        handshake.sta1_handle_commit(commit2.to_rx());
669        handshake.sta1_handle_confirm(confirm2.clone().to_rx());
670
671        // STA1 has completed it's side of the handshake.
672        let mut sink = vec![];
673        handshake.sta1.handle_confirm(&mut sink, &confirm2.to_rx().msg());
674        assert!(sink.is_empty());
675    }
676
677    #[test]
678    fn bad_first_commit_rejects_auth() {
679        let mut handshake = TestHandshake::new();
680        let commit1_wrong = CommitMsg {
681            group_id: 19,
682            scalar: &[0xab; 32][..],
683            element: &[0xcd; 64][..],
684            anti_clogging_token: None,
685        };
686
687        let mut sink = vec![];
688        handshake.sta1.handle_commit(&mut sink, &commit1_wrong);
689        assert_eq!(sink.len(), 1);
690        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::AuthFailed));
691    }
692
693    #[test]
694    fn bad_second_commit_ignored() {
695        let mut handshake = TestHandshake::new();
696        let commit1 = handshake.sta1_init();
697        let (_commit1, _confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
698        let commit2_wrong = CommitMsg {
699            group_id: 19,
700            scalar: &[0xab; 32][..],
701            element: &[0xcd; 64][..],
702            anti_clogging_token: None,
703        };
704        let mut sink = vec![];
705        handshake.sta1.handle_commit(&mut sink, &commit2_wrong);
706        assert_eq!(sink.len(), 0);
707    }
708
709    #[test]
710    fn reflected_commit_discarded() {
711        let mut handshake = TestHandshake::new();
712        let commit1 = handshake.sta1_init();
713
714        let mut sink = vec![];
715        handshake.sta1.handle_commit(&mut sink, &commit1.to_rx().msg());
716        assert_eq!(sink.len(), 1);
717        assert_matches!(sink.remove(0), SaeUpdate::ResetTimeout(Timeout::Retransmission));
718    }
719
720    #[test]
721    fn maximum_commit_retries() {
722        let mut handshake = TestHandshake::new();
723        let commit1 = handshake.sta1_init();
724        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
725
726        // STA2 should allow MAX_RETRIES_PER_EXCHANGE retry operations before giving up.
727        for i in 0..MAX_RETRIES_PER_EXCHANGE {
728            let (commit2_retry, confirm2_retry) =
729                handshake.sta2_handle_commit(commit1.clone().to_rx());
730            assert_eq!(commit2, commit2_retry);
731            assert_eq!(confirm2_retry.to_rx().msg().send_confirm, i + 2);
732        }
733
734        // The last straw!
735        let mut sink = vec![];
736        handshake.sta2.handle_commit(&mut sink, &commit1.to_rx().msg());
737        assert_eq!(sink.len(), 1);
738        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::TooManyRetries));
739    }
740
741    #[test]
742    fn completed_exchange_fails_after_retries() {
743        let mut handshake = TestHandshake::new();
744        let commit1 = handshake.sta1_init();
745        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
746
747        // STA2 should allow MAX_RETRIES_PER_EXCHANGE retry operations before giving up. We subtract 1
748        // here for the reason explained in the note below.
749        for i in 0..(MAX_RETRIES_PER_EXCHANGE - 1) {
750            let (commit2_retry, confirm2_retry) =
751                handshake.sta2_handle_commit(commit1.clone().to_rx());
752            assert_eq!(commit2, commit2_retry);
753            assert_eq!(confirm2_retry.to_rx().msg().send_confirm, i + 2);
754        }
755
756        let mut sink = vec![];
757
758        // Generate 3 different confirm messages for our testing...
759        let confirm1_sc1 = handshake.sta1_handle_commit(commit2.clone().to_rx());
760        handshake.sta1.handle_commit(&mut sink, &commit2.to_rx().msg());
761        assert_eq!(sink.len(), 3);
762        sink.remove(0);
763        let confirm1_sc2 = expect_confirm(&mut sink);
764        sink.clear();
765        let confirm1_invalid = ConfirmMsg { send_confirm: 3, confirm: &[0xab; 32][..] };
766
767        // STA2 completes the handshake. However, one more indication that STA1 is misbehaving will
768        // immediately kill the authentication.
769        handshake.sta2_handle_confirm(confirm1_sc1.clone().to_rx());
770
771        // NOTE: We run all of the operations here two times. This is because of a quirk in the SAE
772        // state machine: while only certain operations *increment* sync, all invalid operations
773        // will *check* sync. We can test whether sync is being incremented by running twice to see
774        // if this pushes us over the MAX_RETRIES_PER_EXCHANGE threshold.
775
776        // STA2 ignores commits.
777        handshake.sta2.handle_commit(&mut sink, &commit1.to_rx().msg());
778        handshake.sta2.handle_commit(&mut sink, &commit1.to_rx().msg());
779        assert_eq!(sink.len(), 0);
780
781        // STA2 ignores invalid confirm.
782        handshake.sta2.handle_confirm(&mut sink, &confirm1_invalid);
783        handshake.sta2.handle_confirm(&mut sink, &confirm1_invalid);
784        assert_eq!(sink.len(), 0);
785
786        // STA2 ignores old confirm.
787        handshake.sta2.handle_confirm(&mut sink, &confirm1_sc1.to_rx().msg());
788        handshake.sta2.handle_confirm(&mut sink, &confirm1_sc1.to_rx().msg());
789        assert_eq!(sink.len(), 0);
790
791        // But another valid confirm increments sync!
792        handshake.sta2.handle_confirm(&mut sink, &confirm1_sc2.to_rx().msg());
793        assert_eq!(sink.len(), 1);
794        expect_confirm(&mut sink);
795        handshake.sta2.handle_confirm(&mut sink, &confirm1_sc2.to_rx().msg());
796        assert_eq!(sink.len(), 1);
797        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::TooManyRetries));
798    }
799
800    #[test]
801    fn resend_commit_after_retransmission_timeout() {
802        let mut handshake = TestHandshake::new();
803        let commit1 = handshake.sta1_init();
804
805        let mut sink = vec![];
806        handshake.sta1.handle_timeout(&mut sink, Timeout::Retransmission);
807        let commit1_retry = expect_commit(&mut sink);
808        expect_reset_timeout(&mut sink, Timeout::Retransmission);
809        assert_eq!(commit1, commit1_retry);
810    }
811
812    #[test]
813    fn resend_confirm_after_retransmission_timeout() {
814        let mut handshake = TestHandshake::new();
815        let commit1 = handshake.sta1_init();
816        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
817
818        let mut sink = vec![];
819        handshake.sta2.handle_timeout(&mut sink, Timeout::Retransmission);
820        // On timeout we should only send commit and confirm.
821        let confirm2_retry = expect_confirm(&mut sink);
822        expect_reset_timeout(&mut sink, Timeout::Retransmission);
823        assert_eq!(
824            confirm2.to_rx().msg().send_confirm + 1,
825            confirm2_retry.to_rx().msg().send_confirm
826        );
827    }
828
829    #[test]
830    fn abort_commit_after_too_many_timeouts() {
831        let mut handshake = TestHandshake::new();
832        let commit1 = handshake.sta1_init();
833
834        let mut sink = vec![];
835        for i in 0..MAX_RETRIES_PER_EXCHANGE {
836            handshake.sta1.handle_timeout(&mut sink, Timeout::Retransmission);
837            let commit1_retry = expect_commit(&mut sink);
838            expect_reset_timeout(&mut sink, Timeout::Retransmission);
839            assert_eq!(commit1, commit1_retry);
840        }
841
842        // This camel can't hold another straw!
843        handshake.sta1.handle_timeout(&mut sink, Timeout::Retransmission);
844        assert_eq!(sink.len(), 1);
845        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::TooManyRetries));
846    }
847
848    #[test]
849    fn abort_confirm_after_too_many_timeouts() {
850        let mut handshake = TestHandshake::new();
851        let commit1 = handshake.sta1_init();
852        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.clone().to_rx());
853
854        let mut sink = vec![];
855        for i in 0..MAX_RETRIES_PER_EXCHANGE {
856            handshake.sta2.handle_timeout(&mut sink, Timeout::Retransmission);
857            // On timeout we should only send commit and confirm.
858            let confirm2_retry = expect_confirm(&mut sink);
859            expect_reset_timeout(&mut sink, Timeout::Retransmission);
860            assert_eq!(
861                confirm2.to_rx().msg().send_confirm + i + 1,
862                confirm2_retry.to_rx().msg().send_confirm
863            );
864        }
865
866        handshake.sta2.handle_timeout(&mut sink, Timeout::Retransmission);
867        assert_eq!(sink.len(), 1);
868        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::TooManyRetries));
869    }
870
871    #[test]
872    fn ignore_unexpected_retransmit_timeout() {
873        let mut handshake = TestHandshake::new();
874        let mut sink = vec![];
875        // Timeout::Retransmission is ignored while in New state.
876        handshake.sta1.handle_timeout(&mut sink, Timeout::Retransmission);
877        assert!(sink.is_empty());
878
879        let commit1 = handshake.sta1_init();
880        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
881        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
882        let key1 = handshake.sta1_handle_confirm(confirm2.to_rx());
883
884        // Timeout::Retransmission is ignored while in Accepted state.
885        handshake.sta1.handle_timeout(&mut sink, Timeout::Retransmission);
886        assert!(sink.is_empty());
887    }
888
889    #[test]
890    fn fail_on_early_key_expiration() {
891        let mut handshake = TestHandshake::new();
892        handshake.sta1_init();
893
894        // Early key expiration indicates that something has gone very wrong, so we abort.
895        let mut sink = vec![];
896        handshake.sta1.handle_timeout(&mut sink, Timeout::KeyExpiration);
897        assert_eq!(sink.len(), 1);
898        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::InternalError(_)));
899    }
900
901    #[test]
902    fn key_expiration_timeout() {
903        let mut handshake = TestHandshake::new();
904        // Timeout::KeyExpiration is only expected once our handshake has completed.
905        let commit1 = handshake.sta1_init();
906        let (commit2, confirm2) = handshake.sta2_handle_commit(commit1.to_rx());
907        let confirm1 = handshake.sta1_handle_commit(commit2.to_rx());
908        let key1 = handshake.sta1_handle_confirm(confirm2.to_rx());
909
910        let mut sink = vec![];
911        handshake.sta1.handle_timeout(&mut sink, Timeout::KeyExpiration);
912        assert_eq!(sink.len(), 1);
913        assert_matches!(sink.remove(0), SaeUpdate::Reject(RejectReason::KeyExpiration));
914    }
915}