pybricks-api/examples/ev3/bluetooth_pc/pybricks/messaging.py

# SPDX-License-Identifier: MIT
# Copyright (C) 2020 The Pybricks Authors

from _thread import allocate_lock
from errno import ECONNRESET
from struct import pack, unpack

from .bluetooth import (BDADDR_ANY, ThreadingRFCOMMServer,
                        ThreadingRFCOMMClient, StreamRequestHandler)


def resolve(brick):
    """Fake resolver to get address from Bluetooth name.

    To connect to an EV3 server, you must specify its full address.
    """
    return brick


class Mailbox:
    def __init__(self, name, connection, encode=None, decode=None):
        """Object that represents a mailbox for sending an receiving messages
        from other connected devices.

        Arguments:
            name (str):
                The name of this mailbox.
            connection:
                A connection object that implements the mailbox connection
                interface.
            encode:
                A function that encodes an object into a bytes-like object.
            decode:
                A function that decodes an object from a bytes-like object.
        """
        self.name = name
        self._connection = connection

        if encode:
            self.encode = encode

        if decode:
            self.decode = decode

    def encode(self, value):
        return value

    def decode(self, payload):
        return payload

    def read(self):
        """Reads the current value of the mailbox.

        Returns:
            The decoded value or ``None`` if the mailbox has never received
            a value.
        """
        data = self._connection.read_from_mailbox(self.name)
        if data is None:
            return None
        return self.decode(data)

    def send(self, value, destination=None):
        """Sends a value to remote mailboxes with the same name as this
        mailbox.

        Arguments:
            value: The value to send.
            destination: The name or address of a specific device or ``None``
                to broadcast to all connected devices.
        """
        data = self.encode(value)
        self._connection.send_to_mailbox(destination, self.name, data)

    def wait(self):
        """Waits for the mailbox to receive a message."""
        self._connection.wait_for_mailbox_update(self.name)

    def wait_new(self):
        """Waits for the mailbox to receive a message that is different from
        the current contents of the mailbox.

        Returns:
            The new value. (Same as return value of :meth:`read`.)
        """
        old = self.read()
        while True:
            self.wait()
            new = self.read()
            if new != old:
                return new


class LogicMailbox(Mailbox):
    """:class:`Mailbox` that holds a logic or boolean value.

    This is compatible with the "logic" message blocks in the standard
    EV3 firmware.
    """

    def encode(self, value):
        return b'\x01' if value else b'\x00'

    def decode(self, payload):
        return bool(payload[0])


class NumericMailbox(Mailbox):
    """:class:`Mailbox` that holds a numeric or floating point value.

    This is compatible with the "numeric" message blocks in the standard
    EV3 firmware.
    """

    def encode(self, value):
        return pack('<f', value)

    def decode(self, payload):
        return unpack('<f', payload)[0]


class TextMailbox(Mailbox):
    """:class:`Text` that holds a text or string point value.

    This is compatible with the "text" message blocks in the standard
    EV3 firmware.
    """

    def encode(self, value):
        return ('{}\0'.format(value)).encode('utf-8')

    def decode(self, payload):
        return payload.decode().strip('\0')


# EV3 standard firmware is hard-coded to use channel 1
EV3_RFCOMM_CHANNEL = 1

# EV3 VM bytecodes
SYSTEM_COMMAND_NO_REPLY = 0x81
WRITEMAILBOX = 0x9E


class MailboxHandler(StreamRequestHandler):
    def handle(self):
        with self.server._lock:
            self.server._clients[self.client_address[0]] = self.request
        while True:
            try:
                buf = self.rfile.recv(2)
                if len(buf) == 0:
                    break
            except OSError as ex:
                # The client disconnected the connection
                if ex.args[0] == ECONNRESET:
                    break
                raise
            size, = unpack('<H', buf)
            buf = self.rfile.recv(size)
            msg_count, cmd_type, cmd, name_size = unpack('<HBBB', buf[0:5])
            if cmd_type != SYSTEM_COMMAND_NO_REPLY:
                raise ValueError('Bad message type')
            if cmd != WRITEMAILBOX:
                raise ValueError('Bad command')
            mbox = buf[5:5+name_size].decode().strip('\0')
            data_size, = unpack('<H', buf[5+name_size:7+name_size])
            data = buf[7+name_size:7+name_size+data_size]

            with self.server._lock:
                self.server._mailboxes[mbox] = data
                update_lock = self.server._updates.get(mbox)
                if update_lock:
                    update_lock.release()


