OMEMO: Improve handling of newly added devices

3 files changed, 120 insertions, 48 deletions

diff --git a/plugins/omemo/src/logic/database.vala b/plugins/omemo/src/logic/database.vala
index 1f9e1304..6243569f 100644
--- a/plugins/omemo/src/logic/database.vala
+++ b/plugins/omemo/src/logic/database.vala
@@ -30,6 +30,10 @@ public class Database : Qlite.Database {
             return select().with(this.identity_id, "=", identity_id).with(this.address_name, "=", address_name);
         }
 
+        public QueryBuilder get_with_device_id(int identity_id, int device_id) {
+            return select().with(this.identity_id, "=", identity_id).with(this.device_id, "=", device_id);
+        }
+
         public void insert_device_list(int32 identity_id, string address_name, ArrayList<int32> devices) {
             update().with(this.identity_id, "=", identity_id).with(this.address_name, "=", address_name).set(now_active, false).perform();
             foreach (int32 device_id in devices) {
@@ -49,7 +53,22 @@ public class Database : Qlite.Database {
             string identity_key = Base64.encode(bundle.identity_key.serialize());
             RowOption row = with_address(identity_id, address_name).with(this.device_id, "=", device_id).single().row();
             if (row.is_present() && row[identity_key_public_base64] != null && row[identity_key_public_base64] != identity_key) {
-                error("Tried to change the identity key for a known device id. Likely an attack.");
+                critical("Tried to change the identity key for a known device id. Likely an attack.");
+                return -1;
+            }
+            return upsert()
+                    .value(this.identity_id, identity_id, true)
+                    .value(this.address_name, address_name, true)
+                    .value(this.device_id, device_id, true)
+                    .value(this.identity_key_public_base64, identity_key)
+                    .value(this.trust_level, trust).perform();
+        }
+
+        public int64 insert_device_session(int32 identity_id, string address_name, int device_id, string identity_key, TrustLevel trust) {
+            RowOption row = with_address(identity_id, address_name).with(this.device_id, "=", device_id).single().row();
+            if (row.is_present() && row[identity_key_public_base64] != null && row[identity_key_public_base64] != identity_key) {
+                critical("Tried to change the identity key for a known device id. Likely an attack.");
+                return -1;
             }
             return upsert()
                     .value(this.identity_id, identity_id, true)
@@ -86,10 +105,6 @@ public class Database : Qlite.Database {
             return this.with_address(identity_id, address_name)
                 .with(this.device_id, "=", device_id).single().row().inner;
         }
-
-        public QueryBuilder get_with_device_id(int device_id) {
-            return select().with(this.device_id, "=", device_id);
-        }
     }
 
 
@@ -104,10 +119,11 @@ public class Database : Qlite.Database {
             index("trust_idx", {identity_id, address_name}, true);
         }
 
-        public bool get_blind_trust(int32 identity_id, string address_name) {
-            return this.select().with(this.identity_id, "=", identity_id)
-                    .with(this.address_name, "=", address_name)
-                    .with(this.blind_trust, "=", true).count() > 0;
+        public bool get_blind_trust(int32 identity_id, string address_name, bool def = false) {
+            RowOption row = this.select().with(this.identity_id, "=", identity_id)
+                    .with(this.address_name, "=", address_name).single().row();
+            if (row.is_present()) return row[blind_trust];
+            return def;
         }
     }
 
diff --git a/plugins/omemo/src/logic/manager.vala b/plugins/omemo/src/logic/manager.vala
index 53e02e37..2bbd918c 100644
--- a/plugins/omemo/src/logic/manager.vala
+++ b/plugins/omemo/src/logic/manager.vala
@@ -66,7 +66,6 @@ public class Manager : StreamInteractionModule, Object {
         this.trust_manager = trust_manager;
 
         stream_interactor.stream_negotiated.connect(on_stream_negotiated);
-        stream_interactor.account_added.connect(on_account_added);
         stream_interactor.get_module(MessageProcessor.IDENTITY).pre_message_send.connect(on_pre_message_send);
         stream_interactor.get_module(RosterManager.IDENTITY).mutual_subscription.connect(on_mutual_subscription);
     }
@@ -171,14 +170,15 @@ public class Manager : StreamInteractionModule, Object {
         stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).request_user_devicelist.begin((!)stream, jid);
     }
 
-    private void on_account_added(Account account) {
-        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).store_created.connect((store) => on_store_created.begin(account, store));
-        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).device_list_loaded.connect((jid, devices) => on_device_list_loaded(account, jid, devices));
-        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).bundle_fetched.connect((jid, device_id, bundle) => on_bundle_fetched(account, jid, device_id, bundle));
-    }
-
     private void on_stream_negotiated(Account account, XmppStream stream) {
-        stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY).request_user_devicelist.begin(stream, account.bare_jid);
+        StreamModule module = stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY);
+        if (module != null) {
+            module.request_user_devicelist.begin(stream, account.bare_jid);
+            module.device_list_loaded.connect((jid, devices) => on_device_list_loaded(account, jid, devices));
+            module.bundle_fetched.connect((jid, device_id, bundle) => on_bundle_fetched(account, jid, device_id, bundle));
+            module.bundle_fetch_failed.connect((jid) => continue_message_sending(account, jid));
+        }
+        initialize_store.begin(account);
     }
 
     private void on_device_list_loaded(Account account, Jid jid, ArrayList<int32> device_list) {
@@ -202,7 +202,7 @@ public class Manager : StreamInteractionModule, Object {
         //Fetch the bundle for each new device
         int inc = 0;
         foreach (Row row in db.identity_meta.get_unknown_devices(identity_id, jid.bare_jid.to_string())) {
-            module.fetch_bundle(stream, Jid.parse(row[db.identity_meta.address_name]), row[db.identity_meta.device_id]);
+            module.fetch_bundle(stream, Jid.parse(row[db.identity_meta.address_name]), row[db.identity_meta.device_id], false);
             inc++;
         }
         if (inc > 0) {
@@ -245,7 +245,7 @@ public class Manager : StreamInteractionModule, Object {
         int identity_id = db.identity.get_id(account.id);
         if (identity_id < 0) return;
 
-        bool blind_trust = db.trust.get_blind_trust(identity_id, jid.bare_jid.to_string());
+        bool blind_trust = db.trust.get_blind_trust(identity_id, jid.bare_jid.to_string(), true);
 
         //If we don't blindly trust new devices and we haven't seen this key before then don't trust it
         bool untrust = !(blind_trust || db.identity_meta.with_address(identity_id, jid.bare_jid.to_string())
@@ -269,30 +269,44 @@ public class Manager : StreamInteractionModule, Object {
         //Update the database with the appropriate trust information
         db.identity_meta.insert_device_bundle(identity_id, jid.bare_jid.to_string(), device_id, bundle, trusted);
 
-        XmppStream? stream = stream_interactor.get_stream(account);
-        if(stream == null) return;
-        StreamModule? module = ((!)stream).get_module(StreamModule.IDENTITY);
-        if(module == null) return;
+        if (should_start_session(account, jid)) {
+            XmppStream? stream = stream_interactor.get_stream(account);
+            if (stream != null) {
+                StreamModule? module = ((!)stream).get_module(StreamModule.IDENTITY);
+                if (module != null) {
+                    module.start_session(stream, jid, device_id, bundle);
+                }
+            }
+        }
+        continue_message_sending(account, jid);
+    }
 
-        //Get all messages waiting on the bundle and determine if they can now be sent
-        HashSet<Entities.Message> send_now = new HashSet<Entities.Message>();
+    private bool should_start_session(Account account, Jid jid) {
         lock (message_states) {
             foreach (Entities.Message msg in message_states.keys) {
+                if (!msg.account.equals(account)) continue;
+                Gee.List<Jid> occupants = get_occupants(msg.counterpart.bare_jid, account);
+                if (account.bare_jid.equals(jid) || (msg.counterpart != null && (msg.counterpart.equals_bare(jid) || occupants.contains(jid)))) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
 
-                bool session_created = true;
+    private void continue_message_sending(Account account, Jid jid) {
+        //Get all messages waiting and determine if they can now be sent
+        HashSet<Entities.Message> send_now = new HashSet<Entities.Message>();
+        lock (message_states) {
+            foreach (Entities.Message msg in message_states.keys) {
                 if (!msg.account.equals(account)) continue;
                 Gee.List<Jid> occupants = get_occupants(msg.counterpart.bare_jid, account);
 
                 MessageState state = message_states[msg];
 
-                if (trusted == TrustLevel.TRUSTED || trusted == TrustLevel.VERIFIED) {
-                    if(account.bare_jid.equals(jid) || (msg.counterpart != null && (msg.counterpart.equals_bare(jid) || occupants.contains(jid)))) {
-                        session_created = module.start_session(stream, jid, device_id, bundle);
-                    }
-                }
-                if (account.bare_jid.equals(jid) && session_created) {
+                if (account.bare_jid.equals(jid)) {
                     state.waiting_own_sessions--;
-                } else if (msg.counterpart != null && (msg.counterpart.equals_bare(jid) || occupants.contains(jid)) && session_created) {
+                } else if (msg.counterpart != null && (msg.counterpart.equals_bare(jid) || occupants.contains(jid))) {
                     state.waiting_other_sessions--;
                 }
                 if (state.should_retry_now()){
@@ -309,12 +323,15 @@ public class Manager : StreamInteractionModule, Object {
         }
     }
 
-    private async void on_store_created(Account account, Store store) {
+    private async void initialize_store(Account account) {
         // If the account is not yet persisted, wait for that and then continue - without identity.account_id the entry isn't worth much.
         if (account.id == -1) {
-            account.notify["id"].connect(() => on_store_created.callback());
+            account.notify["id"].connect(() => initialize_store.callback());
             yield;
         }
+        StreamModule? module = stream_interactor.module_manager.get_module(account, StreamModule.IDENTITY);
+        if (module == null) return;
+        Store store = module.store;
         Qlite.Row? row = db.identity.row_with(db.identity.account_id, account.id).inner;
         int identity_id = -1;
         bool publish_identity = false;
@@ -354,12 +371,9 @@ public class Manager : StreamInteractionModule, Object {
         }
 
         // Generated new device ID, ensure this gets added to the devicelist
-        if (publish_identity) {
-            XmppStream? stream = stream_interactor.get_stream(account);
-            if (stream == null) return;
-            StreamModule? module = ((!)stream).get_module(StreamModule.IDENTITY);
-            if(module == null) return;
-            module.request_user_devicelist.begin(stream, account.bare_jid);
+        XmppStream? stream = stream_interactor.get_stream(account);
+        if (stream != null) {
+            module.request_user_devicelist.begin((!)stream, account.bare_jid);
         }
     }
 
diff --git a/plugins/omemo/src/logic/trust_manager.vala b/plugins/omemo/src/logic/trust_manager.vala
index 658e55ef..1b8a9436 100644
--- a/plugins/omemo/src/logic/trust_manager.vala
+++ b/plugins/omemo/src/logic/trust_manager.vala
@@ -65,6 +65,7 @@ public class TrustManager {
     private StanzaNode create_encrypted_key_node(uint8[] key, Address address, Store store) throws GLib.Error {
         SessionCipher cipher = store.create_session_cipher(address);
         CiphertextMessage device_key = cipher.encrypt(key);
+        debug("Created encrypted key for %s/%d", address.name, address.device_id);
         StanzaNode key_node = new StanzaNode.build("key", NS_URI)
             .put_attribute("rid", address.device_id.to_string())
             .put_node(new StanzaNode.text(Base64.encode(device_key.serialized)));
@@ -181,7 +182,7 @@ public class TrustManager {
     public bool is_known_address(Account account, Jid jid) {
         int identity_id = db.identity.get_id(account.id);
         if (identity_id < 0) return false;
-        return db.identity_meta.with_address(identity_id, jid.to_string()).count() > 0;
+        return db.identity_meta.with_address(identity_id, jid.to_string()).with(db.identity_meta.last_active, ">", 0).count() > 0;
     }
 
     public Gee.List<int32> get_trusted_devices(Account account, Jid jid) {
@@ -261,7 +262,8 @@ public class TrustManager {
         }
 
         public override async bool run(Entities.Message message, Xmpp.MessageStanza stanza, Conversation conversation) {
-            Store store = stream_interactor.module_manager.get_module(conversation.account, StreamModule.IDENTITY).store;
+            StreamModule module = stream_interactor.module_manager.get_module(conversation.account, StreamModule.IDENTITY);
+            Store store = module.store;
 
             StanzaNode? _encrypted = stanza.stanza.get_subnode("encrypted", NS_URI);
             if (_encrypted == null || MessageFlag.get_flag(stanza) != null || stanza.from == null) return false;
@@ -270,6 +272,7 @@ public class TrustManager {
                 message.body = "[This message is OMEMO encrypted]"; // TODO temporary
             };
             if (!Plugin.ensure_context()) return false;
+            int identity_id = db.identity.get_id(conversation.account.id);
             MessageFlag flag = new MessageFlag();
             stanza.add_flag(flag);
             StanzaNode? _header = encrypted.get_subnode("header");
@@ -278,6 +281,7 @@ public class TrustManager {
             int sid = header.get_attribute_int("sid");
             if (sid <= 0) return false;
             foreach (StanzaNode key_node in header.get_subnodes("key")) {
+                debug("Is ours? %d =? %u", key_node.get_attribute_int("rid"), store.local_registration_id);
                 if (key_node.get_attribute_int("rid") == store.local_registration_id) {
 
                     string? payload = encrypted.get_deep_string_content("payload");
@@ -289,27 +293,63 @@ public class TrustManager {
                     uint8[] iv = Base64.decode((!)iv_node);
                     Gee.List<Jid> possible_jids = new ArrayList<Jid>();
                     if (conversation.type_ == Conversation.Type.CHAT) {
-                        possible_jids.add(stanza.from);
+                        possible_jids.add(stanza.from.bare_jid);
                     } else {
                         Jid? real_jid = message.real_jid;
                         if (real_jid != null) {
-                            possible_jids.add(real_jid);
+                            possible_jids.add(real_jid.bare_jid);
+                        } else if (key_node.get_attribute_bool("prekey")) {
+                            // pre key messages do store the identity key, so we can use that to find the real jid
+                            PreKeySignalMessage msg = Plugin.get_context().deserialize_pre_key_signal_message(Base64.decode((!)key_node_content));
+                            string identity_key = Base64.encode(msg.identity_key.serialize());
+                            foreach (Row row in db.identity_meta.get_with_device_id(identity_id, sid).with(db.identity_meta.identity_key_public_base64, "=", identity_key)) {
+                                possible_jids.add(new Jid(row[db.identity_meta.address_name]));
+                            }
+                            if (possible_jids.size != 1) {
+                                continue;
+                            }
                         } else {
                             // If we don't know the device name (MUC history w/o MAM), test decryption with all keys with fitting device id
-                            foreach (Row row in db.identity_meta.get_with_device_id(sid)) {
+                            foreach (Row row in db.identity_meta.get_with_device_id(identity_id, sid)) {
                                 possible_jids.add(new Jid(row[db.identity_meta.address_name]));
                             }
                         }
                     }
 
+                    if (possible_jids.size == 0) {
+                        debug("Received message from unknown entity with device id %d", sid);
+                    }
+
                     foreach (Jid possible_jid in possible_jids) {
                         try {
-                            Address address = new Address(possible_jid.bare_jid.to_string(), header.get_attribute_int("sid"));
+                            Address address = new Address(possible_jid.to_string(), sid);
                             if (key_node.get_attribute_bool("prekey")) {
+                                Row? device = db.identity_meta.get_device(identity_id, possible_jid.to_string(), sid);
                                 PreKeySignalMessage msg = Plugin.get_context().deserialize_pre_key_signal_message(Base64.decode((!)key_node_content));
+                                string identity_key = Base64.encode(msg.identity_key.serialize());
+                                if (device != null && device[db.identity_meta.identity_key_public_base64] != null) {
+                                    if (device[db.identity_meta.identity_key_public_base64] != identity_key) {
+                                        critical("Tried to use a different identity key for a known device id.");
+                                        continue;
+                                    }
+                                } else {
+                                    debug("Learn new device from incoming message from %s/%d", possible_jid.to_string(), sid);
+                                    bool blind_trust = db.trust.get_blind_trust(identity_id, possible_jid.to_string(), true);
+                                    if (db.identity_meta.insert_device_session(identity_id, possible_jid.to_string(), sid, identity_key, blind_trust ? TrustLevel.TRUSTED : TrustLevel.UNKNOWN) < 0) {
+                                        critical("Failed learning a device.");
+                                        continue;
+                                    }
+                                    XmppStream? stream = stream_interactor.get_stream(conversation.account);
+                                    if (device == null && stream != null) {
+                                        module.request_user_devicelist.begin(stream, possible_jid);
+                                    }
+                                }
+                                debug("Starting new session for decryption with device from %s/%d", possible_jid.to_string(), sid);
                                 SessionCipher cipher = store.create_session_cipher(address);
                                 key = cipher.decrypt_pre_key_signal_message(msg);
+                                // TODO: Finish session
                             } else {
+                                debug("Continuing session for decryption with device from %s/%d", possible_jid.to_string(), sid);
                                 SignalMessage msg = Plugin.get_context().deserialize_signal_message(Base64.decode((!)key_node_content));
                                 SessionCipher cipher = store.create_session_cipher(address);
                                 key = cipher.decrypt_signal_message(msg);
@@ -332,6 +372,7 @@ public class TrustManager {
                             message.encryption = Encryption.OMEMO;
                             flag.decrypted = true;
                         } catch (Error e) {
+                            debug("Decrypting message from %s/%d failed: %s", possible_jid.to_string(), sid, e.message);
                             continue;
                         }
 
@@ -339,10 +380,11 @@ public class TrustManager {
                         if (conversation.type_ == Conversation.Type.GROUPCHAT && message.real_jid == null) {
                             message.real_jid = possible_jid;
                         }
-                        break;
+                        return false;
                     }
                 }
             }
+            debug("Received OMEMO encryped message that could not be decrypted.");
             return false;
         }