diff --git a/tests/chatroom_server.py b/tests/chatroom_server.py new file mode 100644 index 0000000..cc8096d --- /dev/null +++ b/tests/chatroom_server.py @@ -0,0 +1,93 @@ +import aiosched +import logging +import sys + +# An asynchronous chatroom + +clients: dict[aiosched.socket.AsyncSocket, list[str, str]] = {} +names: set[str] = set() + + +async def serve(bind_address: tuple): + """ + Serves asynchronously forever + + :param bind_address: The address to bind the server to represented as a tuple + (address, port) where address is a string and port is an integer + """ + + sock = aiosched.socket.socket() + await sock.bind(bind_address) + await sock.listen(5) + logging.info(f"Serving asynchronously at {bind_address[0]}:{bind_address[1]}") + async with aiosched.with_context() as ctx: + async with sock: + while True: + try: + conn, address_tuple = await sock.accept() + clients[conn] = ["", f"{address_tuple[0]}:{address_tuple[1]}"] + logging.info(f"{address_tuple[0]}:{address_tuple[1]} connected") + await ctx.spawn(handler, conn) + except Exception as err: + # Because exceptions just *work* + logging.info(f"{address_tuple[0]}:{address_tuple[1]} has raised {type(err).__name__}: {err}") + + +async def handler(sock: aiosched.socket.AsyncSocket): + """ + Handles a single client connection + + :param sock: The AsyncSocket object connected to the client + """ + + address = clients[sock][1] + name = "" + async with sock: # Closes the socket automatically + await sock.send_all(b"Welcome to the chatroom pal, may you tell me your name?\n> ") + while True: + while not name.endswith("\n"): + name = (await sock.receive(64)).decode() + name = name[:-1] + if name not in names: + names.add(name) + clients[sock][0] = name + break + else: + await sock.send_all(b"Sorry, but that name is already taken. Try again!\n> ") + await sock.send_all(f"Okay {name}, welcome to the chatroom!\n".encode()) + logging.info(f"{name} has joined the chatroom ({address}), informing clients") + for i, client_sock in enumerate(clients): + if client_sock != sock and clients[client_sock][0]: + await client_sock.send_all(f"{name} joins the chatroom!\n> ".encode()) + while True: + await sock.send_all(b"> ") + data = await sock.receive(1024) + if not data: + break + logging.info(f"Got: {data!r} from {address}") + for i, client_sock in enumerate(clients): + if client_sock != sock and clients[client_sock][0]: + logging.info(f"Sending {data!r} to {':'.join(map(str, await client_sock.getpeername()))}") + if not data.endswith(b"\n"): + data += b"\n" + await client_sock.send_all(f"[{name}] ({address}): {data.decode()}> ".encode()) + logging.info(f"Sent {data!r} to {i} clients") + logging.info(f"Connection from {address} closed") + clients.pop(sock) + names.discard(name) + + +if __name__ == "__main__": + port = int(sys.argv[1]) if len(sys.argv) > 1 else 1501 + logging.basicConfig( + level=20, + format="[%(levelname)s] %(asctime)s %(message)s", + datefmt="%d/%m/%Y %p", + ) + try: + aiosched.run(serve, ("localhost", port)) + except (Exception, KeyboardInterrupt) as error: # Exceptions propagate! + if isinstance(error, KeyboardInterrupt): + logging.info("Ctrl+C detected, exiting") + else: + logging.error(f"Exiting due to a {type(error).__name__}: {error}")