Add support for QUIT command

.env

@@ -0,0 +1 @@
+USER_MAX_CHANNELS=50

paircd/client.py

@@ -1,10 +1,12 @@
-from asyncio import StreamReader, StreamWriter, Queue
+from asyncio import StreamReader, StreamWriter, TimeoutError, Queue
+from asyncio.tasks import wait_for
 from dataclasses import dataclass, field
 from datetime import datetime
 from logging import log, INFO
-from paircd.reply import RPL_CREATED, RPL_MYINFO, RPL_WELCOME, RPL_YOURHOST
+from typing import Any, Set
+
+from paircd.reply import QUIT, RPL_CREATED, RPL_MYINFO, RPL_WELCOME, RPL_YOURHOST
 from paircd.message import Message
-from typing import Set
 
 
 @dataclass
@@ -12,6 +14,7 @@ class Client:
     hostname: str
     reader: StreamReader
     writer: StreamWriter
+
     msg_queue: Queue = field(default_factory=Queue)
 
     nickname: str = ""
@@ -19,6 +22,7 @@ class Client:
     realname: str = ""
     registered: bool = False
     away: str = ""
+    closed = False
 
     modes: Set[str] = field(default_factory=set)
     channels: Set[str] = field(default_factory=set)
@@ -31,11 +35,43 @@ class Client:
     def log(self, msg: str, level: int = INFO) -> None:
         log(level, f"{self.hostname} ({self.id()}) {msg}")
 
-    async def write_forever(self) -> None:
-        while True:
-            msg = await self.msg_queue.get()
-            self.writer.write(msg)
-            await self.writer.drain()
+    async def write_until_closed(self, server: Any) -> None:
+        while not self.closed:
+            msg = None
+            try:
+                msg = await wait_for(self.msg_queue.get(), timeout=1.0)
+            except TimeoutError:
+                pass
+
+            if msg is not None:
+                self.writer.write(msg)
+                try:
+                    await self.writer.drain()
+                except ConnectionResetError:
+                    await self.quit(server, "Connection reset by peer")
+
+    async def quit(self, server: Any, msg: str) -> None:
+        if self.closed:
+            return
+
+        quit_msg = QUIT(msg, prefix=self.id())
+        for client in server.clients_by_nick.values():
+            client.write_message(quit_msg)
+        for channel_name in self.channels:
+            channel = server.get_channel_by_name(channel_name, create=False)
+            if channel is None:
+                continue
+            channel.remove_client_by_nick(self.nickname)
+        server.remove_client_by_name(self.nickname)
+        await self.close()
+
+    async def close(self) -> None:
+        if self.closed:
+            return
+        self.closed = True
+        if not self.writer.is_closing():
+            self.writer.close()
+        await self.writer.wait_closed()
 
     def write_message(self, message: Message) -> None:
         self.msg_queue.put_nowait(message.encode())

paircd/handler/join.py

