diff --git a/src/olm/session/double_ratchet.rs b/src/olm/session/double_ratchet.rs index 124cf0a1..85c59a61 100644 --- a/src/olm/session/double_ratchet.rs +++ b/src/olm/session/double_ratchet.rs @@ -294,7 +294,7 @@ impl Debug for ActiveDoubleRatchet { /// /// It may be unknown, if the ratchet was restored from a pickle /// which didn't track it. -#[derive(Serialize, Deserialize, Clone)] +#[derive(Serialize, Deserialize, Clone, PartialEq, Eq)] pub enum RatchetCount { Known(u64), Unknown(()), @@ -325,3 +325,133 @@ impl Debug for RatchetCount { } } } + +#[cfg(test)] +mod test { + use assert_matches::assert_matches; + + use super::{ + ActiveDoubleRatchet, DoubleRatchet, DoubleRatchetState, InactiveDoubleRatchet, RatchetCount, + }; + use crate::olm::{ + session::test::session_and_libolm_pair, Account, OlmMessage, Session, SessionConfig, + }; + + fn create_session_pair(alice: &Account, bob: &mut Account) -> (Session, Session) { + let bob_otks = bob.generate_one_time_keys(1); + let bob_otk = bob_otks.created.get(0).expect("Couldn't get a one-time-key for bob"); + let bob_identity_key = bob.identity_keys().curve25519; + let mut alice_session = alice.create_outbound_session( + SessionConfig::version_1(), + bob_identity_key, + bob_otk.clone(), + ); + + let message = "It's a secret to everybody"; + let olm_message = alice_session.encrypt(message); + let prekey_message = assert_matches!(olm_message, OlmMessage::PreKey(m) => m); + + let alice_identity_key = alice.identity_keys().curve25519; + let bob_session_creation_result = bob + .create_inbound_session(alice_identity_key.clone(), &prekey_message) + .expect("Unable to create inbound session"); + assert_eq!(bob_session_creation_result.plaintext, message.as_bytes()); + (alice_session, bob_session_creation_result.session) + } + + fn assert_active_ratchet(sending_ratchet: &DoubleRatchet) -> &ActiveDoubleRatchet { + match &sending_ratchet.inner { + DoubleRatchetState::Inactive(_) => panic!("Not an active ratchet"), + DoubleRatchetState::Active(s) => s, + } + } + + fn assert_inactive_ratchet(sending_ratchet: &DoubleRatchet) -> &InactiveDoubleRatchet { + match &sending_ratchet.inner { + DoubleRatchetState::Active(_) => panic!("Not an inactive ratchet"), + DoubleRatchetState::Inactive(s) => s, + } + } + + #[test] + fn ratchet_counts() { + let (mut alice_session, mut bob_session) = + create_session_pair(&Account::new(), &mut Account::new()); + + // Both ratchets should start with count 0 + assert_eq!( + assert_active_ratchet(&alice_session.sending_ratchet).ratchet_count, + RatchetCount::Known(0) + ); + assert_eq!( + assert_inactive_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Known(0) + ); + + // Once Bob replies, the ratchets should bump to 1 + let olm_message = bob_session.encrypt("sssh"); + alice_session.decrypt(&olm_message).expect("Alice could not decrypt message from Bob"); + assert_eq!( + assert_inactive_ratchet(&alice_session.sending_ratchet).ratchet_count, + RatchetCount::Known(1) + ); + assert_eq!( + assert_active_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Known(1) + ); + + // Now Alice replies again + let olm_message = alice_session.encrypt("sssh"); + bob_session.decrypt(&olm_message).expect("Bob could not decrypt message from Alice"); + assert_eq!( + assert_active_ratchet(&alice_session.sending_ratchet).ratchet_count, + RatchetCount::Known(2) + ); + assert_eq!( + assert_inactive_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Known(2) + ); + } + + #[test] + fn ratchet_counts_for_imported_session() { + let (_, _, mut alice_session, bob_libolm_session) = + session_and_libolm_pair().expect("unable to create sessions"); + + // Import the libolm session into a proper vodozmac session + let key = b"DEFAULT_PICKLE_KEY"; + let pickle = + bob_libolm_session.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() }); + let mut bob_session = + Session::from_libolm_pickle(&pickle, key).expect("Should be able to unpickle session"); + + assert_eq!( + assert_inactive_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Unknown(()) + ); + + // Once Bob replies, Alice's count bumps to 1, but Bob's remains unknown + let olm_message = bob_session.encrypt("sssh"); + alice_session.decrypt(&olm_message).expect("Alice could not decrypt message from Bob"); + assert_eq!( + assert_inactive_ratchet(&alice_session.sending_ratchet).ratchet_count, + RatchetCount::Known(1) + ); + assert_eq!( + assert_active_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Unknown(()) + ); + + // Now Alice replies again + let olm_message = alice_session.encrypt("sssh"); + bob_session.decrypt(&olm_message).expect("Bob could not decrypt message from Alice"); + assert_eq!( + assert_active_ratchet(&alice_session.sending_ratchet).ratchet_count, + RatchetCount::Known(2) + ); + assert_eq!( + assert_inactive_ratchet(&bob_session.sending_ratchet).ratchet_count, + RatchetCount::Unknown(()) + ); + } +} diff --git a/src/olm/session/mod.rs b/src/olm/session/mod.rs index 87a5fd78..3d55da0c 100644 --- a/src/olm/session/mod.rs +++ b/src/olm/session/mod.rs @@ -534,7 +534,10 @@ mod test { const PICKLE_KEY: [u8; 32] = [0u8; 32]; - fn sessions() -> Result<(Account, OlmAccount, Session, OlmSession)> { + /// Create a pair of accounts, one using vodozemac and one libolm. + /// + /// Then, create a pair of sessions between the two. + pub fn session_and_libolm_pair() -> Result<(Account, OlmAccount, Session, OlmSession)> { let alice = Account::new(); let bob = OlmAccount::new(); bob.generate_one_time_keys(1); @@ -570,7 +573,7 @@ mod test { #[test] fn out_of_order_decryption() { - let (_, _, mut alice_session, bob_session) = sessions().unwrap(); + let (_, _, mut alice_session, bob_session) = session_and_libolm_pair().unwrap(); let message_1 = bob_session.encrypt("Message 1").into(); let message_2 = bob_session.encrypt("Message 2").into(); @@ -592,7 +595,7 @@ mod test { #[test] fn more_out_of_order_decryption() { - let (_, _, mut alice_session, bob_session) = sessions().unwrap(); + let (_, _, mut alice_session, bob_session) = session_and_libolm_pair().unwrap(); let message_1 = bob_session.encrypt("Message 1").into(); let message_2 = bob_session.encrypt("Message 2").into(); @@ -630,7 +633,7 @@ mod test { #[test] fn max_keys_out_of_order_decryption() { - let (_, _, mut alice_session, bob_session) = sessions().unwrap(); + let (_, _, mut alice_session, bob_session) = session_and_libolm_pair().unwrap(); let mut messages: Vec<messages::OlmMessage> = Vec::new(); for i in 0..(MAX_MESSAGE_KEYS + 2) { @@ -664,7 +667,7 @@ mod test { #[test] fn max_gap_out_of_order_decryption() { - let (_, _, mut alice_session, bob_session) = sessions().unwrap(); + let (_, _, mut alice_session, bob_session) = session_and_libolm_pair().unwrap(); for i in 0..(MAX_MESSAGE_GAP + 1) { bob_session.encrypt(format!("Message {}", i).as_str()); @@ -680,7 +683,7 @@ mod test { #[test] #[cfg(feature = "libolm-compat")] fn libolm_unpickling() { - let (_, _, mut session, olm) = sessions().unwrap(); + let (_, _, mut session, olm) = session_and_libolm_pair().unwrap(); let plaintext = "It's a secret to everybody"; let old_message = session.encrypt(plaintext); @@ -717,7 +720,7 @@ mod test { #[test] fn session_pickling_roundtrip_is_identity() { - let (_, _, session, _) = sessions().unwrap(); + let (_, _, session, _) = session_and_libolm_pair().unwrap(); let pickle = session.pickle().encrypt(&PICKLE_KEY);