Skip to main content

vodozemac/olm/account/
mod.rs

1// Copyright 2021 Damir Jelić, Denis Kasak
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod fallback_keys;
16mod one_time_keys;
17
18use std::collections::HashMap;
19
20use chacha20poly1305::{
21    ChaCha20Poly1305, Nonce,
22    aead::{Aead, KeyInit},
23};
24use cipher::common::Generate;
25use serde::{Deserialize, Serialize};
26use thiserror::Error;
27use zeroize::Zeroize;
28
29pub use self::one_time_keys::OneTimeKeyGenerationResult;
30use self::{
31    fallback_keys::FallbackKeys,
32    one_time_keys::{OneTimeKeys, OneTimeKeysPickle},
33};
34use super::{
35    SessionConfig,
36    messages::PreKeyMessage,
37    session::{DecryptionError, Session},
38    session_keys::SessionKeys,
39    shared_secret::{RemoteShared3DHSecret, Shared3DHSecret},
40};
41use crate::{
42    Ed25519Signature, PickleError, olm,
43    types::{
44        Curve25519Keypair, Curve25519KeypairPickle, Curve25519PublicKey, Curve25519SecretKey,
45        Ed25519Keypair, Ed25519KeypairPickle, Ed25519PublicKey, KeyId,
46    },
47    utilities::{pickle, unpickle},
48};
49
50const PUBLIC_MAX_ONE_TIME_KEYS: usize = 50;
51
52/// Error describing failure modes when creating a Olm [`Session`] from an
53/// incoming Olm message.
54#[derive(Error, Debug)]
55pub enum SessionCreationError {
56    /// The pre-key message contained an unknown one-time key. This happens
57    /// either because we never had such a one-time key, or because it has
58    /// already been used up.
59    #[error("The pre-key message contained an unknown one-time key: {0}")]
60    MissingOneTimeKey(Curve25519PublicKey),
61    /// The pre-key message contains a Curve25519 identity key that doesn't
62    /// match to the identity key that was given.
63    #[error(
64        "The given identity key doesn't match the one in the pre-key message: \
65        expected {0}, got {1}"
66    )]
67    MismatchedIdentityKey(Curve25519PublicKey, Curve25519PublicKey),
68    /// The pre-key message was encrypted with a Session which used an
69    /// unexpected SessionConfig.
70    #[error(
71        "The session config doesn't match the one used for the pre-key message: expected {expected:?}, got {got:?}"
72    )]
73    MismatchedSessionConfig {
74        /// The [`SessionConfig`] we expected.
75        expected: SessionConfig,
76        /// The [`SessionConfig`] the pre-key message was encrypted with.
77        ///
78        /// Will be `None` if we don't understand the received config.
79        got: Option<SessionConfig>,
80    },
81    /// The pre-key message that was used to establish the [`Session`] couldn't
82    /// be decrypted. The message needs to be decryptable, otherwise we will
83    /// have created a Session that wasn't used to encrypt the pre-key
84    /// message.
85    #[error("The message that was used to establish the Session couldn't be decrypted")]
86    Decryption(#[from] DecryptionError),
87    /// One or more keys lacked contributory behavior in the Diffie-Hellman
88    /// operation, resulting in an insecure shared secret.
89    ///
90    /// For more details on contributory behavior please refer to the
91    /// [`x25519_dalek::SharedSecret::was_contributory()`] method.
92    #[error(
93        "The shared secret was derived from an insecure key exchange (non-contributory behaviour)"
94    )]
95    NonContributoryKey,
96}
97
98/// Struct holding the two public identity keys of an [`Account`].
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
100pub struct IdentityKeys {
101    /// The Ed25519 key, used for signing.
102    pub ed25519: Ed25519PublicKey,
103    /// The Curve25519 key, used for to establish shared secrets.
104    pub curve25519: Curve25519PublicKey,
105}
106
107/// Return type for the creation of inbound [`Session`] objects.
108#[derive(Debug)]
109pub struct InboundCreationResult {
110    /// The [`Session`] that was created from a pre-key message.
111    pub session: Session,
112    /// The plaintext of the pre-key message.
113    pub plaintext: Vec<u8>,
114}
115
116/// Return type for the creation of a dehydrated device.
117#[derive(Debug)]
118pub struct DehydratedDeviceResult {
119    /// The encrypted dehydrated device, as a base64-encoded string.
120    pub ciphertext: String,
121    /// The nonce used for encrypting, as a base64-encoded string.
122    pub nonce: String,
123}
124
125/// An Olm [`Account`] manages all cryptographic keys used on a device.
126pub struct Account {
127    /// A permanent Ed25519 key used for signing. Also known as the fingerprint
128    /// key.
129    signing_key: Ed25519Keypair,
130    /// The permanent Curve25519 key used for triple Diffie-Hellman (3DH). Also
131    /// known as the sender key or the identity key.
132    diffie_hellman_key: Curve25519Keypair,
133    /// The ephemeral (one-time) Curve25519 keys used as part of the triple
134    /// Diffie-Hellman (3DH).
135    one_time_keys: OneTimeKeys,
136    /// The ephemeral Curve25519 keys used in lieu of a one-time key as part of
137    /// the 3DH, in case we run out of those. We keep track of both the current
138    /// and the previous fallback key in any given moment.
139    fallback_keys: FallbackKeys,
140}
141
142impl Account {
143    /// Create a new [`Account`] with new random identity keys.
144    pub fn new() -> Self {
145        Self {
146            signing_key: Ed25519Keypair::new(),
147            diffie_hellman_key: Curve25519Keypair::new(),
148            one_time_keys: OneTimeKeys::new(),
149            fallback_keys: FallbackKeys::new(),
150        }
151    }
152
153    /// Get the [`IdentityKeys`] of this Account
154    pub const fn identity_keys(&self) -> IdentityKeys {
155        IdentityKeys { ed25519: self.ed25519_key(), curve25519: self.curve25519_key() }
156    }
157
158    /// Get a copy of the account's public Ed25519 key
159    pub const fn ed25519_key(&self) -> Ed25519PublicKey {
160        self.signing_key.public_key()
161    }
162
163    /// Get a copy of the account's public Curve25519 key
164    pub const fn curve25519_key(&self) -> Curve25519PublicKey {
165        self.diffie_hellman_key.public_key()
166    }
167
168    /// Sign the given message using our Ed25519 fingerprint key.
169    pub fn sign(&self, message: impl AsRef<[u8]>) -> Ed25519Signature {
170        self.signing_key.sign(message.as_ref())
171    }
172
173    /// Get the maximum number of one-time keys the client should keep on the
174    /// server.
175    ///
176    /// **Note**: this differs from the libolm method of the same name, the
177    /// libolm method returned the maximum amount of one-time keys the
178    /// [`Account`] could hold and only half of those should be uploaded.
179    pub const fn max_number_of_one_time_keys(&self) -> usize {
180        // We tell clients to upload a limited amount of one-time keys, this
181        // amount is smaller than what we can store.
182        //
183        // We do this because a client might receive the count of uploaded keys
184        // from the server before they receive all the pre-key messages that
185        // used some of our one-time keys. This would mean that we would forget
186        // private one-time keys, since we're generating new ones, while we
187        // didn't yet receive the pre-key messages that used those one-time
188        // keys.
189        PUBLIC_MAX_ONE_TIME_KEYS
190    }
191
192    /// Create a [`Session`] with the given identity key and one-time key.
193    pub fn create_outbound_session(
194        &self,
195        session_config: SessionConfig,
196        identity_key: Curve25519PublicKey,
197        one_time_key: Curve25519PublicKey,
198    ) -> Result<Session, SessionCreationError> {
199        let base_key = Curve25519SecretKey::new();
200        let public_base_key = Curve25519PublicKey::from(&base_key);
201
202        let shared_secret = Shared3DHSecret::new(
203            self.diffie_hellman_key.secret_key(),
204            &base_key,
205            &identity_key,
206            &one_time_key,
207        )
208        .ok_or(SessionCreationError::NonContributoryKey)?;
209
210        let session_keys = SessionKeys {
211            identity_key: self.curve25519_key(),
212            base_key: public_base_key,
213            one_time_key,
214        };
215
216        Ok(Session::new(session_config, shared_secret, session_keys))
217    }
218
219    /// Try to find a [`Curve25519SecretKey`] that forms a pair with the given
220    /// [`Curve25519PublicKey`].
221    fn find_one_time_key(&self, public_key: &Curve25519PublicKey) -> Option<&Curve25519SecretKey> {
222        self.one_time_keys
223            .get_secret_key(public_key)
224            .or_else(|| self.fallback_keys.get_secret_key(public_key))
225    }
226
227    /// Remove a one-time key that has previously been published but not yet
228    /// used.
229    ///
230    /// **Note**: This function is only rarely useful and you'll know if you
231    /// need it. Notably, you do *not* need to call it manually when using up
232    /// a key via [`Account::create_inbound_session`] since the key is
233    /// automatically removed in that case.
234    #[cfg(feature = "low-level-api")]
235    pub fn remove_one_time_key(
236        &mut self,
237        public_key: Curve25519PublicKey,
238    ) -> Option<Curve25519SecretKey> {
239        self.remove_one_time_key_helper(public_key)
240    }
241
242    fn remove_one_time_key_helper(
243        &mut self,
244        public_key: Curve25519PublicKey,
245    ) -> Option<Curve25519SecretKey> {
246        self.one_time_keys.remove_secret_key(&public_key)
247    }
248
249    /// Create a [`Session`] from the given [`PreKeyMessage`] message and
250    /// identity key
251    pub fn create_inbound_session(
252        &mut self,
253        expected_config: SessionConfig,
254        their_identity_key: Curve25519PublicKey,
255        pre_key_message: &PreKeyMessage,
256    ) -> Result<InboundCreationResult, SessionCreationError> {
257        if their_identity_key != pre_key_message.identity_key() {
258            Err(SessionCreationError::MismatchedIdentityKey(
259                their_identity_key,
260                pre_key_message.identity_key(),
261            ))
262        } else {
263            let config = match pre_key_message.message.version() {
264                olm::messages::message::MAC_TRUNCATED_VERSION => SessionConfig::version_1(),
265                #[cfg(feature = "experimental-session-config")]
266                olm::messages::message::VERSION => SessionConfig::version_2(),
267                _ => {
268                    return Err(SessionCreationError::MismatchedSessionConfig {
269                        expected: expected_config,
270                        got: None,
271                    });
272                }
273            };
274
275            if config != expected_config {
276                return Err(SessionCreationError::MismatchedSessionConfig {
277                    expected: expected_config,
278                    got: Some(config),
279                });
280            }
281
282            // Find the matching private part of the OTK that the message claims
283            // was used to create the session that encrypted it.
284            let public_otk = pre_key_message.one_time_key();
285            let private_otk = self
286                .find_one_time_key(&public_otk)
287                .ok_or(SessionCreationError::MissingOneTimeKey(public_otk))?;
288
289            // Construct a 3DH shared secret from the various curve25519 keys.
290            let shared_secret = RemoteShared3DHSecret::new(
291                self.diffie_hellman_key.secret_key(),
292                private_otk,
293                &pre_key_message.identity_key(),
294                &pre_key_message.base_key(),
295            )
296            .ok_or(SessionCreationError::NonContributoryKey)?;
297
298            // These will be used to uniquely identify the Session.
299            let session_keys = SessionKeys {
300                identity_key: pre_key_message.identity_key(),
301                base_key: pre_key_message.base_key(),
302                one_time_key: pre_key_message.one_time_key(),
303            };
304
305            // Create a Session, AKA a double ratchet, this one will have an
306            // inactive sending chain until we decide to encrypt a message.
307            let mut session = Session::new_remote(
308                config,
309                shared_secret,
310                pre_key_message.message.ratchet_key,
311                session_keys,
312            );
313
314            // Decrypt the message to check if the Session is actually valid.
315            let plaintext = session.decrypt_decoded(&pre_key_message.message)?;
316
317            // We only drop the one-time key now, this is why we can't use a
318            // one-time key type that takes `self`. If we didn't do this,
319            // someone could maliciously pretend to use up our one-time key and
320            // make us drop the private part. Unsuspecting users that actually
321            // try to use such an one-time key won't be able to communicate with
322            // us. This is strictly worse than the one-time key exhaustion
323            // scenario.
324            self.remove_one_time_key_helper(pre_key_message.one_time_key());
325
326            Ok(InboundCreationResult { session, plaintext })
327        }
328    }
329
330    /// Generates the supplied number of one time keys.
331    /// Returns the public parts of the one-time keys that were created and
332    /// discarded.
333    ///
334    /// Our one-time key store inside the [`Account`] has a limited amount of
335    /// places for one-time keys, If we try to generate new ones while the store
336    /// is completely populated, the oldest one-time keys will get discarded
337    /// to make place for new ones.
338    pub fn generate_one_time_keys(&mut self, count: usize) -> OneTimeKeyGenerationResult {
339        self.one_time_keys.generate(count)
340    }
341
342    /// Get the number of one-time keys we have stored locally.
343    ///
344    /// This will be equal or greater to the number of one-time keys we have
345    /// published. Each time a new [`Session`] is created using the
346    /// [`Account::create_inbound_session()`] a one-time key will be used up
347    /// and removed.
348    pub fn stored_one_time_key_count(&self) -> usize {
349        self.one_time_keys.private_keys.len()
350    }
351
352    /// Get the currently unpublished one-time keys.
353    ///
354    /// The one-time keys should be published to a server and marked as
355    /// published using the `mark_keys_as_published()` method.
356    pub fn one_time_keys(&self) -> HashMap<KeyId, Curve25519PublicKey> {
357        self.one_time_keys
358            .unpublished_public_keys
359            .iter()
360            .map(|(key_id, key)| (*key_id, *key))
361            .collect()
362    }
363
364    /// Generate a single new fallback key.
365    ///
366    /// The fallback key will be used by other users to establish a [`Session`]
367    /// if all the one-time keys on the server have been used up.
368    ///
369    /// Returns the public Curve25519 key of the *previous* fallback key, that
370    /// is, the one that will get removed from the [`Account`] when this method
371    /// is called. This return value is mostly useful for logging purposes.
372    pub fn generate_fallback_key(&mut self) -> Option<Curve25519PublicKey> {
373        self.fallback_keys.generate_fallback_key()
374    }
375
376    /// Get the currently unpublished fallback key.
377    ///
378    /// The fallback key should be published just like the one-time keys, after
379    /// it has been successfully published it needs to be marked as published
380    /// using the `mark_keys_as_published()` method as well.
381    pub fn fallback_key(&self) -> HashMap<KeyId, Curve25519PublicKey> {
382        let fallback_key = self.fallback_keys.unpublished_fallback_key();
383
384        if let Some(fallback_key) = fallback_key {
385            HashMap::from([(fallback_key.key_id(), fallback_key.public_key())])
386        } else {
387            HashMap::new()
388        }
389    }
390
391    /// The [`Account`] stores at most two private parts of the fallback key.
392    /// This method lets us forget the previously used fallback key.
393    pub fn forget_fallback_key(&mut self) -> bool {
394        self.fallback_keys.forget_previous_fallback_key().is_some()
395    }
396
397    /// Mark all currently unpublished one-time and fallback keys as published.
398    pub fn mark_keys_as_published(&mut self) {
399        self.one_time_keys.mark_as_published();
400        self.fallback_keys.mark_as_published();
401    }
402
403    /// Convert the account into a struct which implements [`serde::Serialize`]
404    /// and [`serde::Deserialize`].
405    pub fn pickle(&self) -> AccountPickle {
406        AccountPickle {
407            signing_key: self.signing_key.clone().into(),
408            diffie_hellman_key: self.diffie_hellman_key.clone().into(),
409            one_time_keys: self.one_time_keys.clone().into(),
410            fallback_keys: self.fallback_keys.clone(),
411        }
412    }
413
414    /// Restore an [`Account`] from a previously saved [`AccountPickle`].
415    pub fn from_pickle(pickle: AccountPickle) -> Self {
416        pickle.into()
417    }
418
419    /// Create an [`Account`] object by unpickling an account pickle in libolm
420    /// legacy pickle format.
421    ///
422    /// Such pickles are encrypted and need to first be decrypted using
423    /// `pickle_key`.
424    #[cfg(feature = "libolm-compat")]
425    pub fn from_libolm_pickle(
426        pickle: &str,
427        pickle_key: &[u8],
428    ) -> Result<Self, crate::LibolmPickleError> {
429        use self::libolm::Pickle;
430        use crate::utilities::unpickle_libolm;
431
432        const PICKLE_VERSION: u32 = 4;
433        unpickle_libolm::<Pickle, _>(pickle, pickle_key, PICKLE_VERSION)
434    }
435
436    /// Pickle an [`Account`] into a libolm pickle format.
437    ///
438    /// This pickle can be restored using the [`Account::from_libolm_pickle()`]
439    /// method, or can be used in the [`libolm`] C library.
440    ///
441    /// The pickle will be encrypted using the pickle key.
442    ///
443    /// *Note*: This method might be lossy, the vodozemac [`Account`] has the
444    /// ability to hold more one-time keys compared to the [`libolm`]
445    /// variant.
446    ///
447    /// ⚠️  ***Security Warning***: The pickle key will get expanded into both
448    /// an AES key and an IV in a deterministic manner. If the same pickle
449    /// key is reused, this will lead to IV reuse. To prevent this, users
450    /// have to ensure that they always use a globally (probabilistically)
451    /// unique pickle key.
452    ///
453    /// [`libolm`]: https://gitlab.matrix.org/matrix-org/olm/
454    ///
455    /// # Examples
456    /// ```
457    /// use vodozemac::olm::Account;
458    /// use olm_rs::{account::OlmAccount, PicklingMode};
459    /// let account = Account::new();
460    ///
461    /// let export = account
462    ///     .to_libolm_pickle(&[0u8; 32])
463    ///     .expect("We should be able to pickle a freshly created Account");
464    ///
465    /// let unpickled = OlmAccount::unpickle(
466    ///     export,
467    ///     PicklingMode::Encrypted { key: [0u8; 32].to_vec() },
468    /// ).expect("We should be able to unpickle our exported Account");
469    /// ```
470    #[cfg(feature = "libolm-compat")]
471    pub fn to_libolm_pickle(&self, pickle_key: &[u8]) -> Result<String, crate::LibolmPickleError> {
472        use self::libolm::Pickle;
473        use crate::utilities::pickle_libolm;
474        pickle_libolm::<Pickle>(self.into(), pickle_key)
475    }
476
477    #[cfg(all(any(fuzzing, test), feature = "libolm-compat"))]
478    #[doc(hidden)]
479    pub fn from_decrypted_libolm_pickle(pickle: &[u8]) -> Result<Self, crate::LibolmPickleError> {
480        use std::io::Cursor;
481
482        use matrix_pickle::Decode;
483
484        use self::libolm::Pickle;
485
486        let mut cursor = Cursor::new(&pickle);
487        let pickle = Pickle::decode(&mut cursor)?;
488
489        pickle.try_into()
490    }
491
492    /// Create a dehydrated device from the account.
493    ///
494    /// A dehydrated device is a device that is stored encrypted on the server
495    /// that can receive messages when the user has no other active devices.
496    /// Upon login, the user can rehydrate the device (using
497    /// [`Account::from_dehydrated_device`]) and decrypt the messages sent to
498    /// the dehydrated device.
499    ///
500    /// The account must be a newly-created account that does not have any Olm
501    /// sessions, since the dehydrated device format does not store sessions.
502    ///
503    /// Returns the ciphertext and nonce.  `key` is a 256-bit (32-byte) key for
504    /// encrypting the device.
505    ///
506    /// The format used here is defined in
507    /// [MSC3814](https://github.com/matrix-org/matrix-spec-proposals/pull/3814).
508    pub fn to_dehydrated_device(
509        &self,
510        key: &[u8; 32],
511    ) -> Result<DehydratedDeviceResult, crate::DehydratedDeviceError> {
512        use matrix_pickle::Encode;
513
514        use self::dehydrated_device::Pickle;
515        use crate::{DehydratedDeviceError, LibolmPickleError, utilities::base64_encode};
516
517        let pickle: Pickle = self.try_into()?;
518        let mut encoded = pickle
519            .encode_to_vec()
520            .map_err(|e| DehydratedDeviceError::LibolmPickle(LibolmPickleError::Encode(e)))?;
521
522        let cipher = ChaCha20Poly1305::new(key.into());
523        let nonce = Nonce::generate();
524        let ciphertext = cipher.encrypt(&nonce, encoded.as_slice());
525
526        encoded.zeroize();
527
528        let ciphertext = ciphertext?;
529
530        Ok(DehydratedDeviceResult {
531            ciphertext: base64_encode(ciphertext),
532            nonce: base64_encode(nonce),
533        })
534    }
535
536    /// Create an [`Account`] object from a dehydrated device.
537    ///
538    /// `ciphertext` and `nonce` are the ciphertext and nonce returned by
539    /// [`Account::to_dehydrated_device`]. `key` is a 256-bit (32-byte) key for
540    /// decrypting the device, and must be the same key used when
541    /// [`Account::to_dehydrated_device`] was called.
542    pub fn from_dehydrated_device(
543        ciphertext: &str,
544        nonce: &str,
545        key: &[u8; 32],
546    ) -> Result<Self, crate::DehydratedDeviceError> {
547        use self::dehydrated_device::PICKLE_VERSION;
548        use crate::utilities::{base64_decode, get_pickle_version};
549
550        let cipher = ChaCha20Poly1305::new(key.into());
551        let ciphertext = base64_decode(ciphertext)?;
552        let nonce = base64_decode(nonce)?;
553
554        if nonce.len() != 12 {
555            Err(crate::DehydratedDeviceError::InvalidNonce)
556        } else {
557            let mut nonce_array = [0u8; 12];
558            nonce_array.copy_from_slice(nonce.as_slice());
559
560            let nonce = Nonce::from(nonce_array);
561            let mut plaintext = cipher.decrypt(&nonce, ciphertext.as_slice())?;
562
563            let version = get_pickle_version(&plaintext)
564                .ok_or(crate::DehydratedDeviceError::MissingVersion)?;
565
566            if version != PICKLE_VERSION {
567                Err(crate::DehydratedDeviceError::Version(PICKLE_VERSION, version))
568            } else {
569                let pickle = Self::from_decrypted_dehydrated_device(&plaintext);
570                plaintext.zeroize();
571                pickle
572            }
573        }
574    }
575
576    // This function is public for fuzzing, but should not be used by anything
577    // else
578    #[doc(hidden)]
579    pub fn from_decrypted_dehydrated_device(
580        pickle: &[u8],
581    ) -> Result<Self, crate::DehydratedDeviceError> {
582        use std::io::Cursor;
583
584        use matrix_pickle::Decode;
585
586        use self::dehydrated_device::Pickle;
587        use crate::{DehydratedDeviceError, LibolmPickleError};
588
589        let mut cursor = Cursor::new(&pickle);
590        let pickle = Pickle::decode(&mut cursor)
591            .map_err(|e| DehydratedDeviceError::LibolmPickle(LibolmPickleError::Decode(e)))?;
592
593        pickle.try_into()
594    }
595}
596
597impl Default for Account {
598    fn default() -> Self {
599        Self::new()
600    }
601}
602
603/// A format suitable for serialization which implements [`serde::Serialize`]
604/// and [`serde::Deserialize`]. Obtainable by calling [`Account::pickle`].
605#[derive(Serialize, Deserialize)]
606pub struct AccountPickle {
607    signing_key: Ed25519KeypairPickle,
608    diffie_hellman_key: Curve25519KeypairPickle,
609    one_time_keys: OneTimeKeysPickle,
610    fallback_keys: FallbackKeys,
611}
612
613/// A format suitable for serialization which implements [`serde::Serialize`]
614/// and [`serde::Deserialize`]. Obtainable by calling [`Account::pickle`].
615impl AccountPickle {
616    /// Serialize and encrypt the pickle using the given key.
617    ///
618    /// This is the inverse of [`AccountPickle::from_encrypted`].
619    pub fn encrypt(self, pickle_key: &[u8; 32]) -> String {
620        pickle(&self, pickle_key)
621    }
622
623    /// Obtain a pickle from a ciphertext by decrypting and deserializing using
624    /// the given key.
625    ///
626    /// This is the inverse of [`AccountPickle::encrypt`].
627    pub fn from_encrypted(ciphertext: &str, pickle_key: &[u8; 32]) -> Result<Self, PickleError> {
628        unpickle(ciphertext, pickle_key)
629    }
630}
631
632impl From<AccountPickle> for Account {
633    fn from(pickle: AccountPickle) -> Self {
634        Self {
635            signing_key: pickle.signing_key.into(),
636            diffie_hellman_key: pickle.diffie_hellman_key.into(),
637            one_time_keys: pickle.one_time_keys.into(),
638            fallback_keys: pickle.fallback_keys,
639        }
640    }
641}
642
643#[cfg(feature = "libolm-compat")]
644mod libolm {
645    use matrix_pickle::{Decode, DecodeError, Encode, EncodeError};
646    use zeroize::{Zeroize, ZeroizeOnDrop};
647
648    use super::{
649        Account,
650        fallback_keys::{FallbackKey, FallbackKeys},
651        one_time_keys::OneTimeKeys,
652    };
653    use crate::{
654        Curve25519PublicKey, Ed25519Keypair, KeyId,
655        types::{Curve25519Keypair, Curve25519SecretKey},
656        utilities::LibolmEd25519Keypair,
657    };
658
659    #[derive(Encode, Decode, Zeroize, ZeroizeOnDrop)]
660    struct OneTimeKey {
661        key_id: u32,
662        published: bool,
663        public_key: [u8; 32],
664        private_key: Box<[u8; 32]>,
665    }
666
667    impl From<&OneTimeKey> for FallbackKey {
668        fn from(key: &OneTimeKey) -> Self {
669            FallbackKey {
670                key_id: KeyId(key.key_id.into()),
671                key: Curve25519SecretKey::from_slice(&key.private_key),
672                published: key.published,
673            }
674        }
675    }
676
677    #[derive(Zeroize, ZeroizeOnDrop)]
678    struct FallbackKeysArray {
679        fallback_key: Option<OneTimeKey>,
680        previous_fallback_key: Option<OneTimeKey>,
681    }
682
683    impl Decode for FallbackKeysArray {
684        fn decode(reader: &mut impl std::io::Read) -> Result<Self, DecodeError> {
685            let count = u8::decode(reader)?;
686
687            let (fallback_key, previous_fallback_key) = if count >= 1 {
688                let fallback_key = OneTimeKey::decode(reader)?;
689
690                let previous_fallback_key =
691                    if count >= 2 { Some(OneTimeKey::decode(reader)?) } else { None };
692
693                (Some(fallback_key), previous_fallback_key)
694            } else {
695                (None, None)
696            };
697
698            Ok(Self { fallback_key, previous_fallback_key })
699        }
700    }
701
702    impl Encode for FallbackKeysArray {
703        fn encode(&self, writer: &mut impl std::io::Write) -> Result<usize, EncodeError> {
704            let ret = match (&self.fallback_key, &self.previous_fallback_key) {
705                (None, None) => 0u8.encode(writer)?,
706                (Some(key), None) | (None, Some(key)) => {
707                    let mut ret = 1u8.encode(writer)?;
708                    ret += key.encode(writer)?;
709
710                    ret
711                }
712                (Some(key), Some(previous_key)) => {
713                    let mut ret = 2u8.encode(writer)?;
714                    ret += key.encode(writer)?;
715                    ret += previous_key.encode(writer)?;
716
717                    ret
718                }
719            };
720
721            Ok(ret)
722        }
723    }
724
725    #[derive(Encode, Decode, Zeroize, ZeroizeOnDrop)]
726    pub(super) struct Pickle {
727        version: u32,
728        ed25519_keypair: LibolmEd25519Keypair,
729        public_curve25519_key: [u8; 32],
730        private_curve25519_key: Box<[u8; 32]>,
731        one_time_keys: Vec<OneTimeKey>,
732        fallback_keys: FallbackKeysArray,
733        next_key_id: u32,
734    }
735
736    impl TryFrom<&FallbackKey> for OneTimeKey {
737        type Error = ();
738
739        fn try_from(key: &FallbackKey) -> Result<Self, ()> {
740            Ok(OneTimeKey {
741                key_id: key.key_id.0.try_into().map_err(|_| ())?,
742                published: key.published(),
743                public_key: key.public_key().to_bytes(),
744                private_key: key.secret_key().to_bytes(),
745            })
746        }
747    }
748
749    impl From<&Account> for Pickle {
750        fn from(account: &Account) -> Self {
751            let one_time_keys: Vec<_> = account
752                .one_time_keys
753                .secret_keys()
754                .iter()
755                .filter_map(|(key_id, secret_key)| {
756                    Some(OneTimeKey {
757                        key_id: key_id.0.try_into().ok()?,
758                        published: account.one_time_keys.is_secret_key_published(key_id),
759                        public_key: Curve25519PublicKey::from(secret_key).to_bytes(),
760                        private_key: secret_key.to_bytes(),
761                    })
762                })
763                .collect();
764
765            let fallback_keys = FallbackKeysArray {
766                fallback_key: account
767                    .fallback_keys
768                    .fallback_key
769                    .as_ref()
770                    .and_then(|f| f.try_into().ok()),
771                previous_fallback_key: account
772                    .fallback_keys
773                    .previous_fallback_key
774                    .as_ref()
775                    .and_then(|f| f.try_into().ok()),
776            };
777
778            let next_key_id = account.one_time_keys.next_key_id.try_into().unwrap_or_default();
779
780            Self {
781                version: 4,
782                ed25519_keypair: LibolmEd25519Keypair {
783                    private_key: account.signing_key.expanded_secret_key(),
784                    public_key: *account.signing_key.public_key().as_bytes(),
785                },
786                public_curve25519_key: account.diffie_hellman_key.public_key().to_bytes(),
787                private_curve25519_key: account.diffie_hellman_key.secret_key().to_bytes(),
788                one_time_keys,
789                fallback_keys,
790                next_key_id,
791            }
792        }
793    }
794
795    impl TryFrom<Pickle> for Account {
796        type Error = crate::LibolmPickleError;
797
798        fn try_from(pickle: Pickle) -> Result<Self, Self::Error> {
799            let mut one_time_keys = OneTimeKeys::new();
800
801            for key in &pickle.one_time_keys {
802                let secret_key = Curve25519SecretKey::from_slice(&key.private_key);
803                let key_id = KeyId(key.key_id.into());
804                one_time_keys.insert_secret_key(key_id, secret_key, key.published);
805            }
806
807            one_time_keys.next_key_id = u64::from(pickle.next_key_id) + 1;
808
809            let fallback_keys = FallbackKeys {
810                key_id: pickle
811                    .fallback_keys
812                    .fallback_key
813                    .as_ref()
814                    .map(|k| k.key_id.wrapping_add(1))
815                    .unwrap_or(0) as u64,
816                fallback_key: pickle.fallback_keys.fallback_key.as_ref().map(|k| k.into()),
817                previous_fallback_key: pickle
818                    .fallback_keys
819                    .previous_fallback_key
820                    .as_ref()
821                    .map(|k| k.into()),
822            };
823
824            Ok(Self {
825                signing_key: Ed25519Keypair::from_expanded_key(
826                    &pickle.ed25519_keypair.private_key,
827                )?,
828                diffie_hellman_key: Curve25519Keypair::from_secret_key(
829                    &pickle.private_curve25519_key,
830                ),
831                one_time_keys,
832                fallback_keys,
833            })
834        }
835    }
836}
837
838mod dehydrated_device {
839    use matrix_pickle::{Decode, DecodeError, Encode, EncodeError};
840    use zeroize::{Zeroize, ZeroizeOnDrop};
841
842    use super::{
843        Account,
844        fallback_keys::{FallbackKey, FallbackKeys},
845        one_time_keys::OneTimeKeys,
846    };
847    use crate::{
848        DehydratedDeviceError, Ed25519Keypair, KeyId,
849        types::{Curve25519Keypair, Curve25519SecretKey},
850    };
851
852    #[derive(Encode, Decode, Zeroize, ZeroizeOnDrop)]
853    pub(crate) struct OneTimeKey {
854        #[secret]
855        pub(crate) private_key: Box<[u8; 32]>,
856    }
857
858    impl From<&OneTimeKey> for FallbackKey {
859        fn from(key: &OneTimeKey) -> Self {
860            FallbackKey {
861                key_id: KeyId(0),
862                key: Curve25519SecretKey::from_slice(&key.private_key),
863                published: true,
864            }
865        }
866    }
867
868    impl TryFrom<&FallbackKey> for OneTimeKey {
869        type Error = ();
870
871        fn try_from(key: &FallbackKey) -> Result<Self, ()> {
872            Ok(OneTimeKey { private_key: key.secret_key().to_bytes() })
873        }
874    }
875
876    #[derive(Zeroize, ZeroizeOnDrop)]
877    pub(crate) struct OptFallbackKey {
878        pub(crate) fallback_key: Option<OneTimeKey>,
879    }
880
881    impl Decode for OptFallbackKey {
882        fn decode(reader: &mut impl std::io::Read) -> Result<Self, DecodeError> {
883            let present = bool::decode(reader)?;
884
885            let fallback_key = if present {
886                let fallback_key = OneTimeKey::decode(reader)?;
887
888                Some(fallback_key)
889            } else {
890                None
891            };
892
893            Ok(Self { fallback_key })
894        }
895    }
896
897    impl Encode for OptFallbackKey {
898        fn encode(&self, writer: &mut impl std::io::Write) -> Result<usize, EncodeError> {
899            let ret = match &self.fallback_key {
900                None => false.encode(writer)?,
901                Some(key) => {
902                    let mut ret = true.encode(writer)?;
903                    ret += key.encode(writer)?;
904
905                    ret
906                }
907            };
908
909            Ok(ret)
910        }
911    }
912
913    #[derive(Encode, Decode, Zeroize, ZeroizeOnDrop)]
914    /// Pickle used for dehydrated devices.
915    ///
916    /// Dehydrated devices are used for receiving encrypted messages when the
917    /// user has no other devices logged in, and are defined in
918    /// [MSC3814](https://github.com/matrix-org/matrix-spec-proposals/pull/3814).
919    pub(super) struct Pickle {
920        version: u32,
921        #[secret]
922        private_curve25519_key: Box<[u8; 32]>,
923        #[secret]
924        private_ed25519_key: Box<[u8; 32]>,
925        one_time_keys: Vec<OneTimeKey>,
926        opt_fallback_key: OptFallbackKey,
927    }
928
929    pub(super) const PICKLE_VERSION: u32 = 1;
930
931    impl TryFrom<&Account> for Pickle {
932        type Error = DehydratedDeviceError;
933
934        fn try_from(account: &Account) -> Result<Self, Self::Error> {
935            let one_time_keys: Vec<_> = account
936                .one_time_keys
937                .secret_keys()
938                .values()
939                .map(|secret_key| OneTimeKey { private_key: secret_key.to_bytes() })
940                .collect();
941
942            let fallback_key =
943                account.fallback_keys.fallback_key.as_ref().and_then(|f| f.try_into().ok());
944
945            Ok(Self {
946                version: PICKLE_VERSION,
947                private_curve25519_key: account.diffie_hellman_key.secret_key().to_bytes(),
948                private_ed25519_key: account
949                    .signing_key
950                    .unexpanded_secret_key()
951                    .ok_or(DehydratedDeviceError::InvalidAccount)?,
952                one_time_keys,
953                opt_fallback_key: OptFallbackKey { fallback_key },
954            })
955        }
956    }
957
958    impl TryFrom<Pickle> for Account {
959        type Error = DehydratedDeviceError;
960
961        fn try_from(pickle: Pickle) -> Result<Self, Self::Error> {
962            use crate::{DehydratedDeviceError, LibolmPickleError};
963            let mut one_time_keys = OneTimeKeys::new();
964
965            for (num, key) in pickle.one_time_keys.iter().enumerate() {
966                let secret_key = Curve25519SecretKey::from_slice(&key.private_key);
967                let key_id = KeyId(num as u64);
968                one_time_keys.insert_secret_key(key_id, secret_key, true);
969            }
970            one_time_keys.next_key_id = pickle.one_time_keys.len().try_into().unwrap_or_default();
971
972            let fallback_keys = FallbackKeys {
973                key_id: 1,
974                fallback_key: pickle.opt_fallback_key.fallback_key.as_ref().map(|otk| otk.into()),
975                previous_fallback_key: None,
976            };
977
978            Ok(Self {
979                signing_key: Ed25519Keypair::from_unexpanded_key(&pickle.private_ed25519_key)
980                    .map_err(|e| {
981                        DehydratedDeviceError::LibolmPickle(LibolmPickleError::PublicKey(e))
982                    })?,
983                diffie_hellman_key: Curve25519Keypair::from_secret_key(
984                    &pickle.private_curve25519_key,
985                ),
986                one_time_keys,
987                fallback_keys,
988            })
989        }
990    }
991}
992
993#[cfg(test)]
994mod test {
995    use anyhow::{Context, Result, bail};
996    use assert_matches2::assert_matches;
997    use matrix_pickle::{Decode, Encode};
998    use olm_rs::{account::OlmAccount, session::OlmMessage as LibolmOlmMessage};
999
1000    #[cfg(feature = "libolm-compat")]
1001    use super::libolm::Pickle;
1002    use super::{
1003        Account, InboundCreationResult, SessionConfig, SessionCreationError, dehydrated_device,
1004    };
1005    use crate::{
1006        Curve25519PublicKey as PublicKey,
1007        cipher::Mac,
1008        olm::{
1009            AccountPickle,
1010            account::PUBLIC_MAX_ONE_TIME_KEYS,
1011            messages::{OlmMessage, PreKeyMessage},
1012        },
1013    };
1014
1015    const PICKLE_KEY: [u8; 32] = [0u8; 32];
1016
1017    #[test]
1018    fn max_number_of_one_time_keys_matches_global_constant() {
1019        assert_eq!(Account::new().max_number_of_one_time_keys(), PUBLIC_MAX_ONE_TIME_KEYS);
1020    }
1021
1022    #[test]
1023    #[cfg(feature = "low-level-api")]
1024    fn generate_and_remove_one_time_key() {
1025        let mut alice = Account::new();
1026        assert_eq!(alice.stored_one_time_key_count(), 0);
1027
1028        alice.generate_one_time_keys(1);
1029        assert_eq!(alice.stored_one_time_key_count(), 1);
1030
1031        let public_key = alice
1032            .one_time_keys()
1033            .values()
1034            .next()
1035            .copied()
1036            .expect("Should have an unpublished one-time key");
1037        let secret_key_bytes = alice
1038            .find_one_time_key(&public_key)
1039            .expect("Should find secret key for public one-time key")
1040            .to_bytes();
1041        let removed_key_bytes = alice
1042            .remove_one_time_key(public_key)
1043            .expect("Should be able to remove one-time key")
1044            .to_bytes();
1045
1046        assert_eq!(removed_key_bytes, secret_key_bytes);
1047        assert_eq!(alice.stored_one_time_key_count(), 0);
1048    }
1049
1050    #[test]
1051    fn generate_and_forget_fallback_keys() {
1052        let mut alice = Account::default();
1053        assert!(!alice.forget_fallback_key());
1054        alice.generate_fallback_key();
1055        assert!(!alice.forget_fallback_key());
1056        alice.generate_fallback_key();
1057        assert!(alice.forget_fallback_key());
1058    }
1059
1060    #[test]
1061    fn vodozemac_libolm_communication() -> Result<()> {
1062        // vodozemac account
1063        let alice = Account::new();
1064        // libolm account
1065        let bob = OlmAccount::new();
1066
1067        bob.generate_one_time_keys(1);
1068
1069        let one_time_key = bob
1070            .parsed_one_time_keys()
1071            .curve25519()
1072            .values()
1073            .next()
1074            .cloned()
1075            .expect("Didn't find a valid one-time key");
1076
1077        bob.mark_keys_as_published();
1078
1079        let identity_keys = bob.parsed_identity_keys();
1080        let curve25519_key = PublicKey::from_base64(identity_keys.curve25519())?;
1081        let one_time_key = PublicKey::from_base64(&one_time_key)?;
1082        let mut alice_session = alice.create_outbound_session(
1083            SessionConfig::version_1(),
1084            curve25519_key,
1085            one_time_key,
1086        )?;
1087
1088        let message = "It's a secret to everybody";
1089        let olm_message: LibolmOlmMessage = alice_session.encrypt(message).unwrap().into();
1090
1091        if let LibolmOlmMessage::PreKey(m) = olm_message.clone() {
1092            let libolm_session =
1093                bob.create_inbound_session_from(&alice.curve25519_key().to_base64(), m)?;
1094            assert_eq!(alice_session.session_id(), libolm_session.session_id());
1095
1096            let plaintext = libolm_session.decrypt(olm_message)?;
1097            assert_eq!(message, plaintext);
1098
1099            let second_text = "Here's another secret to everybody";
1100            let olm_message = alice_session.encrypt(second_text).unwrap().into();
1101
1102            let plaintext = libolm_session.decrypt(olm_message)?;
1103            assert_eq!(second_text, plaintext);
1104
1105            let reply_plain = "Yes, take this, it's dangerous out there";
1106            let reply = libolm_session.encrypt(reply_plain).into();
1107            let plaintext = alice_session.decrypt(&reply)?;
1108
1109            assert_eq!(plaintext, reply_plain.as_bytes());
1110
1111            let another_reply = "Last one";
1112            let reply = libolm_session.encrypt(another_reply).into();
1113            let plaintext = alice_session.decrypt(&reply)?;
1114            assert_eq!(plaintext, another_reply.as_bytes());
1115
1116            let last_text = "Nope, I'll have the last word";
1117            let olm_message = alice_session.encrypt(last_text).unwrap().into();
1118
1119            let plaintext = libolm_session.decrypt(olm_message)?;
1120            assert_eq!(last_text, plaintext);
1121        } else {
1122            bail!("Received a invalid message type {:?}", olm_message);
1123        }
1124
1125        Ok(())
1126    }
1127
1128    fn vodozemac_vodozemac_communication(session_config: SessionConfig) -> Result<()> {
1129        // Both of these are vodozemac accounts.
1130        let alice = Account::new();
1131        let mut bob = Account::new();
1132
1133        bob.generate_one_time_keys(1);
1134
1135        let mut alice_session = alice.create_outbound_session(
1136            session_config,
1137            bob.curve25519_key(),
1138            *bob.one_time_keys()
1139                .iter()
1140                .next()
1141                .context("Failed getting bob's OTK, which should never happen here.")?
1142                .1,
1143        )?;
1144
1145        assert!(!bob.one_time_keys().is_empty());
1146        bob.mark_keys_as_published();
1147        assert!(bob.one_time_keys().is_empty());
1148
1149        let message = "It's a secret to everybody";
1150        let olm_message = alice_session.encrypt(message).unwrap();
1151
1152        if let OlmMessage::PreKey(m) = olm_message {
1153            assert_eq!(m.session_keys(), alice_session.session_keys());
1154            assert_eq!(m.session_id(), alice_session.session_id());
1155
1156            let InboundCreationResult { session: mut bob_session, plaintext } =
1157                bob.create_inbound_session(session_config, alice.curve25519_key(), &m)?;
1158            assert_eq!(alice_session.session_id(), bob_session.session_id());
1159            assert_eq!(m.session_keys(), bob_session.session_keys());
1160
1161            assert_eq!(message.as_bytes(), plaintext);
1162
1163            let second_text = "Here's another secret to everybody";
1164            let olm_message = alice_session.encrypt(second_text).unwrap();
1165
1166            let plaintext = bob_session.decrypt(&olm_message)?;
1167            assert_eq!(second_text.as_bytes(), plaintext);
1168
1169            let reply_plain = "Yes, take this, it's dangerous out there";
1170            let reply = bob_session.encrypt(reply_plain).unwrap();
1171            let plaintext = alice_session.decrypt(&reply)?;
1172
1173            assert_eq!(plaintext, reply_plain.as_bytes());
1174
1175            let another_reply = "Last one";
1176            let reply = bob_session.encrypt(another_reply).unwrap();
1177            let plaintext = alice_session.decrypt(&reply)?;
1178            assert_eq!(plaintext, another_reply.as_bytes());
1179
1180            let last_text = "Nope, I'll have the last word";
1181            let olm_message = alice_session.encrypt(last_text).unwrap();
1182
1183            let plaintext = bob_session.decrypt(&olm_message)?;
1184            assert_eq!(last_text.as_bytes(), plaintext);
1185        }
1186
1187        Ok(())
1188    }
1189
1190    #[test]
1191    fn vodozemac_vodozemac_communication_session_config_v1() -> Result<()> {
1192        vodozemac_vodozemac_communication(SessionConfig::version_1())
1193    }
1194
1195    #[test]
1196    #[cfg(feature = "experimental-session-config")]
1197    fn vodozemac_vodozemac_communication_session_config_v2() -> Result<()> {
1198        vodozemac_vodozemac_communication(SessionConfig::version_2())
1199    }
1200
1201    #[test]
1202    fn inbound_session_creation() -> Result<()> {
1203        let alice = OlmAccount::new();
1204        let mut bob = Account::new();
1205
1206        bob.generate_one_time_keys(1);
1207
1208        let one_time_key =
1209            bob.one_time_keys().values().next().cloned().expect("Didn't find a valid one-time key");
1210
1211        let alice_session = alice.create_outbound_session(
1212            &bob.curve25519_key().to_base64(),
1213            &one_time_key.to_base64(),
1214        )?;
1215
1216        let text = "It's a secret to everybody";
1217        let message = alice_session.encrypt(text).into();
1218
1219        let identity_key = PublicKey::from_base64(alice.parsed_identity_keys().curve25519())?;
1220
1221        let InboundCreationResult { session, plaintext } = if let OlmMessage::PreKey(m) = &message {
1222            bob.create_inbound_session(SessionConfig::version_1(), identity_key, m)?
1223        } else {
1224            bail!("Got invalid message type from olm_rs {:?}", message);
1225        };
1226
1227        assert_eq!(alice_session.session_id(), session.session_id());
1228        assert!(bob.one_time_keys.private_keys.is_empty());
1229
1230        assert_eq!(text.as_bytes(), plaintext);
1231
1232        Ok(())
1233    }
1234
1235    #[test]
1236    fn inbound_session_creation_using_fallback_keys() -> Result<()> {
1237        let alice = OlmAccount::new();
1238        let mut bob = Account::new();
1239
1240        bob.generate_fallback_key();
1241
1242        let one_time_key =
1243            bob.fallback_key().values().next().cloned().expect("Didn't find a valid fallback key");
1244        assert!(bob.one_time_keys.private_keys.is_empty());
1245
1246        let alice_session = alice.create_outbound_session(
1247            &bob.curve25519_key().to_base64(),
1248            &one_time_key.to_base64(),
1249        )?;
1250
1251        let text = "It's a secret to everybody";
1252
1253        let message = alice_session.encrypt(text).into();
1254        let identity_key = PublicKey::from_base64(alice.parsed_identity_keys().curve25519())?;
1255
1256        if let OlmMessage::PreKey(m) = &message {
1257            let InboundCreationResult { session, plaintext } =
1258                bob.create_inbound_session(SessionConfig::version_1(), identity_key, m)?;
1259
1260            assert_eq!(m.session_keys(), session.session_keys());
1261            assert_eq!(alice_session.session_id(), session.session_id());
1262            assert!(bob.fallback_keys.fallback_key.is_some());
1263
1264            assert_eq!(text.as_bytes(), plaintext);
1265        } else {
1266            bail!("Got invalid message type from olm_rs");
1267        };
1268
1269        Ok(())
1270    }
1271
1272    #[test]
1273    fn account_pickling_roundtrip_is_identity() -> Result<()> {
1274        let mut account = Account::new();
1275
1276        account.generate_one_time_keys(50);
1277
1278        // Generate two fallback keys so the previous fallback key field gets populated.
1279        account.generate_fallback_key();
1280        account.generate_fallback_key();
1281
1282        let pickle = account.pickle().encrypt(&PICKLE_KEY);
1283
1284        let decrypted_pickle = AccountPickle::from_encrypted(&pickle, &PICKLE_KEY)?;
1285        let unpickled_account = Account::from_pickle(decrypted_pickle);
1286        let repickle = unpickled_account.pickle();
1287
1288        assert_eq!(account.identity_keys(), unpickled_account.identity_keys());
1289
1290        let decrypted_pickle = AccountPickle::from_encrypted(&pickle, &PICKLE_KEY)?;
1291        let pickle = serde_json::to_value(decrypted_pickle)?;
1292        let repickle = serde_json::to_value(repickle)?;
1293
1294        assert_eq!(pickle, repickle);
1295
1296        Ok(())
1297    }
1298
1299    #[test]
1300    #[cfg(feature = "libolm-compat")]
1301    fn libolm_unpickling() -> Result<()> {
1302        let olm = OlmAccount::new();
1303        olm.generate_one_time_keys(10);
1304        olm.generate_fallback_key();
1305
1306        let key = b"DEFAULT_PICKLE_KEY";
1307        let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
1308
1309        let unpickled = Account::from_libolm_pickle(&pickle, key)?;
1310
1311        assert_eq!(olm.parsed_identity_keys().ed25519(), unpickled.ed25519_key().to_base64());
1312        assert_eq!(olm.parsed_identity_keys().curve25519(), unpickled.curve25519_key().to_base64());
1313
1314        let mut olm_one_time_keys: Vec<_> =
1315            olm.parsed_one_time_keys().curve25519().values().map(|k| k.to_owned()).collect();
1316        let mut one_time_keys: Vec<_> =
1317            unpickled.one_time_keys().values().map(|k| k.to_base64()).collect();
1318
1319        // We generated 10 one-time keys on the libolm side + 1 fallback key,
1320        // we expect the next key id to be 12.
1321        assert_eq!(unpickled.one_time_keys.next_key_id, 12);
1322
1323        olm_one_time_keys.sort();
1324        one_time_keys.sort();
1325        assert_eq!(olm_one_time_keys, one_time_keys);
1326
1327        let olm_fallback_key =
1328            olm.parsed_fallback_key().expect("libolm should have a fallback key");
1329        assert_eq!(
1330            olm_fallback_key.curve25519(),
1331            unpickled
1332                .fallback_key()
1333                .values()
1334                .next()
1335                .expect("We should have a fallback key")
1336                .to_base64()
1337        );
1338
1339        Ok(())
1340    }
1341
1342    #[test]
1343    #[cfg(feature = "libolm-compat")]
1344    fn libolm_unpickling_without_fallback_key_does_not_invalidate_last_one_time_key() -> Result<()>
1345    {
1346        let olm = OlmAccount::new();
1347        olm.generate_one_time_keys(5);
1348        olm.mark_keys_as_published();
1349
1350        let key = b"DEFAULT_PICKLE_KEY";
1351        let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
1352
1353        let mut unpickled = Account::from_libolm_pickle(&pickle, key)?;
1354        assert_eq!(unpickled.one_time_keys.next_key_id, 6);
1355
1356        let imported: Vec<_> = unpickled.one_time_keys.secret_keys().keys().copied().collect();
1357        unpickled.generate_one_time_keys(1);
1358
1359        // All 5 imported keys must survive, plus the new one under a fresh id.
1360        assert_eq!(unpickled.one_time_keys.secret_keys().len(), 6);
1361        for key_id in imported {
1362            assert!(unpickled.one_time_keys.secret_keys().contains_key(&key_id));
1363        }
1364
1365        Ok(())
1366    }
1367
1368    #[test]
1369    #[cfg(feature = "libolm-compat")]
1370    fn pickle_cycle_with_one_fallback_key() {
1371        let mut alice = Account::new();
1372        alice.generate_fallback_key();
1373
1374        let mut encoded = Vec::<u8>::new();
1375        let pickle = Pickle::from(&alice);
1376        let size = pickle.encode(&mut encoded).expect("Should encode pickle");
1377        assert_eq!(size, encoded.len());
1378
1379        let account =
1380            Account::from_decrypted_libolm_pickle(&encoded).expect("Should unpickle account");
1381
1382        let key_bytes = alice
1383            .fallback_key()
1384            .values()
1385            .next()
1386            .expect("Should have a fallback key before encoding")
1387            .to_bytes();
1388        let decoded_key_bytes = account
1389            .fallback_key()
1390            .values()
1391            .next()
1392            .expect("Should have a fallback key after decoding")
1393            .to_bytes();
1394        assert_eq!(key_bytes, decoded_key_bytes);
1395    }
1396
1397    #[test]
1398    #[cfg(feature = "libolm-compat")]
1399    fn pickle_cycle_with_two_fallback_keys() {
1400        let mut alice = Account::new();
1401        alice.generate_fallback_key();
1402        alice.generate_fallback_key();
1403
1404        let mut encoded = Vec::<u8>::new();
1405        let pickle = Pickle::from(&alice);
1406        let size = pickle.encode(&mut encoded).expect("Should encode pickle");
1407        assert_eq!(size, encoded.len());
1408
1409        let account =
1410            Account::from_decrypted_libolm_pickle(&encoded).expect("Should unpickle account");
1411
1412        let key_bytes = alice
1413            .fallback_key()
1414            .values()
1415            .next()
1416            .expect("Should have a fallback key before encoding")
1417            .to_bytes();
1418        let decoded_key_bytes = account
1419            .fallback_key()
1420            .values()
1421            .next()
1422            .expect("Should have a fallback key after decoding")
1423            .to_bytes();
1424        assert_eq!(key_bytes, decoded_key_bytes);
1425    }
1426
1427    #[test]
1428    #[cfg(feature = "libolm-compat")]
1429    fn signing_with_expanded_key() -> Result<()> {
1430        let olm = OlmAccount::new();
1431        olm.generate_one_time_keys(10);
1432        olm.generate_fallback_key();
1433
1434        let key = b"DEFAULT_PICKLE_KEY";
1435        let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
1436
1437        let account_with_expanded_key = Account::from_libolm_pickle(&pickle, key)?;
1438
1439        // The clone is needed since we're later on using the account.
1440        #[allow(clippy::redundant_clone)]
1441        let signing_key_clone = account_with_expanded_key.signing_key.clone();
1442        signing_key_clone.sign("You met with a terrible fate, haven’t you?".as_bytes());
1443        account_with_expanded_key.sign("You met with a terrible fate, haven’t you?".as_bytes());
1444
1445        Ok(())
1446    }
1447
1448    #[test]
1449    fn invalid_session_creation_does_not_remove_otk() -> Result<()> {
1450        let mut alice = Account::new();
1451        let malory = Account::new();
1452        alice.generate_one_time_keys(1);
1453
1454        let mut session = malory.create_outbound_session(
1455            SessionConfig::default(),
1456            alice.curve25519_key(),
1457            *alice.one_time_keys().values().next().expect("Should have one-time key"),
1458        )?;
1459
1460        let message = session.encrypt("Test").unwrap();
1461
1462        if let OlmMessage::PreKey(m) = message {
1463            let mut message = m.to_bytes();
1464            let message_len = message.len();
1465
1466            // We mangle the MAC so decryption fails but creating a Session
1467            // succeeds.
1468            message[message_len - Mac::TRUNCATED_LEN..message_len]
1469                .copy_from_slice(&[0u8; Mac::TRUNCATED_LEN]);
1470
1471            let message = PreKeyMessage::try_from(message)?;
1472
1473            match alice.create_inbound_session(
1474                SessionConfig::version_1(),
1475                malory.curve25519_key(),
1476                &message,
1477            ) {
1478                Err(SessionCreationError::Decryption(_)) => {}
1479                e => bail!("Expected a decryption error, got {:?}", e),
1480            }
1481            assert!(
1482                !alice.one_time_keys.private_keys.is_empty(),
1483                "The one-time key was removed when it shouldn't"
1484            );
1485
1486            Ok(())
1487        } else {
1488            bail!("Invalid message type");
1489        }
1490    }
1491
1492    #[test]
1493    #[cfg(feature = "libolm-compat")]
1494    fn fuzz_corpus_unpickling() {
1495        crate::run_corpus("olm-account-unpickling", |data| {
1496            let _ = Account::from_decrypted_libolm_pickle(data);
1497        });
1498    }
1499
1500    #[test]
1501    #[cfg(feature = "libolm-compat")]
1502    fn libolm_pickle_cycle() -> Result<()> {
1503        let message = "It's a secret to everybody";
1504
1505        let olm = OlmAccount::new();
1506        olm.generate_one_time_keys(10);
1507        olm.generate_fallback_key();
1508
1509        let olm_signature = olm.sign(message);
1510
1511        let key = b"DEFAULT_PICKLE_KEY";
1512        let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
1513
1514        let account = Account::from_libolm_pickle(&pickle, key).unwrap();
1515        let vodozemac_pickle = account.to_libolm_pickle(key).unwrap();
1516        let _ = Account::from_libolm_pickle(&vodozemac_pickle, key).unwrap();
1517
1518        let vodozemac_signature = account.sign(message.as_bytes());
1519        let olm_signature = crate::types::Ed25519Signature::from_base64(&olm_signature)
1520            .expect("We should be able to parse a signature produced by libolm");
1521        account
1522            .identity_keys()
1523            .ed25519
1524            .verify(message.as_bytes(), &olm_signature)
1525            .expect("We should be able to verify the libolm signature with our vodozemac Account");
1526
1527        let unpickled = OlmAccount::unpickle(
1528            vodozemac_pickle,
1529            olm_rs::PicklingMode::Encrypted { key: key.to_vec() },
1530        )
1531        .unwrap();
1532
1533        let utility = olm_rs::utility::OlmUtility::new();
1534        utility
1535            .ed25519_verify(
1536                unpickled.parsed_identity_keys().ed25519(),
1537                message,
1538                vodozemac_signature.to_base64(),
1539            )
1540            .expect("We should be able to verify the signature vodozemac created");
1541        utility
1542            .ed25519_verify(
1543                unpickled.parsed_identity_keys().ed25519(),
1544                message,
1545                olm_signature.to_base64(),
1546            )
1547            .expect("We should be able to verify the original signature from libolm");
1548
1549        assert_eq!(olm.parsed_identity_keys(), unpickled.parsed_identity_keys());
1550
1551        Ok(())
1552    }
1553
1554    #[test]
1555    fn decrypt_with_dehydrated_device() {
1556        let mut alice = Account::new();
1557        let bob = Account::new();
1558        let carol = Account::new();
1559
1560        alice.generate_one_time_keys(alice.max_number_of_one_time_keys());
1561        alice.generate_fallback_key();
1562
1563        let alice_dehydrated_result =
1564            alice.to_dehydrated_device(&PICKLE_KEY).expect("Should be able to dehydrate device");
1565
1566        // encrypt using a one-time key
1567        let mut bob_session = bob
1568            .create_outbound_session(
1569                SessionConfig::version_1(),
1570                alice.curve25519_key(),
1571                *alice
1572                    .one_time_keys()
1573                    .iter()
1574                    .next()
1575                    .expect("Failed getting alice's OTK, which should never happen here.")
1576                    .1,
1577            )
1578            .unwrap();
1579
1580        // encrypt using a fallback key
1581        let mut carol_session = carol
1582            .create_outbound_session(
1583                SessionConfig::version_1(),
1584                alice.curve25519_key(),
1585                *alice
1586                    .fallback_key()
1587                    .iter()
1588                    .next()
1589                    .expect("Failed getting alice's fallback key, which should never happen here.")
1590                    .1,
1591            )
1592            .unwrap();
1593
1594        let message = "It's a secret to everybody";
1595        let bob_olm_message = bob_session.encrypt(message).unwrap();
1596        let carol_olm_message = carol_session.encrypt(message).unwrap();
1597
1598        let mut alice_rehydrated = Account::from_dehydrated_device(
1599            &alice_dehydrated_result.ciphertext,
1600            &alice_dehydrated_result.nonce,
1601            &PICKLE_KEY,
1602        )
1603        .expect("Should be able to rehydrate device");
1604
1605        // make sure we can decrypt both messages
1606        assert_matches!(bob_olm_message, OlmMessage::PreKey(prekey_message));
1607        let InboundCreationResult { session: alice_session, plaintext } = alice_rehydrated
1608            .create_inbound_session(
1609                SessionConfig::version_1(),
1610                bob.curve25519_key(),
1611                &prekey_message,
1612            )
1613            .expect("Alice should be able to create an inbound session from Bob's pre-key message");
1614        assert_eq!(alice_session.session_id(), bob_session.session_id());
1615        assert_eq!(message.as_bytes(), plaintext);
1616
1617        assert_matches!(carol_olm_message, OlmMessage::PreKey(prekey_message));
1618        let InboundCreationResult { session: alice_session, plaintext } = alice_rehydrated
1619            .create_inbound_session(
1620                SessionConfig::version_1(),
1621                carol.curve25519_key(),
1622                &prekey_message,
1623            )
1624            .expect(
1625                "Alice should be able to create an inbound session from Carol's pre-key message",
1626            );
1627
1628        assert_eq!(alice_session.session_id(), carol_session.session_id());
1629        assert_eq!(message.as_bytes(), plaintext);
1630    }
1631
1632    #[test]
1633    fn fails_to_rehydrate_with_wrong_key() {
1634        let mut alice = Account::new();
1635
1636        alice.generate_one_time_keys(alice.max_number_of_one_time_keys());
1637        alice.generate_fallback_key();
1638
1639        let alice_dehydrated_result =
1640            alice.to_dehydrated_device(&PICKLE_KEY).expect("Should be able to dehydrate device");
1641
1642        assert!(
1643            Account::from_dehydrated_device(
1644                &alice_dehydrated_result.ciphertext,
1645                &alice_dehydrated_result.nonce,
1646                &[1; 32],
1647            )
1648            .is_err()
1649        );
1650
1651        assert!(
1652            Account::from_dehydrated_device(
1653                &alice_dehydrated_result.ciphertext,
1654                "WrongNonce123456",
1655                &PICKLE_KEY,
1656            )
1657            .is_err()
1658        );
1659    }
1660
1661    #[derive(Encode, Decode)]
1662    struct OptFallbackPickleTest {
1663        fallback1: dehydrated_device::OptFallbackKey,
1664        fallback2: dehydrated_device::OptFallbackKey,
1665    }
1666
1667    #[test]
1668    fn encodes_optional_fallback_key() {
1669        use std::io::Cursor;
1670
1671        let data_to_pickle = OptFallbackPickleTest {
1672            fallback1: dehydrated_device::OptFallbackKey {
1673                fallback_key: Some(dehydrated_device::OneTimeKey {
1674                    private_key: Box::new([1; 32]),
1675                }),
1676            },
1677            fallback2: dehydrated_device::OptFallbackKey { fallback_key: None },
1678        };
1679
1680        let buffer = Vec::<u8>::new();
1681        let mut cursor = Cursor::new(buffer);
1682        let pickle_length = data_to_pickle.encode(&mut cursor).expect("Can pickle data");
1683        let pickle = cursor.into_inner();
1684        assert_eq!(pickle.len(), pickle_length);
1685
1686        let mut cursor = Cursor::new(&pickle);
1687        let unpickled_data = OptFallbackPickleTest::decode(&mut cursor).expect("Can unpickle");
1688
1689        assert!(unpickled_data.fallback1.fallback_key.is_some());
1690        assert!(unpickled_data.fallback2.fallback_key.is_none());
1691    }
1692
1693    #[test]
1694    fn decrypted_dehydration_cycle() {
1695        use dehydrated_device::Pickle;
1696
1697        let alice = Account::new();
1698
1699        let mut encoded = Vec::<u8>::new();
1700        let pickle = Pickle::try_from(&alice)
1701            .expect("We should be able to create a dehydrated device from the account");
1702        let size = pickle.encode(&mut encoded).expect("Should dehydrate");
1703        assert_eq!(size, encoded.len());
1704
1705        let account =
1706            Account::from_decrypted_dehydrated_device(&encoded).expect("Should rehydrate account");
1707
1708        assert_eq!(alice.identity_keys(), account.identity_keys());
1709    }
1710
1711    #[test]
1712    #[cfg(feature = "experimental-session-config")]
1713    fn create_session_with_incorrect_session_config_upgrade() {
1714        // Both of these are vodozemac accounts.
1715        let alice = Account::new();
1716        let mut bob = Account::new();
1717
1718        bob.generate_one_time_keys(1);
1719        let one_time_key =
1720            *bob.one_time_keys().values().next().expect("Bob should have generated a one-time key");
1721
1722        let mut alice_session = alice
1723            .create_outbound_session(SessionConfig::version_1(), bob.curve25519_key(), one_time_key)
1724            .expect("We should be able to create an outbound session with bob.");
1725
1726        let message = "It's a secret to everybody";
1727        let pre_key_message =
1728            alice_session.encrypt(message).expect("We should be able to encrypt the first message");
1729
1730        assert_matches2::assert_let!(OlmMessage::PreKey(pre_key_message) = pre_key_message);
1731
1732        let result = bob.create_inbound_session(
1733            SessionConfig::version_2(),
1734            alice.curve25519_key(),
1735            &pre_key_message,
1736        );
1737
1738        assert_matches!(
1739            result,
1740            Err(SessionCreationError::MismatchedSessionConfig { .. }),
1741            "We should not create a session if the incorrect session config was used"
1742        );
1743    }
1744
1745    #[test]
1746    #[cfg(feature = "experimental-session-config")]
1747    fn create_session_with_incorrect_session_config_downgrade() {
1748        let alice = Account::new();
1749        let mut bob = Account::new();
1750
1751        bob.generate_one_time_keys(1);
1752        let one_time_key =
1753            *bob.one_time_keys().values().next().expect("Bob should have generated a one-time key");
1754
1755        let mut alice_session = alice
1756            .create_outbound_session(SessionConfig::version_2(), bob.curve25519_key(), one_time_key)
1757            .expect("We should be able to create an outbound session with bob.");
1758
1759        let message = "It's a secret to everybody";
1760        let pre_key_message =
1761            alice_session.encrypt(message).expect("We should be able to encrypt the first message");
1762
1763        assert_matches2::assert_let!(OlmMessage::PreKey(pre_key_message) = pre_key_message);
1764
1765        let result = bob.create_inbound_session(
1766            SessionConfig::version_1(),
1767            alice.curve25519_key(),
1768            &pre_key_message,
1769        );
1770
1771        assert_matches!(
1772            result,
1773            Err(SessionCreationError::MismatchedSessionConfig { .. }),
1774            "We should not create a session if the incorrect session config was used"
1775        );
1776    }
1777
1778    #[test]
1779    fn create_session_with_unsupported_session_config() {
1780        let alice = Account::new();
1781        let mut bob = Account::new();
1782
1783        bob.generate_one_time_keys(1);
1784        let one_time_key =
1785            *bob.one_time_keys().values().next().expect("Bob should have generated a one-time key");
1786
1787        let mut alice_session = alice
1788            .create_outbound_session(SessionConfig::version_1(), bob.curve25519_key(), one_time_key)
1789            .expect("We should be able to create an outbound session with bob.");
1790
1791        let message = "It's a secret to everybody";
1792        let pre_key_message =
1793            alice_session.encrypt(message).expect("We should be able to encrypt the first message");
1794
1795        assert_matches2::assert_let!(OlmMessage::PreKey(mut pre_key_message) = pre_key_message);
1796
1797        // Technically this can't happen as the pre-key message parsing will reject such
1798        // a version, but let's double check if our session creation is robust against
1799        // unknown versions.
1800        pre_key_message.message.version = 0xFF;
1801
1802        let result = bob.create_inbound_session(
1803            SessionConfig::version_1(),
1804            alice.curve25519_key(),
1805            &pre_key_message,
1806        );
1807
1808        assert_matches!(
1809            result,
1810            Err(SessionCreationError::MismatchedSessionConfig { got: None, .. }),
1811            "We should not create a session if an unknown session config was used for the pre-key message"
1812        );
1813    }
1814}