class MailboxHandlerMixIn:
    def __init__(self):
        # protects against concurrent access of other attributes
        self._lock = allocate_lock()
        # map of mailbox name to raw data
        self._mailboxes = {}
        # map of device name/address to object with send() method
        self._clients = {}
        # map of mailbox name to mutex lock
        self._updates = {}
        # map of names to addresses
        self._addresses = {}

    def read_from_mailbox(self, mbox):
        """Reads the current raw data from a mailbox.

        Arguments:
            mbox (str):
                The name of the mailbox.

        Returns:
            bytes:
                The current mailbox raw data or ``None`` if nothing has ever
                been delivered to the mailbox.
        """
        with self._lock:
            return self._mailboxes.get(mbox)

    def send_to_mailbox(self, brick, mbox, payload):
        """Sends a mailbox value using raw bytes data.

        Arguments:
            brick (str):
                The name or address of the brick or ``None``` to broadcast to
                all connected devices
            mbox (str):
                The name of the mailbox.
            payload (bytes):
                A bytes-like object that will be sent to the mailbox.
        """
        mbox_len = len(mbox) + 1
        payload_len = len(payload)
        send_len = 7 + mbox_len + payload_len
        fmt = '<HHBBB{}sH{}s'.format(mbox_len, payload_len)
        data = pack(fmt, send_len, 1, SYSTEM_COMMAND_NO_REPLY, WRITEMAILBOX,
                    mbox_len, mbox.encode('utf-8'), payload_len, payload)
        with self._lock:
            if brick is None:
                for client in self._clients.values():
                    client.send(data)
            else:
                addr = self._addresses.get(brick)
                if addr is None:
                    addr = resolve(brick)
                    self._addresses[brick] = addr
                if addr is None:
                    raise ValueError('no paired devices matching "{}"'.format(brick))
                self._clients[addr].send(data)

    def wait_for_mailbox_update(self, mbox):
        """Waits until ``mbox`` receives a value."""
        lock = allocate_lock()
        lock.acquire()
        with self._lock:
            self._updates[mbox] = lock
        try:
            return lock.acquire()
        finally:
            with self._lock:
                del self._updates[mbox]


class BluetoothMailboxServer(MailboxHandlerMixIn, ThreadingRFCOMMServer):
    def __init__(self):
        """Object that represents an incoming Bluetooth connection from another
        EV3.

        The remote EV3 can either be running MicroPython or the standard EV3
        firmare.
        """
        super().__init__()
        super(ThreadingRFCOMMServer, self).__init__(
            (BDADDR_ANY, EV3_RFCOMM_CHANNEL), MailboxHandler)

    def wait_for_connection(self, count=1):
        """Waits for a :class:`BluetoothMailboxClient` on a remote device to
        connect.

        Arguments:
            count (int):
                The number of remote connections to wait for.

        Raises:
            OSError:
                There was a problem establishing the connection.
        """
        for _ in range(count):
            self.handle_request()


class MailboxRFCOMMClient(ThreadingRFCOMMClient):
    def __init__(self, parent, bdaddr):
        self.parent = parent
        super().__init__((bdaddr, EV3_RFCOMM_CHANNEL), MailboxHandler)

    def send(self, data):
        self.socket.send(data)

    def shutdown_request(self, request):
        request.close()

    def finish_request(self, request, client_address):
        self.RequestHandlerClass(request, client_address, self.parent)


class BluetoothMailboxClient(MailboxHandlerMixIn):
    """Object that represents outgoing Bluetooth connections to one or more
    remote EV3s.

    The remote EV3s can either be running MicroPython or the standard EV3
    firmare.
    """

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def connect(self, brick):
        """Connects to a :class:`BluetoothMailboxServer` on another device.

        The remote device must be paired and waiting for a connection. See
        :meth:`BluetoothMailboxServer.wait_for_connection`.

        Arguments:
            brick (str):
                The name or address of the remote EV3 to connect to.

        Raises:
            TypeError:
                ``brick`` is not a string
            ValueError:
                There are no paired Bluetooth devices that match ``brick``
                or connection to ``brick`` already exists.
            OSError:
                There was a problem establishing the connection.
        """
        addr = resolve(brick)
        if addr is None:
            raise ValueError('no paired devices matching "{}"'.format(brick))
        client = MailboxRFCOMMClient(self, addr)
        if self._clients.setdefault(addr, client) is not client:
            raise ValueError('connection with this address already exists')
        try:
            client.handle_request()
        except Exception:
            del self._clients[addr]
            raise

    def close(self):
        """Closes the connections."""
        for client in self._clients.values():
            client.client_close()
        self._clients.clear()