#!/usr/bin/env python3
# SàT plugin for managing Stream-Management
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
from libervia.backend import G
from libervia.backend.core.i18n import _
from libervia.backend.core.constants import Const as C
from libervia.backend.core import exceptions
from libervia.backend.core.log import getLogger
from twisted.words.protocols.jabber import client as jabber_client
from twisted.words.protocols.jabber import xmlstream
from twisted.words.xish import domish
from twisted.internet import defer
from twisted.internet import task, reactor
from functools import partial
from tx_xmpp import disco, itx_xmpp
from zope.interface import implementer
import collections
import time
log = getLogger(__name__)
PLUGIN_INFO = {
C.PI_NAME: "Stream Management",
C.PI_IMPORT_NAME: "XEP-0198",
C.PI_TYPE: "XEP",
C.PI_MODES: [C.PLUG_MODE_CLIENT],
C.PI_PROTOCOLS: ["XEP-0198"],
C.PI_DEPENDENCIES: [],
C.PI_RECOMMENDATIONS: ["XEP-0045", "XEP-0313"],
C.PI_MAIN: "XEP_0198",
C.PI_HANDLER: "yes",
C.PI_DESCRIPTION: _("""Implementation of Stream Management"""),
}
NS_SM = "urn:xmpp:sm:3"
SM_ENABLED = '/enabled[@xmlns="' + NS_SM + '"]'
SM_RESUMED = '/resumed[@xmlns="' + NS_SM + '"]'
SM_FAILED = '/failed[@xmlns="' + NS_SM + '"]'
SM_R_REQUEST = '/r[@xmlns="' + NS_SM + '"]'
SM_A_REQUEST = '/a[@xmlns="' + NS_SM + '"]'
SM_H_REQUEST = '/h[@xmlns="' + NS_SM + '"]'
# Max number of stanza to send before requesting ack
MAX_STANZA_ACK_R = 5
# Max number of seconds before requesting ack
MAX_DELAY_ACK_R = 30
MAX_COUNTER = 2**32
RESUME_MAX = 5 * 60
# if we don't have an answer to ACK REQUEST after this delay, connection is aborted
ACK_TIMEOUT = 35
class ProfileSessionData(object):
out_counter = 0
in_counter = 0
session_id = None
location = None
session_max = None
# True when an ack answer is expected
ack_requested = False
last_ack_r = 0
disconnected_time = None
def __init__(self, callback, **kw):
self.buffer = collections.deque()
self.buffer_idx = 0
self._enabled = False
self.timer = None
# time used when doing a ack request
# when it times out, connection is aborted
self.req_timer = None
self.callback_data = (callback, kw)
@property
def enabled(self):
return self._enabled
@enabled.setter
def enabled(self, enabled):
if enabled:
if self._enabled:
raise exceptions.InternalError("Stream Management can't be enabled twice")
self._enabled = True
callback, kw = self.callback_data
self.timer = task.LoopingCall(callback, **kw)
self.timer.start(MAX_DELAY_ACK_R, now=False)
else:
self._enabled = False
if self.timer is not None:
self.timer.stop()
self.timer = None
@property
def resume_enabled(self):
return self.session_id is not None
def reset(self):
self.enabled = False
self.buffer.clear()
self.buffer_idx = 0
self.in_counter = self.out_counter = 0
self.session_id = self.location = None
self.ack_requested = False
self.last_ack_r = 0
if self.req_timer is not None:
if self.req_timer.active():
log.error("req_timer has been called/cancelled but not reset")
else:
self.req_timer.cancel()
self.req_timer = None
def get_buffer_copy(self):
return list(self.buffer)
class XEP_0198(object):
# FIXME: location is not handled yet
def __init__(self, host):
log.info(_("Plugin Stream Management initialization"))
self.host = host
host.register_namespace("sm", NS_SM)
host.trigger.add("stream_hooks", self.add_hooks)
host.trigger.add("xml_init", self._xml_init_trigger)
host.trigger.add("disconnecting", self._disconnecting_trigger)
host.trigger.add("disconnected", self._disconnected_trigger)
self._ack_timeout = G.config.common.ack_timeout
if not self._ack_timeout:
log.info(_("Ack timeout disabled"))
else:
log.info(_("Ack timeout set to {timeout}s").format(timeout=self._ack_timeout))
def profile_connecting(self, client):
client._xep_0198_session = ProfileSessionData(
callback=self.check_acks, client=client
)
def get_handler(self, client):
return XEP_0198_handler(self)
def add_hooks(self, client, receive_hooks, send_hooks):
"""Add hooks to handle in/out stanzas counters"""
receive_hooks.append(partial(self.on_receive, client=client))
send_hooks.append(partial(self.on_send, client=client))
return True
def _xml_init_trigger(self, client):
"""Enable or resume a stream mangement"""
if not (NS_SM, "sm") in client.xmlstream.features:
log.warning(
_(
"Your server doesn't support stream management ({namespace}), this is "
"used to improve connection problems detection (like network outages). "
"Please ask your server administrator to enable this feature.".format(
namespace=NS_SM
)
)
)
return True
session = client._xep_0198_session
# a disconnect timer from a previous disconnection may still be active
try:
disconnect_timer = session.disconnect_timer
except AttributeError:
pass
else:
if disconnect_timer.active():
disconnect_timer.cancel()
del session.disconnect_timer
if session.resume_enabled:
# we are resuming a session
resume_elt = domish.Element((NS_SM, "resume"))
resume_elt["h"] = str(session.in_counter)
resume_elt["previd"] = session.session_id
client.send(resume_elt)
session.resuming = True
# session.enabled will be set on reception
return False
else:
# we start a new session
assert session.out_counter == 0
enable_elt = domish.Element((NS_SM, "enable"))
enable_elt["resume"] = "true"
client.send(enable_elt)
session.enabled = True
return True
def _disconnecting_trigger(self, client):
if client.is_component:
return True
session = client._xep_0198_session
if session.enabled:
self.send_ack(client)
# This is a requested disconnection, so we can reset the session
# to disable resuming and close normally the stream
session.reset()
return True
def _disconnected_trigger(self, client, reason):
if client.is_component:
return True
session = client._xep_0198_session
session.enabled = False
if session.resume_enabled:
session.disconnected_time = time.time()
session.disconnect_timer = reactor.callLater(
session.session_max, client.disconnect_profile, reason
)
# disconnect_profile must not be called at this point
# because session can be resumed
return False
else:
return True
def check_acks(self, client):
"""Request ack if needed"""
session = client._xep_0198_session
# log.debug("check_acks (in_counter={}, out_counter={}, buf len={}, buf idx={})"
# .format(session.in_counter, session.out_counter, len(session.buffer),
# session.buffer_idx))
if session.ack_requested or not session.buffer:
return
if (
session.out_counter - session.buffer_idx >= MAX_STANZA_ACK_R
or time.time() - session.last_ack_r >= MAX_DELAY_ACK_R
):
self.request_ack(client)
session.ack_requested = True
session.last_ack_r = time.time()
def update_buffer(self, session, server_acked):
"""Update buffer and buffer_index"""
if server_acked > session.buffer_idx:
diff = server_acked - session.buffer_idx
try:
for i in range(diff):
session.buffer.pop()
except IndexError:
log.error(
"error while cleaning buffer, invalid index (buffer is empty):\n"
"diff = {diff}\n"
"server_acked = {server_acked}\n"
"buffer_idx = {buffer_id}".format(
diff=diff, server_acked=server_acked, buffer_id=session.buffer_idx
)
)
session.buffer_idx += diff
def replay_buffer(self, client, buffer_, discard_results=False):
"""Resend all stanza in buffer
@param buffer_(collection.deque, list): buffer to replay
the buffer will be cleared by this method
@param discard_results(bool): if True, don't replay IQ result stanzas
"""
while True:
try:
stanza = buffer_.pop()
except IndexError:
break
else:
if (
discard_results
and stanza.name == "iq"
and stanza.getAttribute("type") == "result"
):
continue
client.send(stanza)
def send_ack(self, client):
"""Send an answer element with current IN counter"""
a_elt = domish.Element((NS_SM, "a"))
a_elt["h"] = str(client._xep_0198_session.in_counter)
client.send(a_elt)
def request_ack(self, client):
"""Send a request element"""
session = client._xep_0198_session
r_elt = domish.Element((NS_SM, "r"))
client.send(r_elt)
if session.req_timer is not None:
raise exceptions.InternalError("req_timer should not be set")
if self._ack_timeout:
session.req_timer = reactor.callLater(
self._ack_timeout, self.on_ack_time_out, client
)
def _connectionFailed(self, failure_, connector):
normal_host, normal_port = connector.normal_location
del connector.normal_location
log.warning(
_(
"Connection failed using location given by server (host: {host}, port: "
"{port}), switching to normal host and port (host: {normal_host}, port: "
"{normal_port})".format(
host=connector.host,
port=connector.port,
normal_host=normal_host,
normal_port=normal_port,
)
)
)
connector.host, connector.port = normal_host, normal_port
connector.connectionFailed = connector.connectionFailed_ori
del connector.connectionFailed_ori
return connector.connectionFailed(failure_)
def on_enabled(self, enabled_elt, client):
session = client._xep_0198_session
session.in_counter = 0
# we check that resuming is possible and that we have a session id
resume = C.bool(enabled_elt.getAttribute("resume"))
session_id = enabled_elt.getAttribute("id")
if not session_id:
log.warning(_('Incorrect element received, no "id" attribute'))
if not resume or not session_id:
log.warning(
_(
"You're server doesn't support session resuming with stream management, "
"please contact your server administrator to enable it"
)
)
return
session.session_id = session_id
# XXX: we disable resource binding, which must not be done
# when we resume the session.
client.factory.authenticator.res_binding = False
# location, in case server want resuming session to be elsewhere
try:
location = enabled_elt["location"]
except KeyError:
pass
else:
# TODO: handle IPv6 here (in brackets, cf. XEP)
try:
domain, port = location.split(":", 1)
port = int(port)
except ValueError:
log.warning(
_("Invalid location received: {location}").format(location=location)
)
else:
session.location = (domain, port)
# we monkey patch connector to use the new location
connector = client.xmlstream.transport.connector
connector.normal_location = connector.host, connector.port
connector.host = domain
connector.port = port
connector.connectionFailed_ori = connector.connectionFailed
connector.connectionFailed = partial(
self._connectionFailed, connector=connector
)
# resuming time
try:
max_s = int(enabled_elt["max"])
except (ValueError, KeyError) as e:
if isinstance(e, ValueError):
log.warning(_('Invalid "max" attribute'))
max_s = RESUME_MAX
log.info(
_("Using default session max value ({max_s} s).".format(max_s=max_s))
)
log.info(_("Stream Management enabled"))
else:
log.info(
_(
"Stream Management enabled, with a resumption time of {res_m:.2f} min".format(
res_m=max_s / 60
)
)
)
session.session_max = max_s
def on_resumed(self, enabled_elt, client):
session = client._xep_0198_session
assert not session.enabled
del session.resuming
server_acked = int(enabled_elt["h"])
self.update_buffer(session, server_acked)
resend_count = len(session.buffer)
# we resend all stanza which have not been received properly
self.replay_buffer(client, session.buffer)
# now we can continue the session
session.enabled = True
d_time = time.time() - session.disconnected_time
log.info(
_(
"Stream session resumed (disconnected for {d_time} s, {count} "
"stanza(s) resent)"
).format(d_time=int(d_time), count=resend_count)
)
def on_failed(self, failed_elt, client):
session = client._xep_0198_session
condition_elt = failed_elt.firstChildElement()
buffer_ = session.get_buffer_copy()
session.reset()
try:
del session.resuming
except AttributeError:
# stream management can't be started at all
msg = _("Can't use stream management")
if condition_elt is None:
log.error(msg + ".")
else:
log.error(_("{msg}: {reason}").format(msg=msg, reason=condition_elt.name))
else:
# only stream resumption failed, we can try full session init
# XXX: we try to start full session init from this point, with many
# variables/attributes already initialised with a potentially different
# jid. This is experimental and may not be safe. It may be more
# secured to abord the connection and restart everything with a fresh
# client.
msg = _("stream resumption not possible, restarting full session")
if condition_elt is None:
log.warning("{msg}.".format(msg=msg))
else:
log.warning("{msg}: {reason}".format(msg=msg, reason=condition_elt.name))
# stream resumption failed, but we still can do normal stream management
# we restore attributes as if the session was new, and init stream
# we keep everything initialized, and only do binding, roster request
# and initial presence sending.
if client.conn_deferred.called:
client.conn_deferred = defer.Deferred()
else:
log.error("conn_deferred should be called at this point")
plg_0045 = self.host.plugins.get("XEP-0045")
plg_0313 = self.host.plugins.get("XEP-0313")
# FIXME: we should call all loaded plugins with generic callbacks
# (e.g. prepareResume and resume), so a hot resuming can be done
# properly for all plugins.
if plg_0045 is not None:
# we have to remove joined rooms
muc_join_args = plg_0045.pop_rooms(client)
# we need to recreate roster
client.handlers.remove(client.roster)
client.roster = client.roster.__class__(self.host)
client.roster.setHandlerParent(client)
# bind init is not done when resuming is possible, so we have to do it now
bind_init = jabber_client.BindInitializer(client.xmlstream)
bind_init.required = True
d = bind_init.start()
# we set the jid, which may have changed
d.addCallback(
lambda __: setattr(client.factory.authenticator, "jid", client.jid)
)
# we call the trigger who will send the element
d.addCallback(lambda __: self._xml_init_trigger(client))
# then we have to re-request the roster, as changes may have occured
d.addCallback(lambda __: defer.ensureDeferred(client.roster.request_roster()))
# we add got_roster to be sure to have roster before sending initial presence
d.addCallback(lambda __: client.roster.got_roster)
# initial presence must be sent manually
d.addCallback(lambda __: client.presence.available())
if plg_0313 is not None:
# we retrieve one2one MAM archives (fire and forget, don't block)
defer.ensureDeferred(plg_0313.resume(client))
if plg_0045 is not None:
# we re-join MUC rooms
muc_d_list = defer.DeferredList(
[defer.ensureDeferred(plg_0045.join(*args)) for args in muc_join_args]
)
d.addCallback(lambda __: muc_d_list)
# at the end we replay the buffer, as those stanzas have probably not
# been received
d.addCallback(
lambda __: self.replay_buffer(client, buffer_, discard_results=True)
)
def on_receive(self, element, client):
if not client.is_component:
session = client._xep_0198_session
if session.enabled and element.name.lower() in C.STANZA_NAMES:
session.in_counter += 1 % MAX_COUNTER
def on_send(self, obj, client):
if not client.is_component:
session = client._xep_0198_session
if (
session.enabled
and domish.IElement.providedBy(obj)
and obj.name.lower() in C.STANZA_NAMES
):
session.out_counter += 1 % MAX_COUNTER
session.buffer.appendleft(obj)
self.check_acks(client)
def on_ack_request(self, r_elt, client):
self.send_ack(client)
def on_ack_answer(self, a_elt, client):
session = client._xep_0198_session
session.ack_requested = False
if self._ack_timeout:
if session.req_timer is None:
log.error("req_timer should be set")
else:
session.req_timer.cancel()
session.req_timer = None
try:
server_acked = int(a_elt["h"])
except ValueError:
log.warning(
_(
"Server returned invalid ack element, disabling stream "
"management: {xml}"
).format(xml=a_elt)
)
session.enabled = False
return
if server_acked > session.out_counter:
log.error(
_(
"Server acked more stanzas than we have sent, disabling stream "
"management."
)
)
session.reset()
return
self.update_buffer(session, server_acked)
self.check_acks(client)
def on_ack_time_out(self, client):
"""Called when a requested ACK has not been received in time"""
log.info(_("Ack was not received in time, aborting connection"))
try:
xmlstream = client.xmlstream
except AttributeError:
log.warning("xmlstream has already been terminated")
else:
transport = xmlstream.transport
if transport is None:
log.warning("transport was already removed")
else:
transport.abortConnection()
client._xep_0198_session.req_timer = None
@implementer(itx_xmpp.IDisco)
class XEP_0198_handler(xmlstream.XMPPHandler):
def __init__(self, plugin_parent):
self.plugin_parent = plugin_parent
self.host = plugin_parent.host
def connectionInitialized(self):
self.xmlstream.addObserver(
SM_ENABLED, self.plugin_parent.on_enabled, client=self.parent
)
self.xmlstream.addObserver(
SM_RESUMED, self.plugin_parent.on_resumed, client=self.parent
)
self.xmlstream.addObserver(
SM_FAILED, self.plugin_parent.on_failed, client=self.parent
)
self.xmlstream.addObserver(
SM_R_REQUEST, self.plugin_parent.on_ack_request, client=self.parent
)
self.xmlstream.addObserver(
SM_A_REQUEST, self.plugin_parent.on_ack_answer, client=self.parent
)
def getDiscoInfo(self, requestor, target, nodeIdentifier=""):
return [disco.DiscoFeature(NS_SM)]
def getDiscoItems(self, requestor, target, nodeIdentifier=""):
return []