From 65ab65cbb1b85562ae27cc213b2eb901e5b5eecf Mon Sep 17 00:00:00 2001
From: Valentin Lorentz <progval+git@progval.net>
Date: Sun, 20 Jun 2021 23:59:51 +0200
Subject: [PATCH] irclib: Fix crashes when ERROR is part of a batch.

---
 src/irclib.py       | 21 ++++++++++++++++++---
 test/test_irclib.py | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/src/irclib.py b/src/irclib.py
index 07d5427c7..34eba36ee 100644
--- a/src/irclib.py
+++ b/src/irclib.py
@@ -708,8 +708,11 @@ class IrcState(IrcCommandDispatcher, log.Firewalled):
         self.channels = channels
         self.nicksToHostmasks = nicksToHostmasks
 
-        # Batches should always finish and be way shorter than 3600s, but
-        # let's just make sure to avoid leaking memory.
+        # Batches usually finish and are way shorter than 3600s, but
+        # we need to:
+        # * keep them in case the connection breaks (and reset() can't
+        #   clear the list itself)
+        # * make sure to avoid leaking memory in general
         self.batches = ExpiringDict(timeout=3600)
 
     def reset(self):
@@ -721,12 +724,24 @@ class IrcState(IrcCommandDispatcher, log.Firewalled):
         self.channels.clear()
         self.supported.clear()
         self.nicksToHostmasks.clear()
-        self.batches.clear()
         self.capabilities_req = set()
         self.capabilities_ack = set()
         self.capabilities_nak = set()
         self.capabilities_ls = {}
 
+        # Don't clear batches right now. reset() is called on ERROR messages,
+        # which may be part of a BATCH so we need to remember that batch.
+        # At worst, the batch will expire in the near future, as self.batches
+        # is an instance of ExpiringDict.
+        # If we did clear the batch, then this would happen:
+        # 1. IrcState.addMsg() would crash on the ERROR, because its batch
+        #    server-tag references an unknown batch, so it would not set the
+        #    'batch' supybot-tag
+        # 2. Irc.doBatch would crash on the closing BATCH, for the same reason
+        # 3. Owner.doBatch would crash because it expects the batch
+        #    supybot-tag to be set, but it wasn't because of 1
+        #self.batches.clear()
+
     def __reduce__(self):
         return (self.__class__, (self.history, self.supported,
                                  self.nicksToHostmasks, self.channels))
diff --git a/test/test_irclib.py b/test/test_irclib.py
index efd173a82..6de1f9308 100644
--- a/test/test_irclib.py
+++ b/test/test_irclib.py
@@ -1042,8 +1042,6 @@ class IrcTestCase(SupyTestCase):
             repr(c.batch)
         )
 
-    maxDiff = None
-
     def testBatchNested(self):
         self.irc.reset()
         logs = textwrap.dedent('''
@@ -1107,6 +1105,42 @@ class IrcTestCase(SupyTestCase):
         self.assertEqual(msg5.tagged('batch'), outer)
         self.assertEqual(self.irc.state.getParentBatches(msg5), [outer])
 
+    def testBatchError(self):
+        # Checks getting an ERROR message in a batch does not cause issues
+        # due to deinitializing the connection at the same time.
+        self.irc.reset()
+        m1 = ircmsgs.IrcMsg(':host BATCH +name batchtype')
+        self.irc.feedMsg(m1)
+        m2 = ircmsgs.IrcMsg('@batch=name :someuser2 NOTICE * :oh no')
+        self.irc.feedMsg(m2)
+        m3 = ircmsgs.IrcMsg('@batch=name ERROR :bye')
+        self.irc.feedMsg(m3)
+        class Callback(irclib.IrcCallback):
+            batch = None
+            def __call__(self, *args, **kwargs):
+                return super().__call__(*args, **kwargs)
+            def name(self):
+                return 'testcallback'
+            def doBatch(self2, irc, msg):
+                self2.batch = msg.tagged('batch')
+
+        # would usually be called by the driver upon reconnect() trigged
+        # by the ERROR:
+        self.irc.reset()
+
+        c = Callback()
+        self.irc.addCallback(c)
+        try:
+            m4 = ircmsgs.IrcMsg(':host BATCH -name')
+            self.irc.feedMsg(m4)
+        finally:
+            self.irc.removeCallback(c.name())
+        self.assertEqual(
+            c.batch,
+            irclib.Batch('name', 'batchtype', (), [m1, m2, m3, m4], None),
+            repr(c.batch)
+        )
+
     def testTruncate(self):
         self.irc = irclib.Irc('test')