Merge pull request CloudBotIRC#83 from linuxdaemon/gonzobot+chan-track-cleanup

Fix race conditions in chan_track.py

Author: edwardslabs

plugins/chan_track.py

@@ -4,26 +4,21 @@
 Requires:
 server_info.py
 """
-import re
 import weakref
 from collections import defaultdict
 from operator import attrgetter
 from weakref import WeakValueDictionary
 
 from cloudbot import hook
-
-NUH_RE = re.compile(r'(?P<nick>.+?)(?:!(?P<user>.+?))?(?:@(?P<host>.+?))?')
+from cloudbot.util.parsers.irc import Prefix
 
 
 class WeakDict(dict):
     # Subclass dict to allow it to be weakly referenced
     pass
 
 
-class DataDict(defaultdict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
+class KeyFoldMixin:
     def __contains__(self, item):
         return super().__contains__(item.casefold())
 
@@ -36,8 +31,18 @@ def __setitem__(self, key, value):
     def __delitem__(self, key):
         super().__delitem__(key.casefold())
 
+
+class KeyFoldDict(dict, KeyFoldMixin):
+    pass
+
+
+class KeyFoldWeakValueDict(WeakValueDictionary, KeyFoldMixin):
+    pass
+
+
+class ChanDict(defaultdict, KeyFoldMixin):
     def __missing__(self, key):
-        data = WeakDict({"name": key, "users": {}})
+        data = WeakDict(name=key, users=KeyFoldDict())
         self[key] = data
         return data
 
@@ -55,7 +60,7 @@ def update_conn_data(conn):
 
 @hook.on_cap_available("userhost-in-names", "multi-prefix")
 def do_caps():
-    pass
+    return True
 
 
 @hook.on_start
@@ -67,36 +72,30 @@ def get_chan_data(bot):
 
 @hook.connect
 def init_chan_data(conn):
-    chan_data = conn.memory.setdefault("chan_data", DataDict())
+    chan_data = conn.memory.setdefault("chan_data", ChanDict())
     chan_data.clear()
 
-    users = conn.memory.setdefault("users", WeakValueDictionary())
+    users = conn.memory.setdefault("users", KeyFoldWeakValueDict())
     users.clear()
 
 
-def parse_nuh(mask):
-    match = NUH_RE.fullmatch(mask)
-    if not match:
-        return None, None, None
-
-    nick = match.group('nick')
-    user = match.group('user')
-    host = match.group('host')
-    return nick, user, host
+def add_user_membership(user, chan, membership):
+    chans = user.setdefault("channels", KeyFoldWeakValueDict())
+    chans[chan] = membership
 
 
 def replace_user_data(conn, chan_data):
     statuses = {status.prefix: status for status in set(conn.memory["server_info"]["statuses"].values())}
     users = conn.memory["users"]
     old_users = chan_data['users']
     new_data = chan_data.pop("new_users", [])
-    new_users = {}
+    new_users = KeyFoldDict()
     caps = conn.memory.get("server_caps", {})
     has_uh_i_n = caps.get("userhost-in-names", False)
     has_multi_pfx = caps.get("multi-prefix", False)
     for name in new_data:
         user_data = WeakDict()
-        memb_data = WeakDict({"user": user_data, "chan": weakref.proxy(chan_data)})
+        memb_data = WeakDict(user=user_data, chan=weakref.proxy(chan_data))
         user_statuses = []
         while name[:1] in statuses:
             status, name = name[:1], name[1:]
@@ -111,24 +110,23 @@ def replace_user_data(conn, chan_data):
         memb_data["status"] = user_statuses
 
         if has_uh_i_n:
-            nick, user, host = parse_nuh(name)
-            user_data.update({"nick": nick, "ident": user, "host": host})
+            pfx = Prefix.parse(name)
+            user_data.update({"nick": pfx.nick, "ident": pfx.user, "host": pfx.host})
         else:
             user_data["nick"] = name
 
-        nick_cf = user_data["nick"].casefold()
-        new_users[nick_cf] = memb_data
-        if nick_cf in old_users:
-            old_data = old_users[nick_cf]
+        nick = user_data["nick"]
+        new_users[nick] = memb_data
+        if nick in old_users:
+            old_data = old_users[nick]
             old_data.update(memb_data)  # New data takes priority over old data
             memb_data.update(old_data)
 
-        old_user_data = users.setdefault(nick_cf, user_data)
+        old_user_data = users.setdefault(nick, user_data)
         old_user_data.update(user_data)
         user_data = old_user_data
         memb_data["user"] = user_data
-        user_chans = user_data.setdefault("channels", WeakValueDictionary())
-        user_chans[chan_data["name"].casefold()] = memb_data
+        add_user_membership(user_data, chan_data["name"], memb_data)
 
     old_users.clear()
     old_users.update(new_users)  # Reassigning the dict would break other references to the data, so just update instead
@@ -181,11 +179,10 @@ def perm_check(chan, conn, nick):
         return False
 
     chan_data = chans[chan]
-    nick_cf = nick.casefold()
-    if nick_cf not in chan_data["users"]:
+    if nick not in chan_data["users"]:
         return False
 
-    memb = chan_data["users"][nick_cf]
+    memb = chan_data["users"][nick]
     status = memb["status"]
     if status and status[0].level > 1:
         return True
@@ -218,58 +215,31 @@ def updateusers(bot):
     return "Updating all channel data"
 
 
-@hook.irc_raw('JOIN')
+@hook.irc_raw(['JOIN', 'MODE'], singlethread=True)
+def on_join_mode(chan, nick, user, host, conn, irc_command, irc_paramlist):
+    """
+    Both JOIN and MODE are handled in one hook with Hook:singlethread=True
+    to ensure they are handled in order, avoiding a possible race condition
+    """
+    if irc_command == 'JOIN':
+        return on_join(chan, nick, user, host, conn)
+    elif irc_command == 'MODE':
+        return on_mode(chan, irc_paramlist, conn)
+
+
 def on_join(chan, nick, user, host, conn):
     if chan.startswith(':'):
         chan = chan[1:]
 
     users = conn.memory['users']
-    user_data = WeakDict({"nick": nick, "user": user, "host": host})
-    user_data = users.setdefault(nick.casefold(), user_data)
+    user_data = WeakDict(nick=nick, user=user, host=host)
+    user_data = users.setdefault(nick, user_data)
     chan_data = conn.memory["chan_data"][chan]
-    memb_data = WeakDict({"chan": weakref.proxy(chan_data), "user": user_data, "status": []})
-    chan_data["users"][nick.casefold()] = memb_data
-    user_chans = user_data.setdefault("channels", WeakValueDictionary())
-    user_chans[chan.casefold()] = memb_data
-
-
-@hook.irc_raw('PART')
-def on_part(chan, nick, conn):
-    if chan.startswith(':'):
-        chan = chan[1:]
-
-    channels = conn.memory["chan_data"]
-    nick_cf = nick.casefold()
-    if nick_cf == conn.nick.casefold():
-        try:
-            del channels[chan]
-        except KeyError:
-            pass
-    else:
-        chan_data = channels[chan]
-        try:
-            del chan_data["users"][nick_cf]
-        except KeyError:
-            pass
+    memb_data = WeakDict(chan=weakref.proxy(chan_data), user=user_data, status=[])
+    chan_data["users"][nick] = memb_data
+    add_user_membership(user_data, chan, memb_data)
 
 
-@hook.irc_raw('KICK')
-def on_kick(chan, target, conn):
-    on_part(chan, target, conn)
-
-
-@hook.irc_raw('QUIT')
-def on_quit(nick, conn):
-    nick_cf = nick.casefold()
-    users = conn.memory["users"]
-    if nick_cf in users:
-        user = users[nick_cf]
-        for memb in user.get("channels", {}).values():
-            chan = memb["chan"]
-            chan["users"].pop(nick_cf)
-
-
-@hook.irc_raw('MODE')
 def on_mode(chan, irc_paramlist, conn):
     if chan.startswith(':'):
         chan = chan[1:]
@@ -306,7 +276,7 @@ def on_mode(chan, irc_paramlist, conn):
                 param = mode_params.pop(0)
 
                 if is_status:
-                    memb = chan_data["users"][param.casefold()]
+                    memb = chan_data["users"][param]
                     status = statuses[c]
                     if adding:
                         memb["status"].append(status)
@@ -316,18 +286,50 @@ def on_mode(chan, irc_paramlist, conn):
                             memb["status"].remove(status)
 
 
+@hook.irc_raw('PART')
+def on_part(chan, nick, conn):
+    if chan.startswith(':'):
+        chan = chan[1:]
+
+    channels = conn.memory["chan_data"]
+    if nick.casefold() == conn.nick.casefold():
+        try:
+            del channels[chan]
+        except KeyError:
+            pass
+    else:
+        chan_data = channels[chan]
+        try:
+            del chan_data["users"][nick]
+        except KeyError:
+            pass
+
+
+@hook.irc_raw('KICK')
+def on_kick(chan, target, conn):
+    on_part(chan, target, conn)
+
+
+@hook.irc_raw('QUIT')
+def on_quit(nick, conn):
+    users = conn.memory["users"]
+    if nick in users:
+        user = users[nick]
+        for memb in user.get("channels", {}).values():
+            chan = memb["chan"]
+            chan["users"].pop(nick)
+
+
 @hook.irc_raw('NICK')
 def on_nick(nick, irc_paramlist, conn):
     users = conn.memory["users"]
-    nick_cf = nick.casefold()
     new_nick = irc_paramlist[0]
     if new_nick.startswith(':'):
         new_nick = new_nick[1:]
 
-    new_nick_cf = new_nick.casefold()
-    user = users.pop(nick_cf)
-    users[new_nick_cf] = user
+    user = users.pop(nick)
+    users[new_nick] = user
     user["nick"] = nick
-    for memb in user["channels"].values():
+    for memb in user.get("channels", {}).values():
         chan_users = memb["chan"]["users"]
-        chan_users[new_nick_cf] = chan_users.pop(nick_cf)
+        chan_users[new_nick] = chan_users.pop(nick)