Exemplo n.º 1
0
class VMessageDecrypter(object):
    """Decrypter for encrypted messages with data integrity check.

    Decodes encrypted plaintext messages in the format encrypted by
    :class:`VMessageEncrypter`\ .

    :param decrypter:   transform for decryption
    :type  decrypter:   :class:`versile.crypto.VBlockTransform`
    :param hash_cls:    hash class for message integrity hash
    :type  hash_cls:    :class:`versile.crypto.VHash`
    :param mac_secret:  secret data for package authentication
    :type  mac_secret:  bytes

    """
    def __init__(self, decrypter, hash_cls, mac_secret):
        self._decrypter = decrypter
        self._hash_cls = hash_cls
        self._mac_secret = mac_secret
        self._max_plaintext_len = 0x10000  # HARDCODED 2-byte message length
        self._cipher_blocksize = decrypter.blocksize
        self._hash_len = hash_cls.digest_size()
        self._read_buf = VByteBuffer()
        self._in_buf = VByteBuffer()
        self._msg_buf = VByteBuffer()
        self._have_len = False
        self._plaintext_blocksize = None
        self._plaintext_len = None
        self._invalid = False
        self._result = None
        self._msg_num = 0

    def reset(self):
        """Resets the decrypter to read a new message.

        :raises: :class:`VCryptoException`

        Raises an exception if ongoing decryption is not completed.

        """
        if self._result is None:
            raise VCryptoException('Ongoing decryption not completed')
        self._read_buf.remove()
        self._in_buf.remove()
        self._msg_buf.remove()
        self._have_len = False
        self._plaintext_len = None
        self._invalid = False
        self._result = None

    def read(self, data):
        """Reads encrypted message data.

        :param data: input data to decrypt and decode
        :type  data: bytes, :class:`versile.common.util.VByteBuffer`
        :returns:    number of bytes read
        :rtype:      int

        Reads only as much data as is required to complete processing
        a complete single message. If data is of type
        :class:`versile.common.util.VByteBuffer` then the data that
        was read will be popped off the buffer.

        .. note::

            When decryption of one message has completed,
            :meth:`reset` must be called before a new message can be
            read.

        """
        if isinstance(data, bytes):
            read_buf = self._read_buf
            read_buf.remove()
            read_buf.append(data)
        elif isinstance(data, VByteBuffer):
            read_buf = data
        else:
            raise TypeError('Input must be bytes or VByteBuffer')

        num_read = 0
        _pbsize = self._plaintext_blocksize
        _cbsize = self._cipher_blocksize
        while read_buf and self._result is None and not self._invalid:
            # First decode single block to get blocksize
            if not self._have_len:
                max_read = _cbsize - len(self._in_buf)
                enc_data = read_buf.pop(max_read)
                self._in_buf.append(enc_data)
                num_read += len(enc_data)
                if len(self._in_buf) == _cbsize:
                    enc_data = self._in_buf.pop()
                    block = self._decrypter(enc_data)
                    if _pbsize is None:
                        _pbsize = len(block)
                        self._plaintext_blocksize = _pbsize
                    self._msg_buf.append(block)
                    len_bytes = self._msg_buf.peek(2)
                    if _pyver == 2:
                        self._plaintext_len = 1 + (
                            (_b_ord(len_bytes[0]) << 8) + _b_ord(len_bytes[1]))
                    else:
                        self._plaintext_len = 1 + (
                            (len_bytes[0] << 8) + len_bytes[1])
                    self._have_len = True

            # If we have first block, decrypt more blocks as available/needed
            if self._have_len:
                msg_len = 2 + self._plaintext_len + self._hash_len
                pad_len = msg_len % _pbsize
                if pad_len:
                    pad_len = _pbsize - pad_len
                    msg_len += pad_len
                msg_left = msg_len - len(self._msg_buf)
                blocks_left = msg_left // _pbsize
                input_left = (blocks_left * _cbsize - len(self._in_buf))
                in_data = read_buf.pop(input_left)
                num_read += len(in_data)
                self._in_buf.append(in_data)
                num_decode = len(self._in_buf)
                num_decode -= num_decode % _cbsize
                if num_decode > 0:
                    enc_data = self._in_buf.pop(num_decode)
                    self._msg_buf.append(self._decrypter(enc_data))
                elif len(self._msg_buf) != msg_len:
                    break

            if self._have_len and len(self._msg_buf) == msg_len:
                len_bytes = self._msg_buf.pop(2)
                plaintext = self._msg_buf.pop(self._plaintext_len)
                padding = self._msg_buf.pop(pad_len)
                msg_hash = self._msg_buf.pop(self._hash_len)

                _mac_msg = b''.join((posint_to_bytes(self._msg_num), len_bytes,
                                     plaintext, padding))
                if msg_hash == self._hash_cls.hmac(self._mac_secret, _mac_msg):
                    self._result = plaintext
                    self._msg_num += 1
                else:
                    self._invalid = True

        return num_read

    def done(self):
        """Returns True if decryption and decoding of a message was done.

        :returns: True if reading is done
        :rtype:   bool
        :raises:  :exc:`versile.crypto.VCryptoException`

        Raises an exception if the message failed to verify against
        the message hash, meaning the message cannot be trusted and
        could have been tampered with.

        """
        if self._invalid:
            raise VCryptoException('Message failed to verify')
        return (self._result is not None)

    def result(self):
        """Returns plaintext of a decrypted and decoded message.

        :returns: decoded plaintext
        :rtype:   bytes
        :raises:  :exc:`versile.crypto.VCryptoException`

        Should only be called if :meth:`done` indicates message
        parsing was completed, otherwise an exception may be raised
        due to incomplete message.

        Raises an exception if the message failed to verify against
        the message hash, meaning the message cannot be trusted and
        may have been tampered with.

        """
        if self._invalid:
            raise VCryptoException('Message failed to verify')
        elif self._result is None:
            raise VCryptoException('Message not yet fully decoded')
        return self._result

    def has_data(self):
        """Returns True if object holds any data

        :returns: True if holds data
        :rtype:   bool

        Returns False only if no data has been read since the object
        was instantiated or since the most recent :meth:`reset`\ .

        """
        return bool(self._read_buf or self._in_buf or self._msg_buf
                    or self._result)

    @property
    def max_plaintext_len(self):
        """Maximum plaintext length allowed in a message."""
        return self._max_plaintext_len