@@ -1,5 +1,6 @@
 from paircd.reply import (
     ERR_NOSUCHCHANNEL,
+    ERR_TOOMANYCHANNELS,
     JOIN,
     RPL_ENDOFNAMES,
     RPL_NAMREPLY,
@@ -11,6 +12,7 @@ from paircd.client import Client
 from paircd.command_handler import CommandHandler
 from paircd.message import Message
 from paircd.server import Server
+from paircd.settings import USER_MAX_CHANNELS
 
 
 class JoinHandler(CommandHandler):
@@ -20,6 +22,10 @@ class JoinHandler(CommandHandler):
     async def handle(self, server: Server, client: Client, msg: Message) -> None:
         channel_name = msg.args[0]
 
+        if USER_MAX_CHANNELS > 0 and len(client.channels) == USER_MAX_CHANNELS:
+            client.write_message(ERR_TOOMANYCHANNELS(client.nickname))
+            return
+
         channel = server.get_channel_by_name(channel_name)
         if channel is None:
             client.write_message(ERR_NOSUCHCHANNEL(client.nickname, channel_name))

paircd/handler/quit.py

@@ -0,0 +1,15 @@
+from paircd.client import Client
+from paircd.command_handler import CommandHandler
+from paircd.message import Message
+from paircd.server import Server
+
+
+class QuitHandler(CommandHandler):
+    def __init__(self) -> None:
+        super().__init__("QUIT")
+
+    async def handle(self, server: Server, client: Client, msg: Message) -> None:
+        quit_msg = ""
+        if len(msg.args) == 1:
+            quit_msg = msg.args[0]
+        await client.quit(server, f"Quit: {quit_msg}")

paircd/handlers.py

@@ -15,6 +15,8 @@ from paircd.handler.nick import NickHandler
 from paircd.handler.notice import NoticeHandler
 from paircd.handler.ping import PingHandler
 from paircd.handler.privmsg import PrivmsgHandler
+
+from paircd.handler.quit import QuitHandler
 from paircd.handler.topic import TopicHandler
 from paircd.handler.user import UserHandler
 from paircd.handler.who import WhoHandler
@@ -27,6 +29,7 @@ HANDLER_CLASSES = [
     NoticeHandler,
     PingHandler,
     PrivmsgHandler,
+    QuitHandler,
     TopicHandler,
     UserHandler,
     WhoHandler,

paircd/main.py

@@ -1,19 +1,31 @@
-from asyncio import create_task, run, start_server
+from asyncio import create_task, run, start_server, IncompleteReadError
 from asyncio.streams import StreamReader, StreamWriter
 from logging import basicConfig, info, INFO
+import logging
 from os import getenv
+from sys import exc_info
 
-from paircd.client import Client
-from paircd.handlers import handle_cmd, register_cmd_handlers
-from paircd.message import parse_message
-from paircd.server import Server
+from dotenv import load_dotenv
+
+load_dotenv("../.env")
+
+from paircd.client import Client  # noqa: E402
+from paircd.handlers import handle_cmd, register_cmd_handlers  # noqa: E402
+from paircd.message import parse_message  # noqa: E402
+from paircd.server import Server  # noqa: E402
 
 basicConfig(format="%(asctime)s [%(levelname)s] - %(message)s", level=INFO)
 
 
 async def read_forever(server: Server, client: Client) -> None:
-    while True:
-        raw_msg = await client.reader.readuntil(b"\r\n")
+    while not client.closed:
+        try:
+            raw_msg = await client.reader.readuntil(b"\r\n")
+        except IncompleteReadError:
+            logging.warning("client connection closed", exc_info=exc_info())
+            await client.quit(server, "Connection reset by peer")
+            break
+
         msg = parse_message(raw_msg)
         await handle_cmd(server, client, msg)
 
@@ -32,7 +44,7 @@ async def serve() -> None:
             writer=writer,
         )
         create_task(read_forever(irc_server, client))
-        create_task(client.write_forever())
+        create_task(client.write_until_closed(irc_server))
 
     server = await start_server(
         register_client,

paircd/reply.py

@@ -30,6 +30,7 @@ NICK = cmd_fn("NICK", "{0}")
 NOTICE = cmd_fn("NOTICE", "{0} :{1}")
 PONG = cmd_fn("PONG", ":{0}")
 PRIVMSG = cmd_fn("PRIVMSG", "{0} :{1}")
+QUIT = cmd_fn("QUIT", ":{0}")
 TOPIC = cmd_fn("TOPIC", "{0} :{1}")
 
 # Error replies

paircd/settings.py

@@ -0,0 +1,3 @@
+import os
+
+USER_MAX_CHANNELS = int(os.getenv("USER_MAX_CHANNELS") or "0")

poetry.lock

@@ -188,6 +188,17 @@ wcwidth = "*"
 checkqa-mypy = ["mypy (==v0.761)"]
 testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "requests", "xmlschema"]
 
+[[package]]
+name = "python-dotenv"
+version = "0.18.0"
+description = "Read key-value pairs from a .env file and set them as environment variables"
+category = "main"
+optional = false
+python-versions = "*"
+
+[package.extras]
+cli = ["click (>=5.0)"]
+
 [[package]]
 name = "pyyaml"
 version = "5.4.1"
@@ -249,7 +260,7 @@ python-versions = "*"
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.8"
-content-hash = "298291104b6892022cb47c34680f925bc85fff59e768a225230ef84444acb002"
+content-hash = "16dd6949f4ee094d7e2d9ddf66b05ac484047dae65acbdd61c29e8a31862624c"
 
 [metadata.files]
 appdirs = [
@@ -345,6 +356,10 @@ pytest = [
     {file = "pytest-5.4.3-py3-none-any.whl", hash = "sha256:5c0db86b698e8f170ba4582a492248919255fcd4c79b1ee64ace34301fb589a1"},
     {file = "pytest-5.4.3.tar.gz", hash = "sha256:7979331bfcba207414f5e1263b5a0f8f521d0f457318836a7355531ed1a4c7d8"},
 ]
+python-dotenv = [
+    {file = "python-dotenv-0.18.0.tar.gz", hash = "sha256:effaac3c1e58d89b3ccb4d04a40dc7ad6e0275fda25fd75ae9d323e2465e202d"},
+    {file = "python_dotenv-0.18.0-py2.py3-none-any.whl", hash = "sha256:dd8fe852847f4fbfadabf6183ddd4c824a9651f02d51714fa075c95561959c7d"},
+]
 pyyaml = [
     {file = "PyYAML-5.4.1-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:3b2b1824fe7112845700f815ff6a489360226a5609b96ec2190a45e62a9fc922"},
     {file = "PyYAML-5.4.1-cp27-cp27m-win32.whl", hash = "sha256:129def1b7c1bf22faffd67b8f3724645203b79d8f4cc81f674654d9902cb4393"},

pyproject.toml

@@ -6,6 +6,7 @@ authors = ["Forest Belton <65484+forestbelton@users.noreply.github.com>"]
 
 [tool.poetry.dependencies]
 python = "^3.8"
+python-dotenv = "^0.18.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^5.2"