--- /home/cmeng/source/smack/Smack-4.4.git/smack-omemo/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java	2022-02-18 08:54:57.000000000 
+++ /home/cmeng/workspace/android/atalk-android/aTalk/src/main/java/org/jivesoftware/smackx/omemo/OmemoManager.java	2024-03-14 06:19:01.000000000 
@@ -32,24 +32,29 @@
 import java.util.WeakHashMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import org.jivesoftware.smack.ConnectionListener;
 import org.jivesoftware.smack.Manager;
+import org.jivesoftware.smack.SmackConfiguration;
 import org.jivesoftware.smack.SmackException;
 import org.jivesoftware.smack.XMPPConnection;
 import org.jivesoftware.smack.XMPPException;
+import org.jivesoftware.smack.packet.ExtensionElement;
 import org.jivesoftware.smack.packet.Message;
 import org.jivesoftware.smack.packet.MessageBuilder;
 import org.jivesoftware.smack.packet.Stanza;
+import org.jivesoftware.smack.packet.StanzaBuilder;
 import org.jivesoftware.smack.util.Async;
 
+import org.jivesoftware.smack.util.stringencoder.Base64;
 import org.jivesoftware.smackx.carbons.CarbonManager;
 import org.jivesoftware.smackx.carbons.packet.CarbonExtension;
 import org.jivesoftware.smackx.disco.ServiceDiscoveryManager;
 import org.jivesoftware.smackx.hints.element.StoreHint;
+import org.jivesoftware.smackx.jet.JingleEnvelopeManager;
 import org.jivesoftware.smackx.mam.MamManager;
 import org.jivesoftware.smackx.muc.MultiUserChat;
 import org.jivesoftware.smackx.muc.MultiUserChatManager;
 import org.jivesoftware.smackx.muc.RoomInfo;
 import org.jivesoftware.smackx.omemo.element.OmemoBundleElement;
 import org.jivesoftware.smackx.omemo.element.OmemoDeviceListElement;
@@ -76,34 +81,40 @@
 import org.jivesoftware.smackx.pubsub.packet.PubSub;
 
 import org.jxmpp.jid.BareJid;
 import org.jxmpp.jid.DomainBareJid;
 import org.jxmpp.jid.EntityBareJid;
 import org.jxmpp.jid.EntityFullJid;
+import org.jxmpp.jid.FullJid;
 
 /**
  * Manager that allows sending messages encrypted with OMEMO.
  * This class also provides some methods useful for a client that implements OMEMO.
  *
  * @author Paul Schaub
+ * @author Eng Chong Meng
  */
 
