From ba0b81f7504fbbfcccea2bbe38e985974ee0f63c Mon Sep 17 00:00:00 2001
From: Paul Schaub <vanitasvitae@fsfe.org>
Date: Thu, 11 Jan 2018 13:21:11 +0100
Subject: [PATCH] Add method to decrypt OmemoElements from a MamQueryResult

---
 .../smackx/omemo/OmemoManager.java            | 40 +++++++--
 .../smackx/omemo/OmemoService.java            | 82 ++++++++++++++++++-
 2 files changed, 115 insertions(+), 7 deletions(-)

diff --git a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java
index e91f8ae37..a9f1a7599 100644
--- a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java
+++ b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java
@@ -25,6 +25,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.SortedSet;
 import java.util.TreeMap;
@@ -49,7 +50,9 @@ import org.jivesoftware.smackx.carbons.CarbonManager;
 import org.jivesoftware.smackx.carbons.packet.CarbonExtension;
 import org.jivesoftware.smackx.disco.ServiceDiscoveryManager;
 import org.jivesoftware.smackx.eme.element.ExplicitMessageEncryptionElement;
+import org.jivesoftware.smackx.forward.packet.Forwarded;
 import org.jivesoftware.smackx.hints.element.StoreHint;
+import org.jivesoftware.smackx.mam.MamManager;
 import org.jivesoftware.smackx.muc.MultiUserChat;
 import org.jivesoftware.smackx.muc.MultiUserChatManager;
 import org.jivesoftware.smackx.muc.RoomInfo;
@@ -396,11 +399,39 @@ public final class OmemoManager extends Manager {
         return getOmemoService().decryptMessage(managerGuard, sender, omemoElement);
     }
 
+    /**
+     * Decrypt OmemoMessages of a {@link org.jivesoftware.smackx.mam.MamManager.MamQueryResult}.
+     * Return a map of Stanzas and their decrypted counterparts.
+     *
+     * Note: This method does not repair broken sessions.
+     *
+     * @param result MamQueryResult
+     * @return map of encrypted stanzas and decrypted messages.
+     *
+     * @throws SmackException.NotLoggedInException if the OmemoManager is not authenticated
+     */
+    public Map<Stanza, OmemoMessage.Received> decryptMAMQueryResult(MamManager.MamQueryResult result)
+            throws SmackException.NotLoggedInException
+    {
+        LoggedInOmemoManager managerGuard = new LoggedInOmemoManager(this);
+        HashMap<Stanza, OmemoMessage.Received> decryptedMessages = new HashMap<>();
+
+        for (Forwarded forwarded : result.forwardedMessages) {
+            Stanza stanza = forwarded.getForwardedStanza();
+            OmemoMessage.Received decrypted = getOmemoService().decryptStanza(stanza, managerGuard);
+            if (decrypted != null) {
+                decryptedMessages.put(stanza, decrypted);
+            }
+        }
+
+        return decryptedMessages;
+    }
+
     /**
      * Trust that a fingerprint belongs to an OmemoDevice.
      * The fingerprint must be the lowercase, hexadecimal fingerprint of the identityKey of the device and must
      * be of length 64.
-     * 
+     *
      * @param device device
      * @param fingerprint fingerprint
      */
@@ -489,11 +520,8 @@ public final class OmemoManager extends Manager {
 
             // Set MAM Storage hint
             StoreHint.set(message);
-
-            if (OmemoConfiguration.getAddEmeEncryptionHint()) {
-                message.addExtension(new ExplicitMessageEncryptionElement(
-                        ExplicitMessageEncryptionElement.ExplicitMessageEncryptionProtocol.omemoVAxolotl));
-            }
+            message.addExtension(new ExplicitMessageEncryptionElement(
+                    ExplicitMessageEncryptionElement.ExplicitMessageEncryptionProtocol.omemoVAxolotl));
             connection().sendStanza(message);
         }
     }