Exemplo n.º 2
0
class VByteChannel(object):
    """Producer/consumer end-point for byte data.

    :param reactor:  reactor driving the socket's event handling
    :param buf_len:  buffer length for read operations
    :type  buf_len:  int

    This class is primarily intended for debugging byte producer/consumer
    I/O chains.

    """
    def __init__(self, reactor, buf_len=4096):
        self.__reactor = reactor
        self.__buf_len = buf_len

        self.__bc_consumed = 0
        self.__bc_consume_lim = 0
        self.__bc_producer = None
        self.__bc_eod = False
        self.__bc_eod_clean = None
        self.__bc_rbuf = VByteBuffer()
        self.__bc_rbuf_len = buf_len
        self.__bc_reader = None
        self.__bc_aborted = False
        self.__bc_cond = Condition()
        self.__bc_scheduled_lim_update = False

        self.__bp_produced = 0
        self.__bp_produce_lim = 0
        self.__bp_consumer = None
        self.__bp_eod = False
        self.__bp_eod_clean = None
        self.__bp_wbuf = VByteBuffer()
        self.__bp_wbuf_len = buf_len
        self.__bp_writer = None
        self.__bp_sent_eod = False
        self.__bp_aborted = False
        self.__bp_cond = Condition()
        self.__bp_scheduled_produce = False

        self.__bc_iface = self.__bp_iface = None

        # Set up a local logger for convenience
        self.__logger = VLogger(prefix='ByteChannel')
        self.__logger.add_watcher(self.reactor.log)

    def __del__(self):
        self.__logger.debug('Dereferenced')

    def recv(self, max_read, timeout=None):
        """Receive input data from byte channel.

        :param max_read: max bytes to read (unlimited if None)
        :type  max_read: int
        :param timeout:  timeout in seconds (blocking if None)
        :type  timeout:  float
        :returns:        data read (empty if input was closed)
        :rtype:          bytes
        :raises:         :exc:`versile.reactor.io.VIOTimeout`\ ,
                         :exc:`versile.reactor.io.VIOError`

        """
        if timeout:
            start_time = time.time()
        with self.__bc_cond:
            while True:
                if self.__bc_rbuf:
                    if max_read is None:
                        result = self.__bc_rbuf.pop()
                    elif max_read > 0:
                        result = self.__bc_rbuf.pop(max_read)
                    else:
                        result = b''

                    # Trigger updating can_produce in reactor thread
                    if not self.__bc_scheduled_lim_update:
                        self.__bc_scheduled_lim_update = True
                        self.reactor.schedule(0.0, self.__bc_lim_update)

                    return result
                elif self.__bc_aborted:
                    raise VIOError('Byte input was aborted')
                elif self.__bc_eod:
                    if self.__bc_eod_clean:
                        return b''
                    else:
                        raise VIOError('Byte input was closed but not cleanly')

                if timeout == 0.0:
                    raise VIOTimeout()
                elif timeout is not None and timeout > 0.0:
                    current_time = time.time()
                    if current_time > start_time + timeout:
                        raise VIOTimeout()
                    wait_time = start_time + timeout - current_time
                    self.__bc_cond.wait(wait_time)
                else:
                    self.__bc_cond.wait()

    def send(self, data, timeout=None):
        """Receive input data from byte channel.

        :param data:     data to write
        :type  data:     bytes
        :type  max_read: int
        :param timeout:  timeout in seconds (blocking if None)
        :type  timeout:  float
        :returns:        number bytes written
        :rtype:          int
        :raises:         :exc:`versile.reactor.io.VIOTimeout`\ ,
                         :exc:`versile.reactor.io.VIOError`

        """
        if timeout:
            start_time = time.time()
        with self.__bp_cond:
            while True:
                if self.__bp_aborted:
                    raise VIOError('Byte output was aborted')
                elif self.__bp_eod:
                    raise VIOError('Byte output was closed')
                if not data:
                    return 0
                max_write = self.__bp_wbuf_len - len(self.__bp_wbuf)
                if max_write > 0:
                    write_data = data[:max_write]
                    self.__bp_wbuf.append(write_data)

                    # Trigger reactor production
                    if not self.__bp_scheduled_produce:
                        self.__bp_scheduled_produce = True
                        self.reactor.schedule(0.0, self.__bp_do_produce)

                    return len(write_data)

                if timeout == 0.0:
                    raise VIOTimeout()
                elif timeout is not None and timeout > 0.0:
                    current_time = time.time()
                    if current_time > start_time + timeout:
                        raise VIOTimeout()
                    wait_time = start_time + timeout - current_time
                    self.__bc_cond.wait(wait_time)
                else:
                    self.__bc_cond.wait()

    def close(self):
        """Closes the connection."""
        def _close():
            if not self.__bp_aborted and not self.__bp_eod:
                self.__bp_eod = True
                self.__bp_eod_clean = True
                self.__bp_do_produce()
            if not self.__bc_aborted and not self.__bc_eod:
                self.__bc_eod = True
                self.__bc_eod_clean = True

        self.reactor.schedule(0.0, _close)

    def abort(self):
        """Aborts the connection."""
        def _abort():
            self._bc_abort()
            self._bp_abort()

        self.reactor.schedule(0.0, _abort)

    @property
    def byte_consume(self):
        """Holds the Byte Consumer interface to the serializer."""
        cons = None
        if self.__bc_iface:
            cons = self.__bc_iface()
        if not cons:
            cons = _VByteConsumer(self)
            self.__bc_iface = weakref.ref(cons)
        return cons

    @property
    def byte_produce(self):
        """Holds the Byte Producer interface to the serializer."""
        prod = None
        if self.__bp_iface:
            prod = self.__bp_iface()
        if not prod:
            prod = _VByteProducer(self)
            self.__bp_iface = weakref.ref(prod)
        return prod

    @property
    def byte_io(self):
        """Byte interface (\ :class:`versile.reactor.io.VByteIOPair`\ )."""
        return VByteIOPair(self.byte_consume, self.byte_produce)

    @property
    def reactor(self):
        """Holds the object's reactor."""
        return self.__reactor

    @peer
    def _bc_consume(self, data, clim):
        if self.__bc_eod:
            raise VIOError('Consumer already received end-of-data')
        elif not self._bc_producer:
            raise VIOError('No connected producer')
        elif not data:
            raise VIOError('No data to consume')
        max_cons = self.__lim(self.__bc_consumed, self.__bc_consume_lim)
        if max_cons == 0:
            raise VIOError('Consume limit exceeded')
        if clim is not None and clim > 0:
            max_cons = min(max_cons, clim)

        with self.__bc_cond:
            buf_len = len(self.__bc_rbuf)
            self.__bc_rbuf.append_list(data.pop_list(max_cons))
            self.__bc_consumed += len(self.__bc_rbuf) - buf_len

            # Update consume limit
            max_add = self.__lim(len(self.__bc_rbuf), self.__bc_rbuf_len)
            if max_add >= 0:
                self.__bc_consume_lim = self.__bc_consumed + max_add
            else:
                self.__bc_consume_lim = -1

            # Notify data is available
            self.__bc_cond.notify_all()

        return self.__bc_consume_lim

    @peer
    def _bc_end_consume(self, clean):
        if self.__bc_eod:
            return
        self.__bc_eod = True
        self.__bc_eod_clean = clean

        with self.__bc_cond:
            self.__bc_cond.notify_all()

    def _bc_abort(self):
        if not self.__bc_aborted:
            with self.__bc_cond:
                self.__bc_aborted = True
                self.__bc_eod = True
                if self.__bc_producer:
                    self.__bc_producer.abort()
                    self._bc_detach()
                self.__bc_cond.notify_all()

    def _bc_attach(self, producer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._bc_attach, producer, rthread=True)
            return

        if self.__bc_producer is producer:
            return
        if self.__bc_eod:
            raise VIOError('Consumer already received end-of-data')
        elif self.__bc_producer:
            raise VIOError('Producer already connected')
        self.__bc_producer = producer
        self.__bc_consumed = 0
        self.__bc_consume_lim = self.__lim(len(self.__bc_rbuf),
                                           self.__bc_rbuf_len)
        producer.attach(self.byte_consume)
        producer.can_produce(self.__bc_consume_lim)

        # Notify attached chain
        try:
            producer.control.notify_consumer_attached(self.byte_consume)
        except VIOMissingControl:
            pass

    def _bc_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._bc_detach, rthread=True)
            return

        if self.__bc_producer:
            prod, self.__bc_producer = self.__bc_producer, None
            self.__bc_consumed = self.__bc_consume_lim = 0
            prod.detach()

    @peer
    def _bp_can_produce(self, limit):
        if not self.__bp_consumer:
            raise VIOError('No attached consumer')
        if limit is None or limit < 0:
            if (not self.__bp_produce_lim is None
                    and not self.__bp_produce_lim < 0):
                self.__bp_produce_lim = limit
                if not self.__bp_scheduled_produce:
                    self.__bp_scheduled_produce = True
                    self.reactor.schedule(0.0, self.__bp_do_produce)
        else:
            if (self.__bp_produce_lim is not None
                    and 0 <= self.__bp_produce_lim < limit):
                self.__bp_produce_lim = limit
                if not self.__bp_scheduled_produce:
                    self.__bp_scheduled_produce = True
                    self.reactor.schedule(0.0, self.__bp_do_produce)

    def _bp_abort(self):
        if not self.__bp_aborted:
            with self.__bp_cond:
                self.__bp_aborted = True
                self.__bp_wbuf.remove()
                if self.__bp_consumer:
                    self.__bp_consumer.abort()
                    self._bp_detach()
                self.__bp_cond.notify_all()

    def _bp_attach(self, consumer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._bp_attach, consumer, rthread=True)
            return

        if self.__bp_consumer is consumer:
            return
        if self.__bp_consumer:
            raise VIOError('Consumer already attached')
        elif self.__bp_eod:
            raise VIOError('Producer already reached end-of-data')
        self.__bp_consumer = consumer
        self.__bp_produced = self.__bp_produce_lim = 0
        consumer.attach(self.byte_produce)

        # Notify attached chain
        try:
            consumer.control.notify_producer_attached(self.byte_produce)
        except VIOMissingControl:
            pass

    def _bp_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._bp_detach, rthread=True)
            return

        if self.__bp_consumer:
            cons, self.__bp_consumer = self.__bp_consumer, None
            cons.detach()
            self.__bp_produced = self.__bp_produce_lim = 0

    @property
    def _bc_control(self):
        return VIOControl()

    @property
    def _bc_producer(self):
        return self.__bc_producer

    @property
    def _bc_flows(self):
        return (self.entity_produce, )

    @property
    def _bc_twoway(self):
        return True

    @property
    def _bc_reverse(self):
        return self.byte_produce

    @property
    def _bp_control(self):
        return VIOControl()

    @property
    def _bp_consumer(self):
        return self.__bp_consumer

    @property
    def _bp_flows(self):
        return (self.entity_consume, )

    @property
    def _bp_twoway(self):
        return True

    @property
    def _bp_reverse(self):
        return self.byte_consume

    def __bc_lim_update(self):
        self.__bc_scheduled_lim_update = False

        if not self.__bc_producer or self.__bc_aborted or self.__bc_eod:
            return

        old_lim = self.__bc_consume_lim
        self.__bc_consume_lim = self.__lim(len(self.__bc_rbuf),
                                           self.__bc_rbuf_len)
        if old_lim != self.__bc_consume_lim:
            self.__bc_producer.can_produce(self.__bc_consume_lim)

    def __bp_do_produce(self):
        self.__bp_scheduled_produce = False

        if not self.__bp_consumer:
            return

        with self.__bp_cond:
            if self.__bp_eod:
                # If end-of-data was reached notify consumer
                if self.__bp_consumer and not self.__bp_sent_eod:
                    self.__bp_consumer.end_consume(self.__bp_eod_clean)
                    self.__bp_sent_eod = True
                return

            if (self.__bp_produce_lim is not None
                    and 0 <= self.__bp_produce_lim <= self.__bp_produced):
                return

            old_lim = self.__bp_produce_lim
            max_write = self.__lim(self.__bp_produced, self.__bp_produce_lim)
            if max_write != 0 and self.__bp_wbuf:
                old_len = len(self.__bp_wbuf)
                new_lim = self.__bp_consumer.consume(self.__bp_wbuf)
                self.__bp_produced += self.__bp_wbuf_len - len(self.__bp_wbuf)
                self.__bp_produce_lim = new_lim
                if old_len != len(self.__bp_wbuf):
                    self.__bp_cond.notify_all()

            # Schedule another if produce limit was updated and buffer has data
            if self.__bp_wbuf and self.__bp_produce_lim != old_lim:
                if not self.__bp_scheduled_produce:
                    self.__bp_scheduled_produce = True
                    self.reactor.schedule(0.0, self.__bp_do_produce)

    @classmethod
    def __lim(self, base, *lims):
        """Return smallest (lim-base) limit, or -1 if all limits are <0"""
        result = -1
        for lim in lims:
            if lim is not None and lim >= 0:
                lim = max(lim - base, 0)
                if result < 0:
                    result = lim
                result = min(result, lim)
        return result
