imaplib: preserve partial reads on exception

This ensures that short IDLE durations / burst() intervals
won't risk corrupting response lines that span multiple packets.

Author: Forest

Lib/imaplib.py

@@ -201,7 +201,7 @@ def __init__(self, host='', port=IMAP4_PORT, timeout=None):
         self.tagnum = 0
         self._tls_established = False
         self._mode_ascii()
-        self._readbuf = b''
+        self._readbuf = []
 
         # Open socket to server.
 
@@ -346,16 +346,24 @@ def read(self, size):
         # https://github.com/python/cpython/issues/51571
 
         parts = []
-        while True:
-            if len(self._readbuf) >= size:
-                parts.append(self._readbuf[:size])
-                self._readbuf = self._readbuf[size:]
-                break
-            parts.append(self._readbuf)
-            size -= len(self._readbuf)
-            self._readbuf = self.sock.recv(DEFAULT_BUFFER_SIZE)
-            if not self._readbuf:
+
+        while size > 0:
+
+            if len(parts) < len(self._readbuf):
+                buf = self._readbuf[len(parts)]
+            else:
+                buf = self.sock.recv(DEFAULT_BUFFER_SIZE)
+                if not buf:
+                    break
+                self._readbuf.append(buf)
+
+            if len(buf) >= size:
+                parts.append(buf[:size])
+                self._readbuf = [buf[size:]]
                 break
+            parts.append(buf)
+            size -= len(buf)
+
         return b''.join(parts)
 
 
@@ -366,18 +374,25 @@ def readline(self):
         LF = b'\n'
         parts = []
         length = 0
+
         while length < _MAXLINE:
-            pos = self._readbuf.find(LF)
+
+            if len(parts) < len(self._readbuf):
+                buf = self._readbuf[len(parts)]
+            else:
+                buf = self.sock.recv(DEFAULT_BUFFER_SIZE)
+                if not buf:
+                    break
+                self._readbuf.append(buf)
+
+            pos = buf.find(LF)
             if pos != -1:
                 pos += 1
-                parts.append(self._readbuf[:pos])
-                self._readbuf = self._readbuf[pos:]
-                break
-            parts.append(self._readbuf)
-            length += len(parts[-1])
-            self._readbuf = self.sock.recv(DEFAULT_BUFFER_SIZE)
-            if not self._readbuf:
+                parts.append(buf[:pos])
+                self._readbuf = [buf[pos:]]
                 break
+            parts.append(buf)
+            length += len(buf)
 
         line = b''.join(parts)
         if len(line) > _MAXLINE:

Lib/test/test_imaplib.py

@@ -239,6 +239,23 @@ def cmd_IDLE(self, tag, args):
             self._send_tagged(tag, 'BAD', 'Expected DONE')
 
 
+class IdleCmdDelayedPacketHandler(SimpleIMAPHandler):
+    capabilities = 'IDLE'
+    def cmd_IDLE(self, tag, args):
+        self._send_textline('+ idling')
+        # response line spanning multiple packets, the last one delayed
+        self._send(b'* 1 EX')
+        time.sleep(0.2)
+        self._send(b'IS')
+        time.sleep(1)
+        self._send(b'TS\r\n')
+        r = yield
+        if r == b'DONE\r\n':
+            self._send_tagged(tag, 'OK', 'Idle completed')
+        else:
+            self._send_tagged(tag, 'BAD', 'Expected DONE')
+
+
 class NewIMAPTestsMixin():
     client = None
 
@@ -583,6 +600,18 @@ def test_idle_burst(self):
         _, data = client.response('RECENT')
         self.assertEqual(data, [b'1', b'9'])
 
+    def test_idle_delayed_packet(self):
+        client, _ = self._setup(IdleCmdDelayedPacketHandler)
+        client.login('user', 'pass')
+        # If our readline() implementation fails to preserve line fragments
+        # when idle timeouts trigger, a response spanning delayed packets
+        # can be corrupted, leaving the protocol stream in a bad state.
+        try:
+            with client.idle(0.5) as idler:
+                self.assertRaises(StopIteration, next, idler)
+        except client.abort as err:
+            self.fail('multi-packet response was corrupted by idle timeout')
+
     def test_login(self):
         client, _ = self._setup(SimpleIMAPHandler)
         typ, data = client.login('user', 'pass')