Mark messages received from a session as trusted

3 files changed, 25 insertions, 10 deletions

diff --git a/src/omemo/omemo.c b/src/omemo/omemo.c
index ba17fe51..96d2d65a 100644
--- a/src/omemo/omemo.c
+++ b/src/omemo/omemo.c
@@ -890,15 +890,26 @@ omemo_on_message_recv(const char *const from_jid, uint32_t sid,
         goto out;
     }
 
-    omemo_ctx.identity_key_store.recv = true;
-
     if (key->prekey) {
         log_debug("OMEMO: decrypting message with prekey");
         pre_key_signal_message *message;
+        ec_public_key *their_identity_key;
+        signal_buffer *identity_buffer = NULL;
+
+        omemo_ctx.identity_key_store.recv = true;
 
         pre_key_signal_message_deserialize(&message, key->data, key->length, omemo_ctx.signal);
+        their_identity_key = pre_key_signal_message_get_identity_key(message);
 
         res = session_cipher_decrypt_pre_key_signal_message(cipher, message, NULL, &plaintext_key);
+
+        omemo_ctx.identity_key_store.recv = false;
+
+        /* Perform a real check of the identity */
+        ec_public_key_serialize(&identity_buffer, their_identity_key);
+        *trusted = is_trusted_identity(&address, signal_buffer_data(identity_buffer),
+                signal_buffer_len(identity_buffer), &omemo_ctx.identity_key_store);
+
         /* Replace used pre_key in bundle */
         uint32_t pre_key_id = pre_key_signal_message_get_pre_key_id(message);
         ec_key_pair *ec_pair;
@@ -918,18 +929,18 @@ omemo_on_message_recv(const char *const from_jid, uint32_t sid,
     } else {
         log_debug("OMEMO: decrypting message with existing session");
         signal_message *message = NULL;
+
         res = signal_message_deserialize(&message, key->data, key->length, omemo_ctx.signal);
+
         if (res < 0) {
             log_error("OMEMO: cannot deserialize message");
         } else {
             res = session_cipher_decrypt_signal_message(cipher, message, NULL, &plaintext_key);
+            *trusted = true;
             SIGNAL_UNREF(message);
         }
     }
 
-    omemo_ctx.identity_key_store.recv = false;
-    *trusted = omemo_ctx.identity_key_store.trusted_msg;
-
     session_cipher_free(cipher);
     if (res != 0) {
         log_error("OMEMO: cannot decrypt message key");
diff --git a/src/omemo/store.c b/src/omemo/store.c
index 61f39c8b..3e44be43 100644
--- a/src/omemo/store.c
+++ b/src/omemo/store.c
@@ -362,9 +362,16 @@ save_identity(const signal_protocol_address *address, uint8_t *key_data,
 {
     identity_key_store_t *identity_key_store = (identity_key_store_t *)user_data;
 
-    if (identity_key_store->recv && !identity_key_store->trusted_msg) {
+    if (identity_key_store->recv) {
         /* Do not trust identity automatically */
-        return SG_SUCCESS;
+        /* Instead we perform a real trust check */
+        identity_key_store->recv = false;
+        int trusted = is_trusted_identity(address, key_data, key_len, user_data);
+        identity_key_store->recv = true;
+        if (trusted == 0) {
+            /* If not trusted we just don't save the identity */
+            return SG_SUCCESS;
+        }
     }
 
     signal_buffer *buffer = signal_buffer_create(key_data, key_len);
@@ -398,7 +405,6 @@ is_trusted_identity(const signal_protocol_address *address, uint8_t *key_data,
     GHashTable *trusted = g_hash_table_lookup(identity_key_store->trusted, address->name);
     if (!trusted) {
         if (identity_key_store->recv) {
-            identity_key_store->trusted_msg = false;
             return 1;
         } else {
             return 0;
@@ -414,7 +420,6 @@ is_trusted_identity(const signal_protocol_address *address, uint8_t *key_data,
 
 
     if (identity_key_store->recv) {
-        identity_key_store->trusted_msg = ret;
         return 1;
     } else {
         return ret;
diff --git a/src/omemo/store.h b/src/omemo/store.h
index a1524c62..f3a9274a 100644
--- a/src/omemo/store.h
+++ b/src/omemo/store.h
@@ -49,7 +49,6 @@ typedef struct {
    uint32_t registration_id;
    GHashTable *trusted;
    bool recv;
-   bool trusted_msg;
 } identity_key_store_t;
 
 GHashTable * session_store_new(void);