-public final class OmemoManager extends Manager {
+public final class OmemoManager extends Manager implements JingleEnvelopeManager {
     private static final Logger LOGGER = Logger.getLogger(OmemoManager.class.getName());
 
+    public static final int SMACK_REPLY_OMEMO_PUBLISH = 15000;
+
     private static final Integer UNKNOWN_DEVICE_ID = -1;
 
     private static final WeakHashMap<XMPPConnection, TreeMap<Integer, OmemoManager>> INSTANCES = new WeakHashMap<>();
     private final OmemoService<?, ?, ?, ?, ?, ?, ?, ?, ?> service;
 
     private final HashSet<OmemoMessageListener> omemoMessageListeners = new HashSet<>();
     private final HashSet<OmemoMucMessageListener> omemoMucMessageListeners = new HashSet<>();
 
     private final PepManager pepManager;
 
     private OmemoTrustCallback trustCallback;
+
+    private Set<Integer> publishedDeviceList;
 
     private BareJid ownJid;
     private Integer deviceId;
 
     /**
      * Private constructor.
@@ -259,12 +270,13 @@
                 try {
                     initialize();
                     finishedCallback.initializationFinished(OmemoManager.this);
                 } catch (Exception e) {
                     finishedCallback.initializationFailed(e);
                 }
+                getConnection().setReplyTimeout(SmackConfiguration.getDefaultReplyTimeout());
             }
         });
     }
 
     /**
      * Return a set of all OMEMO capable devices of a contact.
@@ -935,12 +947,17 @@
      * @return the connection of this manager
      */
     XMPPConnection getConnection() {
         return connection();
     }
 
+    @Override
+    public String getJingleEnvelopeNamespace() {
+        return OMEMO_NAMESPACE_V_AXOLOTL;
+    }
+
     /**
      * Return the OMEMO service object.
      *
      * @return the OmemoService object related to this OmemoManager.
      */
     OmemoService<?, ?, ?, ?, ?, ?, ?, ?, ?> getOmemoService() {
@@ -1005,25 +1022,31 @@
                             "IOException while processing OMEMO PEP device updates. Message: " + message,
                                 e);
             return;
         }
         final OmemoDeviceListElement_VAxolotl newDeviceList = new OmemoDeviceListElement_VAxolotl(deviceList);
 
-        if (!newDeviceList.copyDeviceIds().equals(receivedDeviceList.copyDeviceIds())) {
+        if (!newDeviceList.copyDeviceIds().equals(receivedDeviceList.copyDeviceIds())
+                && !newDeviceList.copyDeviceIds().equals(publishedDeviceList)) {
             LOGGER.log(Level.FINE, "Republish deviceList due to changes:" +
                             " Received: " + Arrays.toString(receivedDeviceList.copyDeviceIds().toArray()) +
                             " Published: " + Arrays.toString(newDeviceList.copyDeviceIds().toArray()));
             Async.go(new Runnable() {
                 @Override
                 public void run() {
                     try {
+                        // LOGGER.log(Level.INFO, "received (new) DeviceList: " + receivedDeviceList.getDeviceIds()
+                        //        + " (" + newDeviceList.getDeviceIds() + ")");
+                        getConnection().setReplyTimeout(SMACK_REPLY_OMEMO_PUBLISH);
                         OmemoService.publishDeviceList(connection(), newDeviceList);
+                        publishedDeviceList = newDeviceList.copyDeviceIds();
                     } catch (InterruptedException | XMPPException.XMPPErrorException |
                                     SmackException.NotConnectedException | SmackException.NoResponseException | PubSubException.NotALeafNodeException e) {
                         LOGGER.log(Level.WARNING, "Could not publish our deviceList upon an received update.", e);
                     }
+                    getConnection().setReplyTimeout(SmackConfiguration.getDefaultReplyTimeout());
                 }
             });
         }
     };
 
     /**
@@ -1095,7 +1118,46 @@
                 manager.setDeviceId(storedDeviceIds.first());
             } else {
                 manager.setDeviceId(randomDeviceId());
             }
         }
     }
+
+    @Override
+    public ExtensionElement encryptJingleTransfer(FullJid recipient, byte[] keyData) throws JingleEncryptionException, InterruptedException, NoSuchAlgorithmException, SmackException.NotConnectedException, SmackException.NoResponseException {
+        BareJid bareJid = recipient.asBareJid();
+        Message sendMessage;
+        try {
+            OmemoMessage.Sent encryptedMessage = encrypt(bareJid, Base64.encodeToString(keyData));
+            MessageBuilder messageBuilder = StanzaBuilder.buildMessage();
+            sendMessage = encryptedMessage.buildMessage(messageBuilder, bareJid);
+
+        } catch (CryptoFailedException | UndecidedOmemoIdentityException | SmackException.NotLoggedInException | IOException e) {
+            throw new JingleEncryptionException(e);
+        }
+
+        ExtensionElement encryptionElement = sendMessage.getExtensionElement(OmemoElement.NAME_ENCRYPTED, OMEMO_NAMESPACE_V_AXOLOTL);
+        if (encryptionElement == null) {
+            throw new AssertionError("OmemoElement MUST NOT be null.");
+        }
+        return encryptionElement;
+    }
+
+    @Override
+    public byte[] decryptJingleTransfer(FullJid sender, ExtensionElement envelope) throws JingleEncryptionException, InterruptedException, XMPPException.XMPPErrorException, SmackException.NotConnectedException, SmackException.NoResponseException {
+        if (!envelope.getNamespace().equals(OMEMO_NAMESPACE_V_AXOLOTL)
+                || !envelope.getElementName().equals(OmemoElement.NAME_ENCRYPTED)) {
+            throw new IllegalArgumentException("Passed ExtensionElement MUST be an OmemoElement!");
+        }
+
+        OmemoElement omemoElement = (OmemoElement) envelope;
+
+        OmemoMessage.Received decryptedPseudoMessage;
+        try {
+            decryptedPseudoMessage = decrypt(sender.asBareJid(), omemoElement);
+        } catch (CryptoFailedException | CorruptedOmemoKeyException | NoRawSessionException | SmackException.NotLoggedInException | IOException e) {
+            throw new JingleEncryptionException(e);
+        }
+
+        return Base64.decode(decryptedPseudoMessage.getBody());
+    }
 }