Exemplo n.º 3
0
class VClientSocketAgent(VClientSocket):
    """A :class:`VClientSocket` with a byte producer/consumer interface.

    :param max_read:  max bytes fetched per socket read
    :type  max_read:  int
    :param max_write: max bytes written per socket send
    :type  max_write: int
    :param wbuf_len:  buffer size of data held for writing (or None)
    :type  wbuf_len:  int

    *max_read* is also the maximum size of the buffer for data read
    from the socket (so maximum bytes read in one read operation is
    *max_read* minus the amount of data currently held in the receive
    buffer).

    If *wbuf_len* is None then *max_write* is used as the buffer size.

    """
    def __init__(self,
                 reactor,
                 sock=None,
                 hc_pol=None,
                 close_cback=None,
                 connected=False,
                 max_read=0x4000,
                 max_write=0x4000,
                 wbuf_len=None):
        self._max_read = max_read
        self._max_write = max_write
        self._wbuf = VByteBuffer()
        if wbuf_len is None:
            wbuf_len = max_write
        self._wbuf_len = wbuf_len

        self._ci = None
        self._ci_eod = False
        self._ci_eod_clean = None
        self._ci_producer = None
        self._ci_consumed = 0
        self._ci_lim_sent = 0
        self._ci_aborted = False

        self._pi = None
        self._pi_closed = False
        self._pi_consumer = None
        self._pi_produced = 0
        self._pi_prod_lim = 0
        self._pi_buffer = VByteBuffer()
        self._pi_aborted = False

        # Parent __init__ must be called after local attributes are
        # initialized due to overloaded methods called during construction
        super_init = super(VClientSocketAgent, self).__init__
        super_init(reactor=reactor,
                   sock=sock,
                   hc_pol=hc_pol,
                   close_cback=close_cback,
                   connected=connected)

    @property
    def byte_consume(self):
        """Holds a :class:`IVByteConsumer` interface to the socket."""
        if not self._ci:
            ci = _VSocketConsumer(self)
            self._ci = weakref.ref(ci)
            return ci
        else:
            ci = self._ci()
            if ci:
                return ci
            else:
                ci = _VSocketConsumer(self)
                self._ci = weakref.ref(ci)
                return ci

    @property
    def byte_produce(self):
        """Holds a :class:`IVByteProducer` interface to the socket."""
        if not self._pi:
            pi = _VSocketProducer(self)
            self._pi = weakref.ref(pi)
            return pi
        else:
            pi = self._pi()
            if pi:
                return pi
            else:
                pi = _VSocketProducer(self)
                self._pi = weakref.ref(pi)
                return pi

    @property
    def byte_io(self):
        """Byte interface (\ :class:`versile.reactor.io.VByteIOPair`\ )."""
        return VByteIOPair(self.byte_consume, self.byte_produce)

    def _can_connect(self, peer):
        """Called internally to validate whether a connection can be made.

        :param peer: peer to connect to
        :type  peer: :class:`versile.common.peer.VPeer`
        :returns:    True if connection is allowed
        :rtype:      bool

        If the socket is connected to a byte consumer, sends a control
        request for 'can_connect(peer)'. If that control message
        returns a boolean, this is used as the _can_connect
        result. Otherwise, True is returned.

        """
        if self._pi_consumer:
            try:
                return self._pi_consumer.control.can_connect(peer)
            except VIOMissingControl:
                return True
        else:
            return True

    def _sock_was_connected(self):
        super(VClientSocketAgent, self)._sock_was_connected()
        if self._p_consumer:
            # Notify producer-connected chain about 'connected' status
            control = self._p_consumer.control

            def notify():
                try:
                    control.connected(self._sock_peer)
                except VIOMissingControl:
                    pass

            self.reactor.schedule(0.0, notify)

    @peer
    def _c_consume(self, buf, clim):
        if self._ci_eod:
            raise VIOClosed('Consumer already reached end-of-data')
        elif not self._ci_producer:
            raise VIOError('No connected producer')
        elif self._ci_consumed >= self._ci_lim_sent:
            raise VIOError('Consume limit exceeded')
        elif not buf:
            raise VIOError('No data to consume')

        max_cons = self._wbuf_len - len(self._wbuf)
        max_cons = min(max_cons, self._ci_lim_sent - self._ci_consumed)
        if clim is not None and clim > 0:
            max_cons = min(max_cons, clim)

        was_empty = not self._wbuf
        indata = buf.pop(max_cons)
        self._wbuf.append(indata)
        self._ci_consumed += len(indata)
        if was_empty:
            self.start_writing(internal=True)
        return self._ci_lim_sent

    def _c_end_consume(self, clean):
        if self._ci_eod:
            return
        self._ci_eod = True
        self._ci_eod_clean = clean

        if not self._wbuf:
            self.close_output(VFIOCompleted())
            if self._ci_producer:
                self._ci_producer.abort()
                self._c_detach()

    def _c_abort(self, force=False):
        if not self._ci_aborted or force:
            self._ci_aborted = True
            self._ci_eod = True
            self._ci_consumed = self._ci_lim_sent = 0
            self._wbuf.clear()
            if not self._sock_out_closed:
                self.close_output(VFIOCompleted())
            if self._ci_producer:
                self._ci_producer.abort()
                self._c_detach()

    def _c_attach(self, producer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._c_attach, producer, rthread=True)
            return

        if self._ci_producer is producer:
            return
        elif self._ci_producer:
            raise VIOError('Producer already attached')

        self._ci_producer = producer
        if self._sock_out_closed:
            self.reactor.schedule(0.0, self._c_abort, True)
        else:
            self._ci_consumed = self._ci_lim_sent = 0
            producer.attach(self.byte_consume)
            self._ci_lim_sent = self._wbuf_len
            producer.can_produce(self._ci_lim_sent)

        # Notify attached chain
        try:
            producer.control.notify_consumer_attached(self.byte_consume)
        except VIOMissingControl:
            pass

    def _c_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._c_detach, rthread=True)
            return

        if self._ci_producer:
            prod, self._ci_producer = self._ci_producer, None
            prod.detach()
            self._ci_consumed = self._ci_lim_sent = 0

    def active_do_write(self):
        if self._wbuf:
            data = self._wbuf.peek(self._max_write)
            try:
                num_written = self.write_some(data)
            except VIOException:
                self._c_abort()
            else:
                if num_written > 0:
                    self._wbuf.remove(num_written)
                    if self._ci_producer:
                        self._ci_lim_sent = (self._ci_consumed +
                                             self._wbuf_len - len(self._wbuf))
                        self._ci_producer.can_produce(self._ci_lim_sent)
                if not self._wbuf:
                    self.stop_writing()
                    if self._ci_eod:
                        self._c_abort()
        else:
            self.stop_writing()

    def _output_was_closed(self, reason):
        # No more output will be written, abort consumer
        self._c_abort()

    @property
    def _c_control(self):
        return self._c_get_control()

    # Implements _c_control in order to be able to override _c_control
    # behavior by overloading as a regular method
    def _c_get_control(self):
        return VIOControl()

    @property
    def _c_producer(self):
        return self._ci_producer

    @property
    def _c_flows(self):
        return tuple()

    @property
    def _c_twoway(self):
        return True

    @property
    def _c_reverse(self):
        return self.byte_produce()

    @peer
    def _p_can_produce(self, limit):
        if not self._pi_consumer:
            raise VIOError('No connected consumer')

        if limit is None or limit < 0:
            if (not self._pi_prod_lim is None and not self._pi_prod_lim < 0):
                if self._pi_produced >= self._pi_prod_lim:
                    self.start_reading(internal=True)
                self._pi_prod_lim = limit
        else:
            if (self._pi_prod_lim is not None
                    and 0 <= self._pi_prod_lim < limit):
                if self._pi_produced >= self._pi_prod_lim:
                    self.start_reading(internal=True)
                self._pi_prod_lim = limit

    def _p_abort(self, force=False):
        if not self._pi_aborted or force:
            self._pi_aborted = True
            self._pi_produced = self._pi_prod_lim = 0
            if not self._sock_in_closed:
                self.close_input(VFIOCompleted())
            if self._pi_consumer:
                self._pi_consumer.abort()
                self._p_detach()

    def _p_attach(self, consumer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._p_attach, consumer, rthread=True)
            return

        if self._pi_consumer is consumer:
            return
        elif self._pi_consumer:
            raise VIOError('Consumer already attached')

        self._pi_produced = self._pi_prod_lim = 0
        self._pi_consumer = consumer
        consumer.attach(self.byte_produce)

        # Notify attached chain
        try:
            consumer.control.notify_producer_attached(self.byte_produce)
        except VIOMissingControl:
            pass

        # If closed, pass notification
        if self._sock_in_closed:
            self.reactor.schedule(0.0, self._p_abort, True)

    def _p_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._p_detach, rthread=True)
            return

        if self._pi_consumer:
            cons, self._pi_consumer = self._pi_consumer, None
            cons.detach()
            self._pi_produced = self._pi_prod_lim = 0

    def active_do_read(self):
        if not self._pi_consumer:
            self.stop_reading()

        if self._pi_prod_lim is not None and self._pi_prod_lim >= 0:
            max_read = self._pi_prod_lim - self._pi_produced
        else:
            max_read = self._max_read
        max_read = min(max_read, self._max_read)
        if max_read <= 0:
            self.stop_reading()
            return
        try:
            data = self.read_some(max_read)
        except Exception as e:
            self._p_abort()
        else:
            self._pi_buffer.append(data)
            if self._pi_buffer:
                self.pi_prod_lim = self._pi_consumer.consume(self._pi_buffer)

    def _input_was_closed(self, reason):
        if self._pi_consumer:
            # Notify consumer about end-of-data
            clean = isinstance(reason, VFIOCompleted)
            self._pi_consumer.end_consume(clean)
        else:
            self._p_abort()

    @property
    def _p_control(self):
        return self._p_get_control()

    # Implements _p_control in order to be able to override _p_control
    # behavior by overloading as a regular method
    def _p_get_control(self):
        class _Control(VIOControl):
            def __init__(self, sock):
                self.__sock = sock

            def req_producer_state(self, consumer):
                # Send notification of socket connect status
                def notify():
                    if self.__sock._sock_peer:
                        try:
                            consumer.control.connected(self.__sock._sock_peer)
                        except VIOMissingControl:
                            pass

                self.__sock.reactor.schedule(0.0, notify)

        return _Control(self)

    @property
    def _p_consumer(self):
        return self._pi_consumer

    @property
    def _p_flows(self):
        return tuple()

    @property
    def _p_twoway(self):
        return True

    @property
    def _p_reverse(self):
        return self.byte_consume()
