571 lines
17 KiB
Python
571 lines
17 KiB
Python
########################################################################
|
|
# File name: service.py
|
|
# This file is part of: aioxmpp
|
|
#
|
|
# LICENSE
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Lesser 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
|
|
# Lesser General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Lesser General Public
|
|
# License along with this program. If not, see
|
|
# <http://www.gnu.org/licenses/>.
|
|
#
|
|
########################################################################
|
|
import asyncio
|
|
import random
|
|
|
|
from datetime import timedelta
|
|
|
|
import aioxmpp
|
|
import aioxmpp.callbacks
|
|
import aioxmpp.errors as errors
|
|
import aioxmpp.service as service
|
|
import aioxmpp.utils as utils
|
|
|
|
from . import xso as ibb_xso
|
|
|
|
|
|
MAX_BLOCK_SIZE = (1 << 16) - 1
|
|
|
|
|
|
class IBBTransport(asyncio.Transport):
|
|
"""
|
|
The transport for IBB sessions.
|
|
|
|
.. note:: Never instantiate this class directly, all instances of
|
|
this class are created the methods
|
|
:meth:`~aioxmpp.ibb.IBBService.open_session` and
|
|
:meth:`~aioxmpp.ibb.IBBService.expect_session` of
|
|
:class:`~aioxmpp.ibb.IBBService`.
|
|
|
|
The following keys are supported for
|
|
:meth:`~asyncio.BaseTransport.get_extra_info`:
|
|
|
|
`block_size`
|
|
The maximal block size of data in a IBB stanza.
|
|
|
|
`peer_jid`
|
|
The JID of the peer.
|
|
|
|
`sid`
|
|
The session id of the unerlying IBB session.
|
|
|
|
`stanza_type`
|
|
The used stanza type.
|
|
"""
|
|
|
|
def __init__(self, service, peer_jid, sid, stanza_type, block_size):
|
|
self._protocol = None
|
|
self._service = service
|
|
self._stanza_type = stanza_type
|
|
self._sid = sid
|
|
self._peer_jid = peer_jid
|
|
self._block_size = block_size
|
|
self._incoming_seq = 0
|
|
self._outgoing_seq = 0
|
|
self._closed = False
|
|
self._closing = False
|
|
|
|
self.set_write_buffer_limits()
|
|
self._write_buffer = b""
|
|
self._can_write = asyncio.Event()
|
|
|
|
self._reading_paused = False
|
|
self._input_buffer = []
|
|
|
|
self._write_task = asyncio.ensure_future(self._write_task_main())
|
|
self._write_task.add_done_callback(
|
|
self._handle_close
|
|
)
|
|
|
|
self._wait_time = self._service.initial_wait_time.total_seconds()
|
|
self._retries = 0
|
|
|
|
def set_write_buffer_limits(self, high=None, low=None):
|
|
if low is None:
|
|
low = 4 * self._block_size
|
|
|
|
if high is None:
|
|
high = 8 * self._block_size
|
|
|
|
if low < 0:
|
|
raise ValueError("the limits must be positive")
|
|
|
|
if high < 0:
|
|
raise ValueError("the limits must be positive")
|
|
|
|
if low > high:
|
|
low = high
|
|
|
|
self._output_buffer_limit_low = low
|
|
self._output_buffer_limit_high = high
|
|
|
|
def get_write_buffer_limits(self):
|
|
return self._output_buffer_limit_low, self._output_buffer_limit_high
|
|
|
|
def get_write_buffer_size(self):
|
|
return len(self._write_buffer)
|
|
|
|
def set_protocol(self, proto):
|
|
self._protocol = proto
|
|
proto.connection_made(self)
|
|
|
|
def get_protocol(self):
|
|
return self._protocol
|
|
|
|
def pause_reading(self):
|
|
self._reading_paused = True
|
|
|
|
def resume_reading(self):
|
|
self._reading_paused = False
|
|
self._protocol.data_received(b"".join(self._input_buffer))
|
|
self._input_buffer.clear()
|
|
|
|
def is_closing(self):
|
|
return self._closing or self._closed
|
|
|
|
def get_extra_info(self, key, default=None):
|
|
return {
|
|
"block_size": self._block_size,
|
|
"peer_jid": self._peer_jid,
|
|
"stanza_type": self._stanza_type,
|
|
"sid": self._sid,
|
|
}.get(key, default)
|
|
|
|
async def _write_task_main(self):
|
|
e = None
|
|
while True:
|
|
await self._can_write.wait()
|
|
|
|
if self._write_buffer:
|
|
data = self._write_buffer[:self._block_size]
|
|
|
|
if self._stanza_type == ibb_xso.IBBStanzaType.IQ:
|
|
stanza = aioxmpp.IQ(
|
|
aioxmpp.IQType.SET,
|
|
to=self._peer_jid,
|
|
payload=ibb_xso.Data(
|
|
self._sid,
|
|
self._outgoing_seq,
|
|
data
|
|
)
|
|
)
|
|
elif self._stanza_type == ibb_xso.IBBStanzaType.MESSAGE:
|
|
# TODO: use some form of tracking for messages
|
|
stanza = aioxmpp.Message(
|
|
aioxmpp.MessageType.NORMAL,
|
|
to=self._peer_jid,
|
|
)
|
|
stanza.xep0047_data = ibb_xso.Data(
|
|
self._sid,
|
|
self._outgoing_seq,
|
|
data
|
|
)
|
|
|
|
try:
|
|
await self._service.client.send(
|
|
stanza
|
|
)
|
|
except errors.XMPPWaitError:
|
|
# wait and try again unless max retries have been reached
|
|
if self._retries < self._service.max_retries:
|
|
await asyncio.sleep(self._wait_time)
|
|
self._wait_time *= self._service.wait_backoff_factor
|
|
self._retries += 1
|
|
continue
|
|
else:
|
|
e = asyncio.TimeoutError()
|
|
break
|
|
except errors.StanzaError as _e:
|
|
# break the loop to close the connection
|
|
e = _e
|
|
break
|
|
# update the internal state after the successful
|
|
# write: remove the written data from the buffer and
|
|
# increment the sequence number
|
|
self._write_buffer = self._write_buffer[len(data):]
|
|
self._outgoing_seq += 1
|
|
self._outgoing_seq &= 0xffff
|
|
|
|
# reset the wait time
|
|
self._wait_time = \
|
|
self._service.initial_wait_time.total_seconds()
|
|
self._retries = 0
|
|
|
|
if len(self._write_buffer) < self._output_buffer_limit_low:
|
|
self._protocol.resume_writing()
|
|
|
|
if not self._write_buffer:
|
|
if self._closing:
|
|
e = None
|
|
break
|
|
self._can_write.clear()
|
|
|
|
close = ibb_xso.Close()
|
|
close.sid = self._sid
|
|
stanza = aioxmpp.IQ(
|
|
aioxmpp.IQType.SET,
|
|
to=self._peer_jid,
|
|
payload=close,
|
|
)
|
|
|
|
try:
|
|
await self._service.client.send(stanza)
|
|
except errors.StanzaError as _e:
|
|
if e is None:
|
|
e = _e
|
|
finally:
|
|
if e is not None:
|
|
raise e
|
|
|
|
def write(self, data):
|
|
"""
|
|
Send `data` over the IBB. If `data` is larger than the block size
|
|
is is chunked and sent in chunks.
|
|
|
|
Chunks from one call of :meth:`write` will always be sent in
|
|
series.
|
|
"""
|
|
|
|
if self.is_closing():
|
|
return
|
|
|
|
self._write_buffer += data
|
|
|
|
if len(self._write_buffer) >= self._output_buffer_limit_high:
|
|
self._protocol.pause_writing()
|
|
|
|
if self._write_buffer:
|
|
self._can_write.set()
|
|
|
|
def _connection_closed(self):
|
|
self._write_task.cancel()
|
|
|
|
def _handle_close(self, fut):
|
|
e = None
|
|
self._service._remove_session(self._peer_jid, self._sid)
|
|
try:
|
|
e = fut.exception()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._protocol.connection_lost(e)
|
|
self._closed = True
|
|
|
|
def close(self):
|
|
"""
|
|
Close the session.
|
|
"""
|
|
if self.is_closing():
|
|
return
|
|
|
|
self._closing = True
|
|
# make sure the writer wakes up
|
|
self._can_write.set()
|
|
|
|
def abort(self):
|
|
"""
|
|
Abort the session.
|
|
"""
|
|
if self.is_closing():
|
|
return
|
|
self._connection_closed()
|
|
|
|
def _data_received(self, data):
|
|
if self._closed:
|
|
return
|
|
|
|
if self._reading_paused:
|
|
self._input_buffer.append(data)
|
|
else:
|
|
self._protocol.data_received(data)
|
|
|
|
def _process_iq(self, payload):
|
|
if payload.seq != self._incoming_seq:
|
|
raise errors.XMPPCancelError(
|
|
condition=errors.ErrorCondition.UNEXPECTED_REQUEST
|
|
)
|
|
self._incoming_seq += 1
|
|
self._incoming_seq &= 0xffff
|
|
self._data_received(payload.content)
|
|
|
|
def _process_msg(self, payload):
|
|
if payload.seq != self._incoming_seq:
|
|
return
|
|
self._incoming_seq += 1
|
|
self._incoming_seq &= 0xffff
|
|
self._data_received(payload.content)
|
|
|
|
|
|
class IBBService(service.Service):
|
|
"""
|
|
A service implementing in-band bytestreams.
|
|
|
|
Methods for establishing sessions:
|
|
|
|
.. automethod:: expect_session
|
|
|
|
.. automethod:: open_session
|
|
|
|
The following attributes control the establishment of sessions due
|
|
to a received request, that was not announced to the service by
|
|
:meth:`expect_session`:
|
|
|
|
.. attribute:: session_limit
|
|
:annotation: = 0
|
|
|
|
The maximal number of sessions to be accepted. If there are
|
|
that many or more active sessions, no new sessions are
|
|
accepted, unless they are whitelisted by
|
|
:meth:`expect_session`. (This means, that by default only
|
|
expected sessions are accepted!).
|
|
|
|
.. attribute:: default_protocol_factory
|
|
|
|
The protocol factory to be used when an unexpected connection
|
|
is established. This *must* be set when changing
|
|
:attr:`session_limit` to a non-zero value.
|
|
|
|
.. signal:: on_session_accepted(transport, protocol)
|
|
|
|
Fires when a session is established due to a received open
|
|
request that was not expected (compare
|
|
:meth:`expect_session`). This can only happen when
|
|
:attr:`session_limit` is set to another value than its default
|
|
value.
|
|
|
|
The following attributes control how the IBB sessions react to
|
|
errors of type wait:
|
|
|
|
.. attribute:: max_retries
|
|
:annotation: = 5
|
|
|
|
The number of times it is tried to resend a data stanza, when a
|
|
:class:`~aioxmpp.errors.XMPPWaitError` is received. When
|
|
:attr:`max_retries` have been tried, the session is closed.
|
|
`connection_lost` of the protocol receives an
|
|
:class:`asyncio.TimeoutError`.
|
|
|
|
.. attribute:: initial_wait_time
|
|
:annotation: = timedelta(seconds=1)
|
|
|
|
The time to wait when receiving a
|
|
:class:`~aioxmpp.errors.XMPPWaitError` for the first time.
|
|
|
|
.. attribute:: wait_backoff_factor
|
|
:annotation: = 1.2
|
|
|
|
The factor by which the wait time is prolonged on each
|
|
successive wait error.
|
|
"""
|
|
|
|
on_session_accepted = aioxmpp.callbacks.Signal()
|
|
|
|
def __init__(self, client, **kwargs):
|
|
super().__init__(client, **kwargs)
|
|
|
|
self._sessions = {}
|
|
|
|
self.session_limit = 0
|
|
self._expected_sessions = {}
|
|
|
|
self.default_protocol_factory = None
|
|
|
|
self.client.on_stream_destroyed.connect(
|
|
self._on_stream_destroyed
|
|
)
|
|
|
|
self.max_retries = 5
|
|
self.initial_wait_time = timedelta(seconds=1)
|
|
self.wait_backoff_factor = 1.2
|
|
|
|
def _on_stream_destroyed(self):
|
|
self._expected_sessions = {}
|
|
|
|
# tear down the remaining open sessions
|
|
for session in list(self._sessions.values()):
|
|
session.abort()
|
|
|
|
def expect_session(self, protocol_factory, peer_jid, sid):
|
|
"""
|
|
Whitelist the session with `peer_jid` and the session id `sid` and
|
|
return it when it is established. This is meant to be used
|
|
with signalling protocols like Jingle and is the counterpart
|
|
to :meth:`open_session`.
|
|
|
|
:returns: an awaitable object, whose result is the tuple
|
|
`(transport, protocol)`
|
|
"""
|
|
def on_done(fut):
|
|
del self._expected_sessions[sid, peer_jid]
|
|
|
|
_, fut = self._expected_sessions[sid, peer_jid] = (
|
|
protocol_factory, asyncio.Future()
|
|
)
|
|
fut.add_done_callback(on_done)
|
|
return fut
|
|
|
|
async def open_session(self, protocol_factory, peer_jid, *,
|
|
stanza_type=ibb_xso.IBBStanzaType.IQ,
|
|
block_size=4096, sid=None):
|
|
"""
|
|
Establish an in-band bytestream session with `peer_jid` and
|
|
return the transport and protocol.
|
|
|
|
:param protocol_factory: the protocol factory
|
|
:type protocol_factory: a nullary callable returning an
|
|
:class:`asyncio.Protocol` instance
|
|
:param peer_jid: the JID with which to establish the byte-stream.
|
|
:type peer_jid: :class:`aioxmpp.JID`
|
|
:param stanza_type: the stanza type to use
|
|
:type stanza_type: class:`~aioxmpp.ibb.IBBStanzaType`
|
|
:param block_size: the maximal size of blocks to transfer
|
|
:type block_size: :class:`int`
|
|
:param sid: the session id to use
|
|
:type sid: :class:`str` (must be a valid NMTOKEN)
|
|
|
|
:returns: the transport and protocol
|
|
:rtype: a tuple of :class:`aioxmpp.ibb.service.IBBTransport`
|
|
and :class:`asyncio.Protocol`
|
|
"""
|
|
if block_size > MAX_BLOCK_SIZE:
|
|
raise ValueError("block_size too large")
|
|
|
|
if sid is None:
|
|
sid = utils.to_nmtoken(random.getrandbits(8*8))
|
|
open_ = ibb_xso.Open()
|
|
open_.stanza = stanza_type
|
|
open_.sid = sid
|
|
open_.block_size = block_size
|
|
|
|
# XXX: retry on XMPPModifyError with RESOURCE_CONSTRAINT
|
|
await self.client.send(
|
|
aioxmpp.IQ(
|
|
aioxmpp.IQType.SET,
|
|
to=peer_jid,
|
|
payload=open_,
|
|
)
|
|
)
|
|
|
|
handle = self._sessions[sid, peer_jid] = IBBTransport(
|
|
self,
|
|
peer_jid,
|
|
sid,
|
|
stanza_type,
|
|
block_size,
|
|
)
|
|
|
|
protocol = protocol_factory()
|
|
handle.set_protocol(protocol)
|
|
return handle, protocol
|
|
|
|
@service.iq_handler(
|
|
aioxmpp.IQType.SET,
|
|
ibb_xso.Open)
|
|
async def _handle_open_request(self, iq):
|
|
peer_jid = iq.from_
|
|
sid = iq.payload.sid
|
|
block_size = iq.payload.block_size
|
|
stanza_type = iq.payload.stanza
|
|
|
|
if block_size > MAX_BLOCK_SIZE:
|
|
raise errors.XMPPModifyError(
|
|
condition=errors.ErrorCondition.RESOURCE_CONSTRAINT
|
|
)
|
|
|
|
try:
|
|
protocol_factory, expected_future = \
|
|
self._expected_sessions[sid, peer_jid]
|
|
except KeyError:
|
|
if len(self._sessions) >= self.session_limit:
|
|
raise errors.XMPPCancelError(
|
|
condition=errors.ErrorCondition.NOT_ACCEPTABLE
|
|
)
|
|
expected_future = None
|
|
protocol_factory = self.default_protocol_factory
|
|
|
|
if (sid, peer_jid) in self._sessions:
|
|
# disallow opening a session twice
|
|
if expected_future is not None:
|
|
# is this correct?
|
|
expected_future.cancel()
|
|
raise errors.XMPPCancelError(
|
|
condition=errors.ErrorCondition.NOT_ACCEPTABLE
|
|
)
|
|
|
|
handle = self._sessions[sid, peer_jid] = IBBTransport(
|
|
self,
|
|
peer_jid,
|
|
sid,
|
|
stanza_type,
|
|
block_size
|
|
)
|
|
|
|
protocol = protocol_factory()
|
|
handle.set_protocol(protocol)
|
|
|
|
if expected_future is None:
|
|
self.on_session_accepted((handle, protocol))
|
|
else:
|
|
expected_future.set_result((handle, protocol))
|
|
|
|
@service.iq_handler(
|
|
aioxmpp.IQType.SET,
|
|
ibb_xso.Close)
|
|
async def _handle_close_request(self, iq):
|
|
peer_jid = iq.from_
|
|
sid = iq.payload.sid
|
|
|
|
try:
|
|
session_handle = self._sessions[sid, peer_jid]
|
|
except KeyError:
|
|
raise errors.XMPPCancelError(
|
|
condition=errors.ErrorCondition.ITEM_NOT_FOUND
|
|
)
|
|
|
|
session_handle._connection_closed()
|
|
|
|
@service.iq_handler(
|
|
aioxmpp.IQType.SET,
|
|
ibb_xso.Data)
|
|
async def _handle_data(self, iq):
|
|
peer_jid = iq.from_
|
|
sid = iq.payload.sid
|
|
|
|
try:
|
|
session_handle = self._sessions[sid, peer_jid]
|
|
except KeyError:
|
|
raise errors.XMPPCancelError(
|
|
condition=errors.ErrorCondition.ITEM_NOT_FOUND
|
|
)
|
|
|
|
session_handle._process_iq(iq.payload)
|
|
|
|
@service.inbound_message_filter
|
|
def _handle_message(self, msg):
|
|
if msg.xep0047_data is None:
|
|
return msg
|
|
|
|
payload = msg.xep0047_data
|
|
peer_jid = msg.from_
|
|
sid = payload.sid
|
|
|
|
try:
|
|
session = self._sessions[sid, peer_jid]
|
|
except KeyError:
|
|
return None
|
|
|
|
session._process_msg(payload)
|
|
|
|
return None
|
|
|
|
def _remove_session(self, peer_jid, sid):
|
|
del self._sessions[sid, peer_jid]
|