diff options
Diffstat (limited to 'slixmpp/plugins/xep_0047/stream.py')
| -rw-r--r-- | slixmpp/plugins/xep_0047/stream.py | 108 |
1 files changed, 45 insertions, 63 deletions
diff --git a/slixmpp/plugins/xep_0047/stream.py b/slixmpp/plugins/xep_0047/stream.py index 817f96a1..3be894eb 100644 --- a/slixmpp/plugins/xep_0047/stream.py +++ b/slixmpp/plugins/xep_0047/stream.py @@ -1,7 +1,6 @@ +import asyncio import socket -import threading import logging -from queue import Queue from slixmpp.stanza import Iq from slixmpp.exceptions import XMPPError @@ -12,11 +11,10 @@ log = logging.getLogger(__name__) class IBBytestream(object): - def __init__(self, xmpp, sid, block_size, jid, peer, window_size=1, use_messages=False): + def __init__(self, xmpp, sid, block_size, jid, peer, use_messages=False): self.xmpp = xmpp self.sid = sid self.block_size = block_size - self.window_size = window_size self.use_messages = use_messages if jid is None: @@ -27,29 +25,20 @@ class IBBytestream(object): self.send_seq = -1 self.recv_seq = -1 - self._send_seq_lock = threading.Lock() - self._recv_seq_lock = threading.Lock() + self.stream_started = False + self.stream_in_closed = False + self.stream_out_closed = False - self.stream_started = threading.Event() - self.stream_in_closed = threading.Event() - self.stream_out_closed = threading.Event() + self.recv_queue = asyncio.Queue() - self.recv_queue = Queue() - - self.send_window = threading.BoundedSemaphore(value=self.window_size) - self.window_ids = set() - self.window_empty = threading.Event() - self.window_empty.set() - - def send(self, data): - if not self.stream_started.is_set() or \ - self.stream_out_closed.is_set(): + @asyncio.coroutine + def send(self, data, timeout=None): + if not self.stream_started or self.stream_out_closed: raise socket.error - data = data[0:self.block_size] - self.send_window.acquire() - with self._send_seq_lock: - self.send_seq = (self.send_seq + 1) % 65535 - seq = self.send_seq + if len(data) > self.block_size: + data = data[:self.block_size] + self.send_seq = (self.send_seq + 1) % 65535 + seq = self.send_seq if self.use_messages: msg = self.xmpp.Message() msg['to'] = self.peer_jid @@ -59,7 +48,6 @@ class IBBytestream(object): msg['ibb_data']['seq'] = seq msg['ibb_data']['data'] = data msg.send() - self.send_window.release() else: iq = self.xmpp.Iq() iq['type'] = 'set' @@ -68,71 +56,65 @@ class IBBytestream(object): iq['ibb_data']['sid'] = self.sid iq['ibb_data']['seq'] = seq iq['ibb_data']['data'] = data - self.window_empty.clear() - self.window_ids.add(iq['id']) - iq.send(callback=self._recv_ack) + yield from iq.send(timeout=timeout) return len(data) - def sendall(self, data): + @asyncio.coroutine + def sendall(self, data, timeout=None): sent_len = 0 while sent_len < len(data): - sent_len += self.send(data[sent_len:]) - - def _recv_ack(self, iq): - self.window_ids.remove(iq['id']) - if not self.window_ids: - self.window_empty.set() - self.send_window.release() - if iq['type'] == 'error': - self.close() + sent_len += yield from self.send(data[sent_len:self.block_size], timeout=timeout) + + @asyncio.coroutine + def sendfile(self, file, timeout=None): + while True: + data = file.read(self.block_size) + if not data: + break + yield from self.send(data, timeout=timeout) def _recv_data(self, stanza): - with self._recv_seq_lock: - new_seq = stanza['ibb_data']['seq'] - if new_seq != (self.recv_seq + 1) % 65535: - self.close() - raise XMPPError('unexpected-request') - self.recv_seq = new_seq + new_seq = stanza['ibb_data']['seq'] + if new_seq != (self.recv_seq + 1) % 65535: + self.close() + raise XMPPError('unexpected-request') + self.recv_seq = new_seq data = stanza['ibb_data']['data'] if len(data) > self.block_size: self.close() raise XMPPError('not-acceptable') - self.recv_queue.put(data) - self.xmpp.event('ibb_stream_data', {'stream': self, 'data': data}) + self.recv_queue.put_nowait(data) + self.xmpp.event('ibb_stream_data', self) if isinstance(stanza, Iq): stanza.reply().send() def recv(self, *args, **kwargs): - return self.read(block=True) + return self.read() - def read(self, block=True, timeout=None, **kwargs): - if not self.stream_started.is_set() or \ - self.stream_in_closed.is_set(): + def read(self): + if not self.stream_started or self.stream_in_closed: raise socket.error - if timeout is not None: - block = True - try: - return self.recv_queue.get(block, timeout) - except: - return None - - def close(self): + return self.recv_queue.get_nowait() + + def close(self, timeout=None): iq = self.xmpp.Iq() iq['type'] = 'set' iq['to'] = self.peer_jid iq['from'] = self.self_jid iq['ibb_close']['sid'] = self.sid - self.stream_out_closed.set() - iq.send(block=False, - callback=lambda x: self.stream_in_closed.set()) + self.stream_out_closed = True + def _close_stream(_): + self.stream_in_closed = True + future = iq.send(timeout=timeout, callback=_close_stream) self.xmpp.event('ibb_stream_end', self) + return future def _closed(self, iq): - self.stream_in_closed.set() - self.stream_out_closed.set() + self.stream_in_closed = True + self.stream_out_closed = True iq.reply().send() self.xmpp.event('ibb_stream_end', self) |