Exemplo n.º 4
0
class VPipeAgent(object):
    """Byte producer/consumer interface to a pipe reader/writer pair.

    :param read_fd:   read pipe file descriptor
    :type  read_fd:   int
    :param write_fd:  write pipe file descriptor
    :type  write_fd:  int
    :param max_read:  max bytes fetched per pipe read
    :type  max_read:  int
    :param max_write: max bytes written per pipe write
    :type  max_write: int
    :param wbuf_len:  buffer size of data held for writing (or None)
    :type  wbuf_len:  int

    The agent creates a :class:`VPipeReader` and :class:`VPipeWriter`
    for the provided pipe read/write descriptors which it uses for
    reactor driven pipe I/O communication.

    *max_read* is also the maximum size of the buffer for data read
    from the socket (so maximum bytes read in one read operation is
    *max_read* minus the amount of data currently held in the receive
    buffer).

    If *wbuf_len* is None then *max_write* is used as the buffer size.

    """
    def __init__(self,
                 reactor,
                 read_fd,
                 write_fd,
                 max_read=0x4000,
                 max_write=0x4000,
                 wbuf_len=None):
        self.__reactor = reactor

        self._reader = _VAgentReader(self, reactor, read_fd)
        self._writer = _VAgentWriter(self, reactor, write_fd)
        self._reader.set_pipe_peer(self._writer)
        self._writer.set_pipe_peer(self._reader)

        self._max_read = max_read
        self._max_write = max_write
        self._wbuf = VByteBuffer()
        if wbuf_len is None:
            wbuf_len = max_write
        self._wbuf_len = wbuf_len

        self._ci = None
        self._ci_eod = False
        self._ci_eod_clean = None
        self._ci_producer = None
        self._ci_consumed = 0
        self._ci_lim_sent = 0
        self._ci_aborted = False

        self._pi = None
        self._pi_closed = False
        self._pi_consumer = None
        self._pi_produced = 0
        self._pi_prod_lim = 0
        self._pi_buffer = VByteBuffer()
        self._pi_aborted = False

    @property
    def byte_consume(self):
        """Holds a :class:`IVByteConsumer` interface to the pipe reader."""
        if not self._ci:
            ci = _VPipeConsumer(self)
            self._ci = weakref.ref(ci)
            return ci
        else:
            ci = self._ci()
            if ci:
                return ci
            else:
                ci = _VPipeConsumer(self)
                self._ci = weakref.ref(ci)
                return ci

    @property
    def byte_produce(self):
        """Holds a :class:`IVByteProducer` interface to the pipe writer."""
        if not self._pi:
            pi = _VPipeProducer(self)
            self._pi = weakref.ref(pi)
            return pi
        else:
            pi = self._pi()
            if pi:
                return pi
            else:
                pi = _VPipeProducer(self)
                self._pi = weakref.ref(pi)
                return pi

    @property
    def byte_io(self):
        """Byte interface (\ :class:`versile.reactor.io.VByteIOPair`\ )."""
        return VByteIOPair(self.byte_consume, self.byte_produce)

    @property
    def reactor(self):
        """Holds the reactor of the associated reader."""
        return self.__reactor

    @peer
    def _c_consume(self, buf, clim):
        if self._ci_eod:
            raise VIOClosed('Consumer already reached end-of-data')
        elif not self._ci_producer:
            raise VIOError('No connected producer')
        elif self._ci_consumed >= self._ci_lim_sent:
            raise VIOError('Consume limit exceeded')
        elif not buf:
            raise VIOError('No data to consume')

        max_cons = self._wbuf_len - len(self._wbuf)
        max_cons = min(max_cons, self._ci_lim_sent - self._ci_consumed)
        if clim is not None and clim > 0:
            max_cons = min(max_cons, clim)

        was_empty = not self._wbuf
        indata = buf.pop(max_cons)
        self._wbuf.append(indata)
        self._ci_consumed += len(indata)
        if was_empty:
            self._writer.start_writing(internal=True)
        return self._ci_lim_sent

    def _c_end_consume(self, clean):
        if self._ci_eod:
            return
        self._ci_eod = True
        self._ci_eod_clean = clean

        if not self._wbuf:
            self._writer.close_output(VFIOCompleted())
            if self._ci_producer:
                self._ci_producer.abort()
                self._c_detach()

    def _c_abort(self):
        if not self._ci_aborted:
            self._ci_aborted = True
            self._ci_eod = True
            self._ci_consumed = self._ci_lim_sent = 0
            self._wbuf.clear()
            self._writer.close_output(VFIOCompleted())
            if self._ci_producer:
                self._ci_producer.abort()
                self._c_detach()

    def _c_attach(self, producer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._c_attach, producer, rthread=True)
            return

        if self._ci_producer is producer:
            return
        elif self._ci_producer:
            raise VIOError('Producer already attached')

        self._ci_producer = producer
        self._ci_consumed = self._ci_lim_sent = 0
        producer.attach(self.byte_consume)
        self._ci_lim_sent = self._wbuf_len
        producer.can_produce(self._ci_lim_sent)

        # Notify attached chain
        try:
            producer.control.notify_consumer_attached(self.byte_consume)
        except VIOMissingControl:
            pass

    def _c_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._c_detach, rthread=True)
            return

        if self._ci_producer:
            prod, self._ci_producer = self._ci_producer, None
            prod.detach()
            self._ci_consumed = self._ci_lim_sent = 0

    def _do_write(self):
        if self._wbuf:
            data = self._wbuf.peek(self._max_write)
            try:
                num_written = self._writer.write_some(data)
            except VIOException:
                self._c_abort()
            else:
                if num_written > 0:
                    self._wbuf.remove(num_written)
                    if self._ci_producer:
                        self._ci_lim_sent = (self._ci_consumed +
                                             self._wbuf_len - len(self._wbuf))
                        self._ci_producer.can_produce(self._ci_lim_sent)
                if not self._wbuf:
                    self._writer.stop_writing()
                    if self._ci_eod:
                        self._c_abort()
        else:
            self._writer.stop_writing()

    def _output_was_closed(self, reason):
        # No more output will be written, abort consumer
        self._c_abort()

    @property
    def _c_control(self):
        return VIOControl()

    @property
    def _c_producer(self):
        return self._ci_producer

    @property
    def _c_flows(self):
        return tuple()

    @property
    def _c_twoway(self):
        return True

    @property
    def _c_reverse(self):
        return self.byte_produce()

    @peer
    def _p_can_produce(self, limit):
        if not self._pi_consumer:
            raise VIOError('No connected consumer')

        if limit is None or limit < 0:
            if (not self._pi_prod_lim is None and not self._pi_prod_lim < 0):
                if self._pi_produced >= self._pi_prod_lim:
                    self._reader.start_reading(internal=True)
                self._pi_prod_lim = limit
        else:
            if (self._pi_prod_lim is not None
                    and 0 <= self._pi_prod_lim < limit):
                if self._pi_produced >= self._pi_prod_lim:
                    self._reader.start_reading(internal=True)
                self._pi_prod_lim = limit

    def _p_abort(self):
        if not self._pi_aborted:
            self._pi_aborted = True
            self._pi_produced = self._pi_prod_lim = 0
            self._reader.close_input(VFIOCompleted())
            if self._pi_consumer:
                self._pi_consumer.abort()
                self._p_detach()

    def _p_attach(self, consumer, rthread=False):
        # Ensure 'attach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._p_attach, consumer, rthread=True)
            return

        if self._pi_consumer is consumer:
            return
        elif self._pi_consumer:
            raise VIOError('Consumer already attached')

        self._pi_produced = self._pi_prod_lim = 0
        self._pi_consumer = consumer
        consumer.attach(self.byte_produce)

        # Notify attached chain
        try:
            consumer.control.notify_producer_attached(self.byte_produce)
        except VIOMissingControl:
            pass

    def _p_detach(self, rthread=False):
        # Ensure 'detach' is performed in reactor thread
        if not rthread:
            self.reactor.execute(self._p_detach, rthread=True)
            return

        if self._pi_consumer:
            cons, self._pi_consumer = self._pi_consumer, None
            cons.detach()
            self._pi_produced = self._pi_prod_lim = 0

    def _do_read(self):
        if not self._pi_consumer:
            self._reader.stop_reading()

        if self._pi_prod_lim is not None and self._pi_prod_lim >= 0:
            max_read = self._pi_prod_lim - self._pi_produced
        else:
            max_read = self._max_read
        max_read = min(max_read, self._max_read)
        if max_read <= 0:
            self._reader.stop_reading()
            return
        try:
            data = self._reader.read_some(max_read)
        except Exception as e:
            self._p_abort()
        else:
            self._pi_buffer.append(data)
            if self._pi_buffer:
                self.pi_prod_lim = self._pi_consumer.consume(self._pi_buffer)

    def _input_was_closed(self, reason):
        if self._pi_consumer:
            # Notify consumer about end-of-data
            clean = isinstance(reason, VFIOCompleted)
            self._pi_consumer.end_consume(clean)
        else:
            self._p_abort()

    @property
    def _p_control(self):
        class _Control(VIOControl):
            def __init__(self, obj):
                self._obj = obj

            def req_producer_state(self, consumer):
                # Send 'connected' notification if pipe is not closed
                def notify():
                    if (self._obj._reader._fd >= 0
                            and self._obj._writer._fd >= 0):
                        try:
                            consumer.control.connected(VPipePeer())
                        except VIOMissingControl:
                            pass

                self._obj.reactor.schedule(0.0, notify)

        return _Control(self)

    @property
    def _p_consumer(self):
        return self._pi_consumer

    @property
    def _p_flows(self):
        return tuple()

    @property
    def _p_twoway(self):
        return True

    @property
    def _p_reverse(self):
        return self.byte_consume()