pep/slixmpp-omemo
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

should_heartbeat: factor out in internal method to avoid reparsing message

Browse source 
Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>
This commit is contained in:
Maxime “pep” Buquet 2021-07-17 20:18:33 +02:00
parent bb52d93241
commit 26665d9e6a
1 changed files with 28 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
slixmpp_omemo/__init__.py
View file

@ -543,6 +543,26 @@ class XEP_0384(BasePlugin):
return lengths return lengths
def _should_heartbeat(self, jid: JID, prekey: bool) -> bool:
"""
Internal helper for :py:func:`XEP_0384.should_heartbeat`.
Returns whether we should send a heartbeat message for JID.
We check if the message is a prekey message, in which case we
assume it's a new session and we want to ACK relatively early.
Otherwise we look at the number of messages since we have last
replied and if above a certain threshold we notify them that we're
still active.
"""
receiving_chain_lengths = self._chain_lengths(jid).get('receiving', [])
lengths = map(lambda d_l: d_l[1], receiving_chain_lengths)
inactive_session = max(lengths, default=0) > self.heartbeat_after
return prekey or inactive_session
def should_heartbeat(self, jid: JID, msg: Union[Message, Encrypted]) -> bool: def should_heartbeat(self, jid: JID, msg: Union[Message, Encrypted]) -> bool:
""" """
Returns whether we should send a heartbeat message for JID. Returns whether we should send a heartbeat message for JID.
@ -554,10 +574,9 @@ class XEP_0384(BasePlugin):
sessions hasn't been answered in a while. sessions hasn't been answered in a while.
""" """
new_session: bool = False prekey: bool = False
inactive_session: bool = False
# Is the message is a prekey message. If so assume it's a new session # Get prekey information from message
encrypted = msg encrypted = msg
if isinstance(msg, Message): if isinstance(msg, Message):
encrypted = msg['omemo_encrypted'] encrypted = msg['omemo_encrypted']
@ -565,17 +584,13 @@ class XEP_0384(BasePlugin):
header = encrypted['header'] header = encrypted['header']
key = header.xml.find("{%s}key[@rid='%s']" % ( key = header.xml.find("{%s}key[@rid='%s']" % (
OMEMO_BASE_NS, self._device_id)) OMEMO_BASE_NS, self._device_id))
# Don't error out. If it's not encrypted to us we don't need to send a
# heartbeat.
if key is not None: if key is not None:
key = Key(key) key = Key(key)
new_session = key['prekey'] in TRUE_VALUES prekey = key['prekey'] in TRUE_VALUES
# Otherwise find how many message we haven't replied to in all of the return self._should_heartbeat(jid, prekey)
# sessions associated to this JID.
receiving_chain_lengths = self._chain_lengths(jid).get('receiving', [])
lengths = map(lambda d_l: d_l[1], receiving_chain_lengths)
inactive_session = max(lengths, default=0) > self.heartbeat_after
return new_session or inactive_session
async def make_heartbeat(self, jid: JID) -> Message: async def make_heartbeat(self, jid: JID) -> Message:
""" """
@ -682,8 +697,9 @@ class XEP_0384(BasePlugin):
finally: finally:
asyncio.ensure_future(self._publish_bundle()) asyncio.ensure_future(self._publish_bundle())
if self.auto_heartbeat and self.should_heartbeat(sender, encrypted): if self.auto_heartbeat and self._should_heartbeat(sender, isPrekeyMessage):
async def send_heartbeat(): async def send_heartbeat():
log.debug('Sending a heartbeat message')
msg = await self.make_heartbeat(JID(jid)) msg = await self.make_heartbeat(JID(jid))
msg.send() msg.send()
asyncio.ensure_future(send_heartbeat()) asyncio.ensure_future(send_heartbeat())