fix: avoid MIME check for already-completed transfers and refactor handler

on_dccmsg was MIME-checking the first chunk even when the transfer was
already marked completed, which could reject re-transfers of small files
where the resumed chunk does not contain the file header. Skip the check
when completed.

Also refactor TransferHandler to keep on_dccmsg and on_dcc_disconnect
focused, reducing per-module cyclomatic complexity back under the B gate.

Author: luni

dccbot/transfer_handler.py

@@ -26,6 +26,79 @@ def __init__(self, bot: IRCBot) -> None:
         """Initialize transfer handler for a specific IRC bot instance."""
         self.bot = bot
 
+    def _touch_channel_activity(self, transfer: dict) -> None:
+        """Bump the activity timestamp for channels associated with this nick."""
+        nick = transfer["nick"].lower()
+        if nick in self.bot.bot_channel_map:
+            now = time.time()
+            for channel in self.bot.bot_channel_map[nick]:
+                self.bot.joined_channels[channel] = now
+
+    def _update_progress(self, transfer: dict) -> None:
+        """Recalculate and log transfer progress/rate if thresholds are met."""
+        now = time.time()
+        percent = int(100 * (transfer["bytes_received"] + transfer["offset"]) / transfer["size"])
+        if transfer["percent"] + 10 > percent and now - transfer["last_progress_update"] < 5:
+            return
+
+        transfer["percent"] = percent
+        elapsed = now - transfer["start_time"]
+        transfer_rate_avg = (transfer["bytes_received"] / elapsed) / 1024 if elapsed > 0 else 0
+        elapsed = now - transfer["last_progress_update"]
+        transferred = transfer["bytes_received"] - transfer["last_progress_bytes_received"]
+        transfer_rate = (transferred / elapsed) / 1024 if elapsed > 0 else 0
+
+        logger.info(
+            "[%s] Downloading %s %d%% @ %.2f KB/s / %.2f KB/s",
+            transfer["nick"],
+            transfer["filename"],
+            transfer["percent"],
+            transfer_rate,
+            transfer_rate_avg,
+        )
+        transfer["last_progress_update"] = now
+        transfer["last_progress_bytes_received"] = transfer["bytes_received"]
+
+    def _check_mime_or_disconnect(self, dcc: AioDCCConnection, transfer: dict, data: bytes) -> bool:
+        """Validate the first chunk's MIME type or close the connection."""
+        if transfer["bytes_received"] != 0 or transfer.get("offset") or not self.bot.allowed_mimetypes:
+            return True
+
+        mime_type = self.bot.mime_checker.from_buffer(data)
+        if mime_type in self.bot.allowed_mimetypes:
+            return True
+
+        logger.warning("[%s] Reject %s: Invalid MIME type (%s)", transfer["nick"], transfer["filename"], mime_type)
+        transfer["status"] = "error"
+        transfer["error"] = f"Invalid MIME type ({mime_type})"
+        transfer["connected"] = False
+        dcc.disconnect()
+        self.bot.current_transfers.pop(dcc, None)
+        return False
+
+    def _write_chunk_or_disconnect(self, dcc: AioDCCConnection, transfer: dict, data: bytes) -> bool:
+        """Append the chunk to the file or mark the transfer as failed."""
+        try:
+            with open(transfer["file_path"], "ab") as f:
+                f.write(data)
+        except Exception as e:
+            logger.error("Error writing to file %s: %s", transfer["file_path"], e)
+            transfer["status"] = "error"
+            transfer["error"] = f"Error writing to file {transfer['file_path']}: {e}"
+            transfer["connected"] = False
+            dcc.disconnect()
+            self.bot.current_transfers.pop(dcc, None)
+            return False
+        return True
+
+    def _send_ack(self, dcc: AioDCCConnection, transfer: dict) -> None:
+        """Send the appropriate 32- or 64-bit acknowledgement."""
+        ack = transfer["bytes_received"] + transfer["offset"]
+        if transfer["size"] >= 1024 * 1024 * 1024 * 4:
+            dcc.send_bytes(struct.pack("!Q", ack))
+        else:
+            dcc.send_bytes(struct.pack("!I", ack))
+
     def on_dccmsg(self, connection: AioDCCConnection, event: irc.client_aio.Event) -> None:
         """Handle incoming DCC data chunk."""
         dcc = connection
@@ -39,60 +112,64 @@ def on_dccmsg(self, connection: AioDCCConnection, event: irc.client_aio.Event) -
         data = event.arguments[0]
 
         if not transfer["completed"]:
