134 lines
5.2 KiB
Python
134 lines
5.2 KiB
Python
import aiosched
|
|
import logging
|
|
import sys
|
|
|
|
# An asynchronous chatroom
|
|
|
|
clients: dict[aiosched.socket.AsyncSocket, list[str, str]] = {}
|
|
names: set[str] = set()
|
|
|
|
|
|
async def message_handler(q: aiosched.Queue):
|
|
"""
|
|
Reads data submitted onto the queue
|
|
"""
|
|
|
|
try:
|
|
logging.info("Message handler spawned")
|
|
while True:
|
|
msg, payload = await q.get()
|
|
logging.info(f"Got message {msg!r} with the following payload: {payload}")
|
|
except Exception as e:
|
|
logging.error(f"An exception occurred in the message handler -> {type(e).__name__}: {e}")
|
|
|
|
|
|
async def serve(bind_address: tuple):
|
|
"""
|
|
Serves asynchronously forever (or until Ctrl+C ;))
|
|
|
|
: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()
|
|
queue = aiosched.Queue()
|
|
await sock.bind(bind_address)
|
|
await sock.listen(5)
|
|
logging.info(f"Serving asynchronously at {bind_address[0]}:{bind_address[1]}")
|
|
async with aiosched.create_pool() as pool:
|
|
await pool.spawn(message_handler, queue)
|
|
async with sock:
|
|
while True:
|
|
try:
|
|
conn, address_tuple = await sock.accept()
|
|
clients[conn] = ["", f"{address_tuple[0]}:{address_tuple[1]}"]
|
|
await queue.put(("connect", clients[conn]))
|
|
logging.info(f"{address_tuple[0]}:{address_tuple[1]} connected")
|
|
await pool.spawn(handler, conn, queue)
|
|
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, q: aiosched.Queue):
|
|
"""
|
|
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> ")
|
|
cond = True
|
|
while cond:
|
|
while not name.endswith("\n"):
|
|
name = (await sock.receive(64)).decode()
|
|
if name == "":
|
|
cond = False
|
|
break
|
|
name = name.rstrip("\n")
|
|
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())
|
|
await q.put(("join", (address, name)))
|
|
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
|
|
decoded = data.decode().rstrip("\n")
|
|
if decoded.startswith("/"):
|
|
logging.info(f"{name} issued server command {decoded}")
|
|
await q.put(("cmd", (name, decoded[1:])))
|
|
match decoded[1:]:
|
|
case "bye":
|
|
await sock.send_all(b"Bye!\n")
|
|
break
|
|
case _:
|
|
await sock.send_all(b"Unknown command\n")
|
|
else:
|
|
await q.put(("msg", (name, data)))
|
|
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")
|
|
await q.put(("leave", name))
|
|
logging.info(f"Connection from {address} closed")
|
|
logging.info(f"{name} has left 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} has left the chatroom\n> ".encode())
|
|
clients.pop(sock)
|
|
names.discard(name)
|
|
logging.info("Handler shutting down")
|
|
|
|
|
|
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, ("0.0.0.0", port), debugger=None)
|
|
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}")
|