PyBitmessage-2024-12-08/src/network/advanceddispatcher.py

"""
src/network/advanceddispatcher.py
=================================
"""
# pylint: disable=attribute-defined-outside-init

import socket
import threading
import time

import network.asyncore_pollchoose as asyncore
import state
from threads import BusyError, nonBlocking


class ProcessingError(Exception):
    """General class for protocol parser exception, use as a base for others."""
    pass


class UnknownStateError(ProcessingError):
    """Parser points to an unknown (unimplemented) state."""
    pass


class AdvancedDispatcher(asyncore.dispatcher):
    """Improved version of asyncore dispatcher, with buffers and protocol state."""
    # pylint: disable=too-many-instance-attributes
    _buf_len = 131072  # 128kB

    def __init__(self, sock=None):
        if not hasattr(self, '_map'):
            asyncore.dispatcher.__init__(self, sock)
        self.read_buf = bytearray()
        self.write_buf = bytearray()
        self.state = "init"
        self.lastTx = time.time()
        self.sentBytes = 0
        self.receivedBytes = 0
        self.expectBytes = 0
        self.readLock = threading.RLock()
        self.writeLock = threading.RLock()
        self.processingLock = threading.RLock()

    def append_write_buf(self, data):
        """Append binary data to the end of stream write buffer."""
        if data:
            if isinstance(data, list):
                with self.writeLock:
                    for chunk in data:
                        self.write_buf.extend(chunk)
            else:
                with self.writeLock:
                    self.write_buf.extend(data)

    def slice_write_buf(self, length=0):
        """Cut the beginning of the stream write buffer."""
        if length > 0:
            with self.writeLock:
                if length >= len(self.write_buf):
                    del self.write_buf[:]
                else:
                    del self.write_buf[0:length]

    def slice_read_buf(self, length=0):
        """Cut the beginning of the stream read buffer."""
        if length > 0:
            with self.readLock:
                if length >= len(self.read_buf):
                    del self.read_buf[:]
                else:
                    del self.read_buf[0:length]

    def process(self):
        """Process (parse) data that's in the buffer, as long as there is enough data and the connection is open."""
        while self.connected and not state.shutdown:
            try:
                with nonBlocking(self.processingLock):
                    if not self.connected or state.shutdown:
                        break
                    if len(self.read_buf) < self.expectBytes:
                        return False
                    try:
                        cmd = getattr(self, "state_" + str(self.state))
                    except AttributeError:
                        self.logger.error(
                            'Unknown state %s', self.state, exc_info=True)
                        raise UnknownStateError(self.state)
                    if not cmd():
                        break
            except BusyError:
                return False
        return False

    def set_state(self, state_str, length=0, expectBytes=0):
        """Set the next processing state."""
        self.expectBytes = expectBytes
        self.slice_read_buf(length)
        self.state = state_str

    def writable(self):
        """Is data from the write buffer ready to be sent to the network?"""
        self.uploadChunk = AdvancedDispatcher._buf_len
        if asyncore.maxUploadRate > 0:
            self.uploadChunk = int(asyncore.uploadBucket)
        self.uploadChunk = min(self.uploadChunk, len(self.write_buf))
        return asyncore.dispatcher.writable(self) and \
            (self.connecting or (self.connected and self.uploadChunk > 0))

    def readable(self):
        """Is the read buffer ready to accept data from the network?"""
        self.downloadChunk = AdvancedDispatcher._buf_len
        if asyncore.maxDownloadRate > 0:
            self.downloadChunk = int(asyncore.downloadBucket)
        try:
            if self.expectBytes > 0 and not self.fullyEstablished:
                self.downloadChunk = min(self.downloadChunk, self.expectBytes - len(self.read_buf))
                if self.downloadChunk < 0:
                    self.downloadChunk = 0
        except AttributeError:
            pass
        return asyncore.dispatcher.readable(self) and \
            (self.connecting or self.accepting or (self.connected and self.downloadChunk > 0))

    def handle_read(self):
        """Append incoming data to the read buffer."""
        self.lastTx = time.time()
        newData = self.recv(self.downloadChunk)
        self.receivedBytes += len(newData)
        asyncore.update_received(len(newData))
        with self.readLock:
            self.read_buf.extend(newData)

    def handle_write(self):
        """Send outgoing data from write buffer."""
        self.lastTx = time.time()
        written = self.send(self.write_buf[0:self.uploadChunk])
        asyncore.update_sent(written)
        self.sentBytes += written
        self.slice_write_buf(written)

    def handle_connect_event(self):
        """Callback for connection established event."""
        try:
            asyncore.dispatcher.handle_connect_event(self)
        except socket.error as e:
            if e.args[0] not in asyncore._DISCONNECTED:  # pylint: disable=protected-access
                raise

    def handle_connect(self):
        """Method for handling connection established implementations."""
        self.lastTx = time.time()

    def state_close(self):
        """Signal to the processing loop to end."""
        # pylint: disable=no-self-use
        return False

    def handle_close(self):
        """Callback for connection being closed, but can also be called directly when you want connection to close."""
        with self.readLock:
            self.read_buf = bytearray()
        with self.writeLock:
            self.write_buf = bytearray()
        self.set_state("close")
        self.close()