From 9714cddd40c83894e619d649490706011e1856bc Mon Sep 17 00:00:00 2001
From: David Schultz <me@zpld.me>
Date: Thu, 24 Mar 2022 10:43:56 -0500
Subject: [PATCH] prune `_cliconns` on netsplit

---
 beryllia/__init__.py | 24 ++++++++++++++++++++----
 beryllia/util.py     | 18 ++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/beryllia/__init__.py b/beryllia/__init__.py
index 0c33c1c..c4a7e1c 100644
--- a/beryllia/__init__.py
+++ b/beryllia/__init__.py
@@ -19,7 +19,7 @@
 from .database  import Database
 from .normalise import RFC1459SearchNormaliser
 
-from .util      import oper_up, pretty_delta, get_statsp, get_klines
+from .util      import oper_up, pretty_delta, get_statsp, get_klines, get_links
 from .util      import try_parse_cidr, try_parse_ip, try_parse_ts
 from .util      import looks_like_glob
 
@@ -30,6 +30,7 @@
 RE_KLINEEXIT = re.compile(r"^\*{3} Notice -- (?:KLINE active for|Disconnecting K-Lined user) (?P<nickname>\S+)\[[^]]+\] .(?P<mask>\S+).$")
 RE_KLINEREJ  = re.compile(r"^\*{3} Notice -- Rejecting K-Lined user (?P<nick>\S+)\[(?P<user>[^]@]+)@(?P<host>[^]]+)\] .(?P<ip>\S+). .(?P<mask>\S+).$")
 RE_NICKCHG   = re.compile(r"^\*{3} Notice -- Nick change: From (?P<old_nick>\S+) to (?P<new_nick>\S+) .(?P<userhost>\S+).$")
+RE_NETSPLIT  = re.compile(r"^\*{3} Notice -- Netsplit (?P<hub>\S+) <-> (?P<leaf>\S+)")
 RE_DATE      = re.compile(r"^(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2})$")
 
 RE_KLINETAG  = re.compile(r"%(\S+)")
@@ -64,7 +65,7 @@ def __init__(self,
         self._database_init: bool = False
         self._wait_for_exit: Dict[str, str] = {}
 
-        self._cliconns: Dict[str, int] = {}
+        self._cliconns: Dict[str, Tuple[int, str]] = {} # nickname: (cliconn_id, server)
 
     def set_throttle(self, rate: int, time: float):
         # turn off throttling
@@ -105,6 +106,15 @@ async def _knag(self, oper: str, nick: str, kline_id: int):
             )
             await self.send(build("NOTICE", [nick, out]))
 
+    async def _prune_cliconns(self):
+        servers = await get_links(self)
+
+        # copy the dict so we don't get a RuntimeError
+        # due to modifying the dict during iteration
+        for nick in self._cliconns.copy().keys():
+            if not self._cliconns[nick][1] in servers:
+                del self._cliconns[nick]
+
     async def line_read(self, line: Line):
         now = time.monotonic()
 
@@ -150,6 +160,7 @@ async def line_read(self, line: Line):
             p_klineexit = RE_KLINEEXIT.search(message)
             p_klinerej  = RE_KLINEREJ.search(message)
             p_nickchg   = RE_NICKCHG.search(message)
+            p_netsplit  = RE_NETSPLIT.search(message)
 
             if p_cliconn is not None:
                 nickname = p_cliconn.group("nick")
@@ -174,14 +185,14 @@ async def line_read(self, line: Line):
                     ip,
                     line.source
                 )
-                self._cliconns[nickname] = cliconn_id
+                self._cliconns[nickname] = (cliconn_id, line.source)
 
             elif p_nickchg is not None:
                 old_nick = p_nickchg.group("old_nick")
                 new_nick = p_nickchg.group("new_nick")
                 if old_nick in self._cliconns:
                     cliconn_id = self._cliconns.pop(old_nick)
-                    self._cliconns[new_nick] = cliconn_id
+                    self._cliconns[new_nick] = (cliconn_id, line.source)
                     await self.database.nick_change.add(cliconn_id, new_nick)
 
             elif p_cliexit is not None:
@@ -278,6 +289,11 @@ async def line_read(self, line: Line):
                             kline_id, nickname, username, hostname, ip
                         )
 
+            elif p_netsplit is not None:
+                # prune people from self._cliconns
+                # when they're on a server that splits away
+                await self._prune_cliconns()
+
         elif (line.command == "PRIVMSG"
                 and line.source is not None
                 and not self.is_me(line.hostmask.nickname)):
diff --git a/beryllia/util.py b/beryllia/util.py
index 48decb4..ce960d1 100644
--- a/beryllia/util.py
+++ b/beryllia/util.py
@@ -16,6 +16,8 @@
 # not in ircstates.numerics
 RPL_STATS      = "249"
 RPL_ENDOFSTATS = "219"
+RPL_LINKS      = "364"
+RPL_ENDOFLINKS = "365"
 
 RE_OPERNAME = re.compile(r"^is opered as (\S+)(?:,|$)")
 
@@ -109,6 +111,22 @@ async def get_statsp(server: Server) -> List[Tuple[str, str]]:
 
     return opers
 
+async def get_links(server: Server) -> List[str]:
+    # :ircd.hub 364 l ircd.leaf ircd.hub :0 test server
+    await server.send(build("LINKS"))
+
+    linked_servers: List[str] = []
+    while True:
+        links_line = await server.wait_for({
+            Response(RPL_LINKS, [SELF, ANY]),
+            Response(RPL_ENDOFLINKS, [SELF, ANY])
+        })
+        if links_line.command == RPL_LINKS:
+            linked_servers.append(links_line.params[1])
+        else:
+            break
+    return linked_servers
+
 RPL_STATSKLINE = "216"
 RPL_ENDOFSTATS = "219"
 ERR_NOPRIVS    = "723"