diff --git a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoService.java b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoService.java
index 4cb392d0e..9e9f3c8a0 100644
--- a/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoService.java
+++ b/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoService.java
@@ -1079,7 +1079,7 @@ public abstract class OmemoService<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
                                           OmemoManager.LoggedInOmemoManager managerGuard) {
         OmemoManager manager = managerGuard.get();
         // Avoid the ratchet being manipulated and the bundle being published multiple times simultaneously
-        synchronized (manager) {
+        synchronized (manager.LOCK) {
             OmemoDevice userDevice = manager.getOwnDevice();
             OmemoElement element = carbonCopy.getExtension(OmemoElement.NAME_ENCRYPTED, OmemoElement_VAxolotl.NAMESPACE);
             if (element == null) {
@@ -1207,6 +1207,86 @@ public abstract class OmemoService<T_IdKeyPair, T_IdKey, T_PreKey, T_SigPreKey,
         }
     }
 
+    /**
+     * Decrypt the OmemoElement inside the given Stanza and return it.
+     * Return null if something goes wrong.
+     *
+     * @param stanza stanza
+     * @param managerGuard authenticated OmemoManager
+     * @return decrypted OmemoMessage or null
+     */
+    OmemoMessage.Received decryptStanza(Stanza stanza, OmemoManager.LoggedInOmemoManager managerGuard) {
+        OmemoManager manager = managerGuard.get();
+        // Avoid the ratchet being manipulated and the bundle being published multiple times simultaneously
+        synchronized (manager.LOCK) {
+            OmemoDevice userDevice = manager.getOwnDevice();
+            OmemoElement element = stanza.getExtension(OmemoElement.NAME_ENCRYPTED, OmemoElement_VAxolotl.NAMESPACE);
+            if (element == null) {
+                return null;
+            }
+
+            OmemoMessage.Received decrypted = null;
+            BareJid sender;
+
+            try {
+                MultiUserChat muc = getMuc(manager.getConnection(), stanza.getFrom());
+                if (muc != null) {
+                    Occupant occupant = muc.getOccupant(stanza.getFrom().asEntityFullJidIfPossible());
+                    Jid occupantJid = occupant.getJid();
+
+                    if (occupantJid == null) {
+                        LOGGER.log(Level.WARNING, "MUC message received, but there is no way to retrieve the senders Jid. " +
+                                stanza.getFrom());
+                        return null;
+                    }
+
+                    sender = occupantJid.asBareJid();
+
+                    // try is for this
+                    decrypted = decryptMessage(managerGuard, sender, element);
+
+                } else {
+                    sender = stanza.getFrom().asBareJid();
+
+                    // and this
+                    decrypted = decryptMessage(managerGuard, sender, element);
+                }
+
+                if (decrypted.isPreKeyMessage() && OmemoConfiguration.getCompleteSessionWithEmptyMessage()) {
+                    LOGGER.log(Level.FINE, "Received a preKeyMessage from " + decrypted.getSenderDevice() + ".\n" +
+                            "Complete the session by sending an empty response message.");
+                    try {
+                        sendRatchetUpdate(managerGuard, decrypted.getSenderDevice());
+                    } catch (CannotEstablishOmemoSessionException e) {
+                        throw new AssertionError("Since we successfully received a message, we MUST be able to " +
+                                "establish a session. " + e);
+                    } catch (NoSuchAlgorithmException | InterruptedException | SmackException.NotConnectedException | SmackException.NoResponseException e) {
+                        LOGGER.log(Level.WARNING, "Cannot send a ratchet update message.", e);
+                    }
+                }
+            } catch (NoRawSessionException e) {
+                OmemoDevice device = e.getDeviceWithoutSession();
+                LOGGER.log(Level.WARNING, "No raw session found for contact " + device + ". ", e);
+
+            } catch (CorruptedOmemoKeyException | CryptoFailedException e) {
+                LOGGER.log(Level.WARNING, "Could not decrypt incoming message: ", e);
+            }
+
+            // Upload fresh bundle.
+            if (getOmemoStoreBackend().loadOmemoPreKeys(userDevice).size() < OmemoConstants.PRE_KEY_COUNT_PER_BUNDLE) {
+                LOGGER.log(Level.FINE, "We used up a preKey. Upload a fresh bundle.");
+                try {
+                    getOmemoStoreBackend().replenishKeys(userDevice);
+                    OmemoBundleElement bundleElement = getOmemoStoreBackend().packOmemoBundle(userDevice);
+                    publishBundle(manager.getConnection(), userDevice, bundleElement);
+                } catch (CorruptedOmemoKeyException | InterruptedException | SmackException.NoResponseException | SmackException.NotConnectedException | XMPPException.XMPPErrorException e) {
+                    LOGGER.log(Level.WARNING, "Could not republish replenished bundle.", e);
+                }
+            }
+            return decrypted;
+        }
+    }
+
     /**
      * Fetch and process a fresh bundle and send an empty preKeyMessage in order to establish a fresh session.
      *