EchoServer

server.py

from typing import Tuple

from loguru import logger
from twisted.internet import reactor
from twisted.internet.posixbase import PosixReactorBase
from twisted.internet.protocol import DatagramProtocol

reactor: PosixReactorBase = reactor


class EchoServer:
    max_dgram_size = 2 * 508

    # 508 bytes is maximum "safe" datagram size over the internet

    def __init__(self):
        logger.debug('Initializing Server...')
        self._accepted_clients = dict()

        class EchoProtocol(DatagramProtocol):
            server = self
            handshake_msg = 'Hola'
            handshake_resp = 'OK'

            def datagramReceived(self,
                                 datagram: bytes,
                                 addr: Tuple[str, int]):
                logger.debug(f'Got message from {addr}')
                try:
                    # first check if message matches handshake.
                    # if it does, reset the client
                    msg = datagram.decode('utf8')
                    assert msg == self.handshake_msg

                    logger.debug('Handshake OK')
                    # accept client, but don't set max size yet
                    self.server._accepted_clients[addr] = None

                    logger.debug('Sending OK to client.')
                    self.transport.write(self.handshake_resp.encode('utf8'),
                                         addr)
                    return
                except (UnicodeError, AssertionError):
                    # could not decode handshake or message didn't match
                    # expected handshake.
                    # Is the client already registered?
                    try:
                        max_size = self.server._accepted_clients[addr]
                    except KeyError:
                        logger.debug('Could not decode handshake and client '
                                     'is not accepted, dropping!')
                        return

                if max_size is None:
                    logger.debug('Client needs to set max size.')
                    # client was accepted but max size has not been
                    # agreed upon.
                    # next message should correspond to a string
                    # specifying the maximum size suggested by the client
                    try:
                        size_msg = datagram.decode('utf8')
                        client_sz = int(size_msg)

                        logger.debug(f'Got max size {client_sz} bytes '
                                     f'suggestion from client.')

                        selected_sz = min(client_sz, self.server.max_dgram_size)

                        logger.debug(f'Selected size {selected_sz} bytes.')
                        self.server._accepted_clients[addr] = selected_sz

                        logger.debug('Sending selected size to client.')
                        self.transport.write(f'{selected_sz:d}'.encode('utf8'),
                                             addr)
                    except (UnicodeError, ValueError) as _:
                        # could not decode size message, drop the client
                        logger.debug(f'Could not decode size message, '
                                     f'dropping {addr}.')
                        self.server._accepted_clients.pop(addr)
                else:
                    # client is accepted, size has been agreed upon
                    dgram_len = len(datagram)
                    if dgram_len > max_size:
                        # client violated size agreement, drop it!
                        logger.debug(f'Client violated maximum size of '
                                     f'{max_size} bytes (sent {dgram_len} '
                                     f'bytes. Dropping {addr}.')
                        self.server._accepted_clients.pop(addr)
                    else:
                        # write back whatever client sent
                        logger.debug(f'Echoing {dgram_len} bytes of data '
                                     f'back to {addr}.')
                        self.transport.write(datagram, addr)

        self._protocol = EchoProtocol

    def serve(self, port: int) -> None:
        reactor.listenUDP(port, self._protocol())
        reactor.run()


if __name__ == '__main__':
    EchoServer().serve(5000)