-            now = time.time()
-            if transfer["nick"].lower() in self.bot.bot_channel_map:
-                for channel in self.bot.bot_channel_map[transfer["nick"].lower()]:
-                    self.bot.joined_channels[channel] = now
-
-            percent = int(100 * (transfer["bytes_received"] + transfer["offset"]) / transfer["size"])
-            if transfer["percent"] + 10 <= percent or now - transfer["last_progress_update"] >= 5:
-                transfer["percent"] = percent
-                elapsed_time = now - transfer["start_time"]
-                transfer_rate_avg = (transfer["bytes_received"] / elapsed_time) / 1024 if elapsed_time > 0 else 0
-                elapsed_time = now - transfer["last_progress_update"]
-                transferred_bytes = transfer["bytes_received"] - transfer["last_progress_bytes_received"]
-                transfer_rate = (transferred_bytes / elapsed_time) / 1024 if elapsed_time > 0 else 0
-
-                logger.info(
-                    "[%s] Downloading %s %d%% @ %.2f KB/s / %.2f KB/s",
-                    transfer["nick"],
-                    transfer["filename"],
-                    transfer["percent"],
-                    transfer_rate,
-                    transfer_rate_avg,
-                )
-                transfer["last_progress_update"] = now
-                transfer["last_progress_bytes_received"] = transfer["bytes_received"]
-
-            if transfer["bytes_received"] == 0 and not transfer.get("offset") and self.bot.allowed_mimetypes:
-                mime_type = self.bot.mime_checker.from_buffer(data)
-                if mime_type not in self.bot.allowed_mimetypes:
-                    logger.warning("[%s] Reject %s: Invalid MIME type (%s)", transfer["nick"], transfer["filename"], mime_type)
-                    transfer["status"] = "error"
-                    transfer["error"] = f"Invalid MIME type ({mime_type})"
-                    transfer["connected"] = False
-                    dcc.disconnect()
-                    self.bot.current_transfers.pop(dcc, None)
-                    return
+            self._touch_channel_activity(transfer)
+            self._update_progress(transfer)
+            if not self._check_mime_or_disconnect(dcc, transfer, data):
+                return
+            if not self._write_chunk_or_disconnect(dcc, transfer, data):
+                return
 
+        transfer["bytes_received"] += len(data)
+        self._send_ack(dcc, transfer)
+
+    def _mark_complete(self, transfer: dict, file_path: str, transfer_rate: float) -> None:
+        """Mark a transfer as complete, queue MD5 and rename incomplete file."""
+        logger.info("[%s] Download %s complete - size: %d bytes, %.2f KB/s", transfer["nick"], transfer["filename"], transfer["size"], transfer_rate)
+        transfer["completed"] = time.time()
+        transfer["status"] = "completed"
+        if transfer.get("md5"):
+            asyncio.create_task(self.bot._add_md5_check_queue_item(transfer))
+
+        incomplete_suffix = get_incomplete_suffix(self.bot.config)
+        if incomplete_suffix and file_path.endswith(incomplete_suffix):
+            target = file_path[: -len(incomplete_suffix)]
             try:
-                with open(transfer["file_path"], "ab") as f:
-                    f.write(data)
+                os.rename(file_path, target)
+                logger.info("Renamed downloaded file to %s", transfer["filename"])
+                transfer["file_path"] = target
             except Exception as e:
-                logger.error("Error writing to file %s: %s", transfer["file_path"], e)
-                transfer["status"] = "error"
-                transfer["error"] = f"Error writing to file {transfer['file_path']}: {e}"
-                transfer["connected"] = False
-                dcc.disconnect()
-                self.bot.current_transfers.pop(dcc, None)
-                return
+                logger.error("Error renaming %s to %s: %s", file_path, target, e)
+
+    def _report_size_mismatch(self, transfer: dict, file_size: int) -> None:
+        """Record a size mismatch failure for a transfer."""
+        logger.error("[%s] Download %s failed: size mismatch %d != %d", transfer["nick"], transfer["filename"], file_size, transfer["size"])
+        if transfer["status"] != "error":
+            transfer["status"] = "failed"
+            transfer["error"] = f"size mismatch {file_size} != {transfer['size']}"
+
+    def _report_missing_file(self, transfer: dict, file_path: str) -> None:
+        """Record a missing output file failure."""
+        logger.error("[%s] Download failed: %s does not exist", transfer["nick"], file_path)
+        if transfer["status"] != "error":
+            transfer["status"] = "error"
+            transfer["error"] = f"[{transfer['nick']}] Download failed: {file_path} does not exist"
+
+    def _finalize_transfer(self, dcc: AioDCCConnection, transfer: dict) -> None:
+        """Check the output file and finalize the transfer state."""
+        file_path = transfer["file_path"]
+        elapsed = time.time() - transfer["start_time"]
+        transfer_rate = (transfer["bytes_received"] / elapsed) / 1024 if elapsed > 0 else 0
 
