figure out fallbacks with omemo source id

2 files changed, 29 insertions, 0 deletions

diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java b/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
index d23e3a88b..ed432bb9b 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/AxolotlService.java
@@ -206,6 +206,19 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
             this.fillMap(store);
         }
 
+        public Set<Jid> findCounterpartsForSourceId(Integer sid) {
+            Set<Jid> candidates = new HashSet<>();
+            synchronized (MAP_LOCK) {
+                for (Map.Entry<String, Map<Integer, XmppAxolotlSession>> entry : map.entrySet()) {
+                    String key = entry.getKey();
+                    if (entry.getValue().containsKey(sid)) {
+                        candidates.add(Jid.of(key));
+                    }
+                }
+            }
+            return candidates;
+        }
+
         private void putDevicesForJid(String bareJid, List<Integer> deviceIds, SQLiteAxolotlStore store) {
             for (Integer deviceId : deviceIds) {
                 SignalProtocolAddress axolotlAddress = new SignalProtocolAddress(bareJid, deviceId);
@@ -316,6 +329,10 @@ public class AxolotlService implements OnAdvancedStreamFeaturesLoaded {
         return keys;
     }
 
+    public Set<Jid> findCounterpartsBySourceId(int sid) {
+        return sessions.findCounterpartsForSourceId(sid);
+    }
+
     public long getNumTrustedKeys(Jid jid) {
         return axolotlStore.getContactNumTrustedKeys(jid.asBareJid().toString());
     }
diff --git a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
index 603eedde8..f5a41e6b3 100644
--- a/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
+++ b/src/main/java/de/pixart/messenger/crypto/axolotl/XmppAxolotlMessage.java
@@ -88,6 +88,18 @@ public class XmppAxolotlMessage {
         }
     }
 
+    public static int parseSourceId(final Element axolotlMessage) throws IllegalArgumentException {
+        final Element header = axolotlMessage.findChild(HEADER);
+        if (header == null) {
+            throw new IllegalArgumentException("No header found");
+        }
+        try {
+            return Integer.parseInt(header.getAttribute(SOURCEID));
+        } catch (NumberFormatException e) {
+            throw new IllegalArgumentException("invalid source id");
+        }
+    }
+
     private XmppAxolotlMessage(final Element axolotlMessage, final Jid from) throws IllegalArgumentException {
         this.from = from;
         Element header = axolotlMessage.findChild(HEADER);