Fix POP3/SMTP correctness and IMAP robustness (#10)

POP3/SMTP proxy correctness fixes.

- RETR now returns the message body (reads the FETCH tuple). Closes #1
- Maildrop snapshotted once; DELE marks survive STAT/LIST/UIDL. Closes #2
- RETR/TOP dot-stuff and CRLF-terminate per RFC 1939. Closes #3
- SMTP relays raw bytes, no UTF-8 mangling. Closes #4
- Message sizes batched in one IMAP round-trip off the event loop; IMAP timeout added. Closes #5
- POP3 TOP implemented. Closes #6

pytest: 10 passed.
Reviewed-on: #10
Co-authored-by: Lyra Thorpe <iam@emmathe.dev>
Co-committed-by: Lyra Thorpe <iam@emmathe.dev>

proxy_server.py

@@ -64,6 +64,41 @@ class Settings:
             raise RuntimeError(f"Missing required environment variables: {', '.join(missing)}")
 
 
+def _split_lines(data):
+    """Split a byte string into lines, normalising CR/CRLF/LF endings."""
+    text = data.replace(b"\r\n", b"\n").replace(b"\r", b"\n")
+    lines = text.split(b"\n")
+    if lines and lines[-1] == b"":
+        lines = lines[:-1]
+    return lines
+
+
+def format_pop3_body(body, max_body_lines=None):
+    """Prepare a message for POP3 RETR/TOP transmission (RFC 1939).
+
+    Normalises line endings to CRLF, byte-stuffs any line beginning with ".",
+    and appends the terminating ".\\r\\n" on its own line. When ``max_body_lines``
+    is given, the full header block is kept and the body is truncated to that
+    many lines (TOP semantics).
+    """
+    lines = _split_lines(body)
+    if max_body_lines is not None:
+        try:
+            separator = lines.index(b"")
+        except ValueError:
+            separator = len(lines)
+        header_lines = lines[: separator + 1]
+        body_lines = lines[separator + 1 :][:max_body_lines]
+        lines = header_lines + body_lines
+    stuffed = []
+    for line in lines:
+        if line.startswith(b"."):
+            line = b"." + line
+        stuffed.append(line)
+    stuffed.append(b".")
+    return b"\r\n".join(stuffed) + b"\r\n"
+
+
 class IMAPBackend:
     """Minimal IMAP wrapper for backend mailbox access."""
 
@@ -76,9 +111,13 @@ class IMAPBackend:
 
     def _connect(self):
         if Settings.BACKEND_IMAP_USE_SSL:
-            client = imaplib.IMAP4_SSL(Settings.BACKEND_IMAP_HOST, Settings.BACKEND_IMAP_PORT)
+            client = imaplib.IMAP4_SSL(
+                Settings.BACKEND_IMAP_HOST, Settings.BACKEND_IMAP_PORT, timeout=30
+            )
         else:
-            client = imaplib.IMAP4(Settings.BACKEND_IMAP_HOST, Settings.BACKEND_IMAP_PORT)
+            client = imaplib.IMAP4(
+                Settings.BACKEND_IMAP_HOST, Settings.BACKEND_IMAP_PORT, timeout=30
+            )
             if Settings.BACKEND_IMAP_USE_STARTTLS:
                 client.starttls(ssl_context=ssl.create_default_context())
         return client
@@ -116,12 +155,30 @@ class IMAPBackend:
         match = re.search(rb"RFC822\.SIZE (\d+)", response)
         return int(match.group(1)) if match else 0
 
+    def fetch_all_sizes(self):
+        """Return a {uid: size} mapping for all messages in one IMAP round-trip."""
+        sizes = {}
+        typ, data = self.connection.uid("fetch", "1:*", "(RFC822.SIZE)")
+        if typ != "OK" or not data:
+            return sizes
+        for response in data:
+            if isinstance(response, tuple):
+                response = response[0]
+            if not isinstance(response, bytes):
+                continue
+            match = self.UID_PATTERN.search(response)
+            if match:
+                sizes[match.group(1)] = int(match.group(2))
+        return sizes
+
     def fetch_message(self, uid):
         typ, data = self.connection.uid("fetch", uid, "(RFC822)")
         if typ != "OK" or not data:
             raise RuntimeError("IMAP message fetch failed")
-        parts = [chunk for chunk in data if isinstance(chunk, bytes)]
-        return b"\r\n".join(parts) + b"\r\n"
+        for chunk in data:
+            if isinstance(chunk, tuple) and len(chunk) > 1:
+                return chunk[1]
+        raise RuntimeError("IMAP message fetch returned no body")
 
     def mark_deleted(self, uid):
         self.connection.uid("STORE", uid, "+FLAGS.SILENT", "(\\Deleted)")
@@ -179,6 +236,8 @@ class POP3Session:
             return await self.handle_list(args)
         if command == "RETR":
             return await self.handle_retr(args)
+        if command == "TOP":
+            return await self.handle_top(args)
         if command == "DELE":
             return await self.handle_dele(args)
         if command == "NOOP":
@@ -220,29 +279,26 @@ class POP3Session:
         backend = IMAPBackend(username, password)
         backend.login()
         self._imap = backend
-        self._refresh_mailbox()
-
-    def _refresh_mailbox(self):
-        self.message_ids = self._imap.list_uids()
-        self.deleted.clear()
+        # Snapshot the maildrop once; it stays static for the session lifetime
+        # (RFC 1939), so DELE marks are not wiped by later STAT/LIST/UIDL.
+        self.message_ids = backend.list_uids()
 
     async def handle_stat(self):
         self._require_auth()
-        self._refresh_mailbox()
+        sizes = await asyncio.to_thread(self._imap.fetch_all_sizes)
         count = len(self.message_ids) - len(self.deleted)
-        total_size = sum(self._imap.fetch_message_size(uid) for uid in self.message_ids if uid not in self.deleted)
+        total_size = sum(sizes.get(uid, 0) for uid in self.message_ids if uid not in self.deleted)
         return await self.send_line(f"+OK {count} {total_size}")
 
     async def handle_list(self, args):
         self._require_auth()
-        self._refresh_mailbox()
+        sizes = await asyncio.to_thread(self._imap.fetch_all_sizes)
         if not args:
             lines = [f"+OK {len(self.message_ids)} messages"]
             for index, uid in enumerate(self.message_ids, start=1):
                 if uid in self.deleted:
                     continue
-                size = self._imap.fetch_message_size(uid)
-                lines.append(f"{index} {size}")
+                lines.append(f"{index} {sizes.get(uid, 0)}")
             lines.append(".")
             return await self.send_lines(lines)
         if len(args) != 1 or not args[0].isdigit():
@@ -253,8 +309,7 @@ class POP3Session:
         uid = self.message_ids[message_number - 1]
         if uid in self.deleted:
             return await self.send_line("-ERR message deleted")
-        size = self._imap.fetch_message_size(uid)
-        return await self.send_line(f"+OK {message_number} {size}")
+        return await self.send_line(f"+OK {message_number} {sizes.get(uid, 0)}")
 
     async def handle_retr(self, args):
         self._require_auth()
@@ -267,9 +322,25 @@ class POP3Session:
         if uid in self.deleted:
             return await self.send_line("-ERR message deleted")
         message = await asyncio.to_thread(self._imap.fetch_message, uid)
+        payload = format_pop3_body(message)
         await self.send_line(f"+OK {len(message)} octets")
-        await self.send_raw(message)
-        await self.send_line(".")
+        await self.send_raw(payload)
+
+    async def handle_top(self, args):
+        self._require_auth()
+        if len(args) != 2 or not args[0].isdigit() or not args[1].isdigit():
+            return await self.send_line("-ERR TOP requires message number and line count")
+        index = int(args[0])
+        line_count = int(args[1])
+        if index < 1 or index > len(self.message_ids):
+            return await self.send_line("-ERR no such message")
+        uid = self.message_ids[index - 1]
+        if uid in self.deleted:
+            return await self.send_line("-ERR message deleted")
+        message = await asyncio.to_thread(self._imap.fetch_message, uid)
+        payload = format_pop3_body(message, max_body_lines=line_count)
+        await self.send_line("+OK top of message follows")
+        await self.send_raw(payload)
 
     async def handle_dele(self, args):
         self._require_auth()
@@ -284,7 +355,6 @@ class POP3Session:
 
     async def handle_uidl(self, args):
         self._require_auth()
-        self._refresh_mailbox()
         if not args:
             lines = ["+OK UID list follows"]
             for index, uid in enumerate(self.message_ids, start=1):
@@ -346,7 +416,6 @@ class SMTPProxyHandler:
 
     def send_message(self, sender, recipients, data):
         """Forward a complete SMTP message to the backend SMTP server."""
-        message = data.decode("utf-8", errors="replace")
         if Settings.BACKEND_SMTP_USE_SSL:
             smtp = smtplib.SMTP_SSL(Settings.BACKEND_SMTP_HOST, Settings.BACKEND_SMTP_PORT, timeout=30)
         else:
@@ -356,7 +425,7 @@ class SMTPProxyHandler:
         try:
             if Settings.BACKEND_SMTP_USER and Settings.BACKEND_SMTP_PASS:
                 smtp.login(Settings.BACKEND_SMTP_USER, Settings.BACKEND_SMTP_PASS)
-            smtp.sendmail(sender, recipients, message)
+            smtp.sendmail(sender, recipients, data)
         finally:
             smtp.quit()