From 4df6099125034b4174536a7f4e3518799daf57be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Mon, 27 Apr 2020 18:27:33 +0200
Subject: [PATCH] crypto: Test the Olm encryption.

---
 src/crypto/machine.rs | 67 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 66 insertions(+), 1 deletion(-)

diff --git a/src/crypto/machine.rs b/src/crypto/machine.rs
index ff03f22d..56e38727 100644
--- a/src/crypto/machine.rs
+++ b/src/crypto/machine.rs
@@ -1380,6 +1380,10 @@ mod test {
     use crate::api::r0::keys;
     use crate::crypto::machine::{OlmMachine, OneTimeKeys};
     use crate::crypto::Device;
+    use crate::events::{
+        to_device::{AnyToDeviceEvent, ToDeviceEncrypted},
+        EventType,
+    };
     use crate::identifiers::{DeviceId, UserId};
 
     use http::Response;
@@ -1459,6 +1463,29 @@ mod test {
         (alice, bob, otk)
     }
 
+    async fn get_machine_pair_with_session() -> (OlmMachine, OlmMachine) {
+        let (mut alice, bob, one_time_keys) = get_machine_pair().await;
+
+        let mut bob_keys = BTreeMap::new();
+
+        let one_time_key = one_time_keys.iter().nth(0).unwrap();
+        let mut keys = BTreeMap::new();
+        keys.insert(one_time_key.0.clone(), one_time_key.1.clone());
+        bob_keys.insert(bob.device_id.clone(), keys);
+
+        let mut one_time_keys = BTreeMap::new();
+        one_time_keys.insert(bob.user_id.clone(), bob_keys);
+
+        let response = keys::claim_keys::Response {
+            failures: BTreeMap::new(),
+            one_time_keys,
+        };
+
+        alice.receive_keys_claim_response(&response).await.unwrap();
+
+        (alice, bob)
+    }
+
     #[tokio::test]
     async fn create_olm_machine() {
         let machine = OlmMachine::new(&user_id(), DEVICE_ID).unwrap();
@@ -1669,7 +1696,7 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_key_claiming() {
+    async fn test_session_creation() {
         let (mut alice_machine, bob_machine, one_time_keys) = get_machine_pair().await;
 
         let mut bob_keys = BTreeMap::new();
@@ -1701,4 +1728,42 @@ mod test {
 
         assert!(!session.lock().await.is_empty())
     }
+
+    #[tokio::test]
+    async fn test_olm_encryption() {
+        let (mut alice, mut bob) = get_machine_pair_with_session().await;
+
+        let session = alice
+            .store
+            .get_sessions(bob.account.identity_keys().curve25519())
+            .await
+            .unwrap()
+            .unwrap()
+            .lock()
+            .await[0]
+            .clone();
+
+        let bob_device = alice
+            .store
+            .get_device(&bob.user_id, &bob.device_id)
+            .await
+            .unwrap()
+            .unwrap();
+
+        let event = ToDeviceEncrypted {
+            sender: alice.user_id.clone(),
+            content: alice
+                .olm_encrypt(session, &bob_device, EventType::Dummy, json!({}))
+                .await
+                .unwrap(),
+        };
+
+        let event = bob.decrypt_to_device_event(&event).await.unwrap();
+
+        if let AnyToDeviceEvent::Dummy(e) = event.deserialize().unwrap() {
+            assert_eq!(e.sender, alice.user_id);
+        } else {
+            panic!("Event had the wrong type");
+        }
+    }
 }