sedbgmux/handlers.py

# This file is a part of sedbgmux, an open source DebugMux client.
# Copyright (c) 2022  Vadim Yanitskiy <axilirator@gmail.com>
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import logging as log
import threading
import socket
import queue
import enum
import abc

from typing import Any, Optional
from construct import Container
from proto import DbgMuxFrame


class DbgMuxFrameHandler(abc.ABC):
    ''' Abstract DebugMux frame handler '''

    _tx_queue: Optional[queue.Queue] = None

    def send(self, msg_type: DbgMuxFrame.MsgType, msg: Any = b''):
        ''' Called by child classes to send some message '''
        assert self._tx_queue is not None
        self._tx_queue.put((msg_type, msg))

    @abc.abstractmethod
    def _handle_frame(self, frame: Container) -> None:
        ''' Handle the given DebugMux frame '''


class ConnState(enum.Enum):
    ''' Connection state for DbgMuxConnHandler '''
    NotEstablished      = enum.auto()
    Establishing        = enum.auto()
    Established         = enum.auto()


class DbgMuxConnHandler(DbgMuxFrameHandler):
    ''' Abstract DebugMux connection handler '''

    def __init__(self):
        self.conn_state: ConnState = ConnState.NotEstablished
        self.ConnRef: int = 0xffff
        self.DPRef: int = 0xffff

    def establish(self, DPRef: int) -> None:
        assert self.conn_state == ConnState.NotEstablished
        log.info("Establishing connection with DPRef=0x%04x", DPRef)
        self.send(DbgMuxFrame.MsgType.ConnEstablish, dict(DPRef=DPRef))
        self.conn_state = ConnState.Establishing
        self.DPRef = DPRef

    def send_data(self, data: bytes) -> None:
        ''' Called by child classes to send connection data '''
        assert self.conn_state == ConnState.Established
        msg = dict(ConnRef=self.ConnRef, Data=data)
        self.send(DbgMuxFrame.MsgType.ConnData, msg)

    def _match(self, frame: Container,
                     msg_type: DbgMuxFrame.MsgType,
                     msg_fields: dict = { }) -> bool:
        if frame['MsgType'] != msg_type:
            return False
        for (key, val) in msg_fields.items():
            if key not in frame['Msg']:
                return False
            if frame['Msg'][key] != val:
                return False
        return True

    def _handle_frame(self, frame: Container) -> None:
        ''' Handle the given DebugMux frame '''

        if self.conn_state == ConnState.Established:
            fields = dict(ConnRef=self.ConnRef)
            if self._match(frame, DbgMuxFrame.MsgType.ConnData, fields):
                self._handle_data(frame['Msg']['Data'])
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration
            elif self._match(frame, DbgMuxFrame.MsgType.ConnTerminated, fields):
                log.info('Connection terminated (DPRef=0x%04x, ConnRef=0x%04x)',
                         self.DPRef, self.ConnRef)
                self.conn_state = ConnState.NotEstablished
                self._handle_terminate()
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration
        elif self.conn_state == ConnState.Establishing:
            # Match ConnEstablished with our DPRef and any ConnRef
            if self._match(frame, DbgMuxFrame.MsgType.ConnEstablished, dict(DPRef=self.DPRef)):
                log.info('Connection established (DPRef=0x%04x, ConnRef=0x%04x)',
                         self.DPRef, frame['Msg']['ConnRef'])
                self.conn_state = ConnState.Established
                self.ConnRef = frame['Msg']['ConnRef']
                self._handle_establish()
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration

    @abc.abstractmethod
    def _handle_data(self, data: bytes) -> None:
        ''' Called on reciept of connection data '''

    @abc.abstractmethod
    def _handle_establish(self) -> None:
        ''' Called on connection establishment '''

    @abc.abstractmethod
    def _handle_terminate(self) -> None:
        ''' Called on connection termination '''


class DbgMuxConnUdpBridge(DbgMuxConnHandler):
    DGRAM_MAX_LEN: int = 1024

    def __init__(self, *args, **kw):
        super().__init__(*args)

        self.raddr: str = kw.get('raddr', '127.0.0.1')
        self.rport: int = kw.get('rport', 9999)
        self.laddr: str = kw.get('laddr', '127.0.0.1')
        self.lport: int = kw.get('lport', 8888)

        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._sock.bind((self.laddr, self.lport))
        self._sock.settimeout(0.5)

        self._thread = threading.Thread(target=self._worker,
                                        daemon=True)
        self._shutdown = threading.Event()

    def _worker(self) -> None:
        while not self._shutdown.is_set():
            try:
                (data, addr) = self._sock.recvfrom(self.DGRAM_MAX_LEN)
                if self.conn_state == ConnState.Established:
                    self.send_data(data)
            except TimeoutError:
                pass

    def _handle_data(self, data: bytes) -> None:
        self._sock.sendto(data, (self.raddr, self.rport))

    def _handle_establish(self) -> None:
        self._shutdown.clear()
        self._thread.start()

    def _handle_terminate(self) -> None:
        self._shutdown.set()