Commit 9b55df68 authored by Maxime Buquet's avatar Maxime Buquet

should_heartbeat: also return True on new sessions

And the docstring now reflects the reality again!

We're parsing the Encrypted dict again, when we just did it in
decrypt_message above, but this function is also part of the API and
doing that for them is the least we can do.

Maybe there should be an internal function that we can call from
decrypt_message, that also gets called by should_heartbeat.
Signed-off-by: Maxime Buquet's avatarMaxime “pep” Buquet <pep@bouah.net>
parent 27da35b7
Pipeline #4061 passed with stage
......@@ -543,7 +543,7 @@ class XEP_0384(BasePlugin):
return lengths
def should_heartbeat(self, jid: JID) -> bool:
def should_heartbeat(self, jid: JID, msg: Union[Message, Encrypted]) -> bool:
"""
Returns whether we should send a heartbeat message for JID.
See notes about heartbeat in
......@@ -554,9 +554,28 @@ class XEP_0384(BasePlugin):
sessions hasn't been answered in a while.
"""
new_session: bool = False
inactive_session: bool = False
# Is the message is a prekey message. If so assume it's a new session
encrypted = msg
if isinstance(msg, Message):
encrypted = msg['omemo_encrypted']
header = encrypted['header']
key = header.xml.find("{%s}key[@rid='%s']" % (
OMEMO_BASE_NS, self._device_id))
if key is not None:
key = Key(key)
new_session = key['prekey'] in TRUE_VALUES
# Otherwise find how many message we haven't replied to in all of the
# 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)
return max(lengths, default=0) > self.heartbeat_after
inactive_session = max(lengths, default=0) > self.heartbeat_after
return new_session or inactive_session
async def make_heartbeat(self, jid: JID) -> Message:
"""
......@@ -663,7 +682,7 @@ class XEP_0384(BasePlugin):
finally:
asyncio.ensure_future(self._publish_bundle())
if self.auto_heartbeat and self.should_heartbeat(sender):
if self.auto_heartbeat and self.should_heartbeat(sender, encrypted):
async def send_heartbeat():
msg = await self.make_heartbeat(JID(jid))
msg.send()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment