From ad64e17a345826140be844f81a90967a28ffbed9 Mon Sep 17 00:00:00 2001 From: kirs-hi Date: Fri, 12 Jun 2026 11:03:04 +0800 Subject: [PATCH] fix: avoid KeyError on cancel and infinite loop in image compression MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two independent robustness fixes: 1. ChatChannel.cancel_session / cancel_all_session raised KeyError when a session existed in self.sessions but no future had been dispatched yet. self.sessions[sid] is created in produce(), but self.futures[sid] is only created later in consume() on first dispatch. Cancelling in that window (e.g. user sends a message then immediately cancels) crashed the cancel path. Use self.futures.get(sid, []) so an absent entry is a no-op. 2. compress_imgfile decremented JPEG quality by 5 with no lower bound. For an image that cannot be compressed below max_size, quality went 0, negative, ... — the loop never terminated and passed invalid quality values to PIL. Add a min_quality floor (10) and return the best effort once reached. Adds tests/test_robustness_fixes.py covering both paths (5 tests). --- channel/chat_channel.py | 7 +- common/utils.py | 6 +- tests/test_robustness_fixes.py | 121 +++++++++++++++++++++++++++++++++ 3 files changed, 131 insertions(+), 3 deletions(-) create mode 100644 tests/test_robustness_fixes.py diff --git a/channel/chat_channel.py b/channel/chat_channel.py index 9104a38e..e2a65c5b 100644 --- a/channel/chat_channel.py +++ b/channel/chat_channel.py @@ -519,7 +519,10 @@ class ChatChannel(Channel): def cancel_session(self, session_id): with self.lock: if session_id in self.sessions: - for future in self.futures[session_id]: + # futures[session_id] is only created in consume() when a task is + # dispatched, so it may be absent if cancel happens right after + # produce() but before the first dispatch. Default to []. + for future in self.futures.get(session_id, []): future.cancel() cnt = self.sessions[session_id][0].qsize() if cnt > 0: @@ -529,7 +532,7 @@ class ChatChannel(Channel): def cancel_all_session(self): with self.lock: for session_id in self.sessions: - for future in self.futures[session_id]: + for future in self.futures.get(session_id, []): future.cancel() cnt = self.sessions[session_id][0].qsize() if cnt > 0: diff --git a/common/utils.py b/common/utils.py index e7264e20..0597d979 100644 --- a/common/utils.py +++ b/common/utils.py @@ -27,10 +27,14 @@ def compress_imgfile(file, max_size): img = Image.open(file) rgb_image = img.convert("RGB") quality = 95 + min_quality = 10 while True: out_buf = io.BytesIO() rgb_image.save(out_buf, "JPEG", quality=quality) - if fsize(out_buf) <= max_size: + if fsize(out_buf) <= max_size or quality <= min_quality: + # Stop at min_quality: further decrements would pass an invalid + # quality (<1) to PIL and the loop would otherwise never terminate + # for images that cannot be compressed below max_size. return out_buf quality -= 5 diff --git a/tests/test_robustness_fixes.py b/tests/test_robustness_fixes.py new file mode 100644 index 00000000..9f8aa453 --- /dev/null +++ b/tests/test_robustness_fixes.py @@ -0,0 +1,121 @@ +# encoding:utf-8 +""" +Unit tests for robustness fixes: + 1. ChatChannel.cancel_session / cancel_all_session must not raise KeyError + when a session has been produced but no task has been dispatched yet + (so self.futures[session_id] does not exist). + 2. common.utils.compress_imgfile must terminate (no infinite loop / invalid + PIL quality) when an image cannot be compressed below max_size. +""" +import io +import os +import sys +import types +import unittest +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + + +# ============================================================================= +# 1. cancel_session / cancel_all_session KeyError regression +# ============================================================================= + +class TestCancelSessionMissingFutures(unittest.TestCase): + """A session may exist in self.sessions before any future is recorded.""" + + def _make_channel(self): + # Import lazily and build a bare object without running __init__, + # to avoid pulling the full channel setup / config. + from channel.chat_channel import ChatChannel + + ch = ChatChannel.__new__(ChatChannel) + import threading + + ch.lock = threading.RLock() + # A produced session whose future has NOT been dispatched yet. + queue = MagicMock() + queue.qsize.return_value = 0 + semaphore = MagicMock() + ch.sessions = {"sid": [queue, semaphore]} + ch.futures = {} # intentionally empty: consume() never ran + return ch + + def test_cancel_session_no_futures_entry(self): + ch = self._make_channel() + # Should not raise KeyError. + try: + ch.cancel_session("sid") + except KeyError: + self.fail("cancel_session raised KeyError when futures entry missing") + + def test_cancel_all_session_no_futures_entry(self): + ch = self._make_channel() + try: + ch.cancel_all_session() + except KeyError: + self.fail("cancel_all_session raised KeyError when futures entry missing") + + def test_cancel_session_cancels_existing_futures(self): + ch = self._make_channel() + fut = MagicMock() + ch.futures["sid"] = [fut] + ch.cancel_session("sid") + fut.cancel.assert_called_once() + + +# ============================================================================= +# 2. compress_imgfile termination +# ============================================================================= + +class TestCompressImgfileTermination(unittest.TestCase): + """compress_imgfile must always return, even for incompressible input.""" + + def setUp(self): + # Skip if Pillow is not available in the test environment. + try: + import PIL # noqa: F401 + except ImportError: + self.skipTest("Pillow not installed") + + def _make_image_buf(self, size=(64, 64)): + from PIL import Image + import random + + img = Image.new("RGB", size) + # Fill with random noise so JPEG cannot compress it well. + pixels = [ + (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) + for _ in range(size[0] * size[1]) + ] + img.putdata(pixels) + buf = io.BytesIO() + img.save(buf, "JPEG", quality=95) + buf.seek(0) + return buf + + def test_returns_when_target_unreachable(self): + from common.utils import compress_imgfile + + buf = self._make_image_buf() + # An impossibly small target that even quality=10 won't reach. + out = compress_imgfile(buf, max_size=10) + self.assertIsInstance(out, io.BytesIO) + # Verify the result is still a valid JPEG (PIL never got invalid quality). + from PIL import Image + + out.seek(0) + img = Image.open(out) + img.verify() + + def test_no_compression_needed_returns_same_object(self): + from common.utils import compress_imgfile + + buf = self._make_image_buf() + size = buf.getbuffer().nbytes + out = compress_imgfile(buf, max_size=size + 1) + self.assertIs(out, buf) + + +if __name__ == "__main__": + unittest.main()