-        transfer["bytes_received"] += len(data)
-        ack = transfer["bytes_received"] + transfer["offset"]
-        if transfer["size"] >= 1024 * 1024 * 1024 * 4:
-            dcc.send_bytes(struct.pack("!Q", ack))
+        if not os.path.exists(file_path):
+            self._report_missing_file(transfer, file_path)
         else:
-            dcc.send_bytes(struct.pack("!I", ack))
+            file_size = os.path.getsize(file_path)
+            if file_size != transfer["size"]:
+                self._report_size_mismatch(transfer, file_size)
+            else:
+                self._mark_complete(transfer, file_path, transfer_rate)
+
+        self.bot.current_transfers.pop(dcc, None)
 
     def on_dcc_disconnect(self, connection: AioDCCConnection, event: irc.client_aio.Event) -> None:
         """Handle DCC disconnect and finalize transfer metadata."""
@@ -104,42 +181,5 @@ def on_dcc_disconnect(self, connection: AioDCCConnection, event: irc.client_aio.
 
         transfer = self.bot.current_transfers[dcc]
         transfer["connected"] = False
-
-        if transfer["nick"].lower() in self.bot.bot_channel_map:
-            for channel in self.bot.bot_channel_map[transfer["nick"].lower()]:
-                self.bot.joined_channels[channel] = time.time()
-
-        file_path = transfer["file_path"]
-        elapsed_time = time.time() - transfer["start_time"]
-        transfer_rate = (transfer["bytes_received"] / elapsed_time) / 1024 if elapsed_time > 0 else 0
-
-        if not os.path.exists(file_path):
-            logger.error("[%s] Download failed: %s does not exist", transfer["nick"], file_path)
-            if transfer["status"] != "error":
-                transfer["status"] = "error"
-                transfer["error"] = f"[{transfer['nick']}] Download failed: {file_path} does not exist"
-        else:
-            file_size = os.path.getsize(file_path)
-            if file_size != transfer["size"]:
-                logger.error("[%s] Download %s failed: size mismatch %d != %d", transfer["nick"], transfer["filename"], file_size, transfer["size"])
-                if transfer["status"] != "error":
-                    transfer["status"] = "failed"
-                    transfer["error"] = f"size mismatch {file_size} != {transfer['size']}"
-            else:
-                logger.info("[%s] Download %s complete - size: %d bytes, %.2f KB/s", transfer["nick"], transfer["filename"], file_size, transfer_rate)
-                transfer["completed"] = time.time()
-                transfer["status"] = "completed"
-                if transfer.get("md5"):
-                    asyncio.create_task(self.bot._add_md5_check_queue_item(transfer))
-
-                incomplete_suffix = get_incomplete_suffix(self.bot.config)
-                if incomplete_suffix and file_path.endswith(incomplete_suffix):
-                    target = file_path[: -len(incomplete_suffix)]
-                    try:
-                        os.rename(file_path, target)
-                        logger.info("Renamed downloaded file to %s", transfer["filename"])
-                        transfer["file_path"] = target
-                    except Exception as e:
-                        logger.error("Error renaming %s to %s: %s", file_path, target, e)
-
-        self.bot.current_transfers.pop(dcc, None)
+        self._touch_channel_activity(transfer)
+        self._finalize_transfer(dcc, transfer)

tests/unit/test_transfer_handler.py

@@ -140,6 +140,26 @@ def test_on_dccmsg_write_failure_sets_error():
     dcc.disconnect.assert_called_once()
 
 
+def test_on_dccmsg_skips_mime_check_for_completed_transfer():
+    """Test transfer handler does not MIME-check already-completed transfers."""
+    dcc = MagicMock()
+    transfer = _make_transfer(size=1024)
+    transfer["completed"] = True
+    bot = _make_bot_with_transfer(dcc, transfer)
+    bot.allowed_mimetypes = ["video/mp4"]
+    bot.mime_checker.from_buffer.return_value = "text/plain"
+    handler = TransferHandler(bot)
+    event = MagicMock()
+    event.arguments = [b"abc"]
+
+    with patch("builtins.open", mock_open()):
+        handler.on_dccmsg(dcc, event)
+
+    bot.mime_checker.from_buffer.assert_not_called()
+    dcc.disconnect.assert_not_called()
+    dcc.send_bytes.assert_called_once()
+
+
 def test_on_dcc_disconnect_unknown_connection():
     """Test unknown disconnect events are ignored safely."""
     dcc = MagicMock()