From 05b5705f22ff1e9ce000a856bb297f3da06488ef Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Tue, 14 Dec 2021 17:20:16 +0100
Subject: [PATCH 1/2] Don't fail on decrypting heartbeats
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
---
 slixmpp_omemo/__init__.py | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/slixmpp_omemo/__init__.py b/slixmpp_omemo/__init__.py
index e5381bd..76d3740 100644
--- a/slixmpp_omemo/__init__.py
+++ b/slixmpp_omemo/__init__.py
@@ -650,9 +650,10 @@ class XEP_0384(BasePlugin):
         allow_untrusted: bool = False,
     ) -> Optional[str]:
         header = encrypted['header']
-        if encrypted['payload']['value'] is None:
-            raise ErroneousPayload('The payload element was empty')
-        payload = b64dec(encrypted['payload']['value'])
+
+        payload = None
+        if encrypted['payload']['value'] is not None:
+            payload = b64dec(encrypted['payload']['value'])
 
         jid = sender.bare
         sid = int(header['sid'])
@@ -674,15 +675,26 @@ class XEP_0384(BasePlugin):
         # XXX: 'cipher' is part of KeyTransportMessages and is used when no payload
         # is passed. We do not implement this yet.
         try:
-            body = self._omemo.decryptMessage(
-                jid,
-                sid,
-                iv,
-                message,
-                isPrekeyMessage,
-                payload,
-                allow_untrusted=allow_untrusted,
-            )
+            if payload is None:
+                self._omemo.decryptRatchetFowardingMessage(
+                    jid,
+                    sid,
+                    iv,
+                    message,
+                    isPrekeyMessage,
+                    allow_untrusted=allow_untrusted,
+                )
+                body = None
+            else:
+                body = self._omemo.decryptMessage(
+                    jid,
+                    sid,
+                    iv,
+                    message,
+                    isPrekeyMessage,
+                    payload,
+                    allow_untrusted=allow_untrusted,
+                )
         except (omemo.exceptions.NoSessionException,):
             # This might happen when the sender is sending using a session
             # that we don't know about (deleted session storage, etc.). In

From e00c646d95689b27e63bc886510775efbcb9bf86 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maxime=20=E2=80=9Cpep=E2=80=9D=20Buquet?= <pep@bouah.net>
Date: Wed, 15 Dec 2021 00:56:19 +0100
Subject: [PATCH 2/2] echo_client: Handle decrypt_message's Optional-ness
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Also comment about it

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
---
 examples/echo_client.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/examples/echo_client.py b/examples/echo_client.py
index a4ae84e..8239b06 100644
--- a/examples/echo_client.py
+++ b/examples/echo_client.py
@@ -150,6 +150,11 @@ class EchoBot(ClientXMPP):
         try:
             encrypted = msg['omemo_encrypted']
             body = self['xep_0384'].decrypt_message(encrypted, mfrom, allow_untrusted)
+            # decrypt_message returns Optional[str]. It is possible to get
+            # body-less OMEMO message (see KeyTransportMessages), currently
+            # used for example to send heartbeats to other devices.
+            if body is None:
+                return None
             decoded = body.decode('utf8')
             if self.is_command(decoded):
                 await self.handle_command(mto, mtype, decoded)