From a2444da9df9177118261e4332187d4f262cd4fc7 Mon Sep 17 00:00:00 2001 From: cobryan05 <13701027+cobryan05@users.noreply.github.com> Date: Thu, 14 Sep 2023 13:51:40 -0500 Subject: [PATCH] Split queries to avoid overflowing device buffers (#502) Several KASA devices seem to have pretty strict buffer size limitations on incoming/outgoing data transfers. Testing on KL125-US and HL103 has shown that sending a request size larger than about ~768 bytes will immediately crash the device. Additionally, a query that generates a response larger than ~4096 bytes will crash the KL125-US. I was unable to generate such a large response to test the HL103. The KL125-US will only return such large queries when its monthly usage stats have been populated. This means that a new bulb would work fine, but after a month of data collection the bulb would break the 4K limit and start to crash. To work around this issue, an estimated worst-case response size is calculated before sending a request by summing up all modules estimated response size. If the estimated size is greater than the device's max_response_payload_size then the query will be split into multiple queries. This PR implements splitting queries expected to have large responses and also removes the module 'skip list' which was a previous workaround to the crash (which worked by simply reducing the number of modules queried, which prevented the overflow) since it is no longer necessary. This PR does not attempt to address the "input buffer size limit." Thus far this limit has not been an issue. --- kasa/modules/module.py | 9 +++++++++ kasa/modules/usage.py | 5 +++++ kasa/smartbulb.py | 5 +++++ kasa/smartdevice.py | 30 ++++++++++++++++++++++-------- kasa/tests/test_smartdevice.py | 15 ++++++++++++++- 5 files changed, 55 insertions(+), 9 deletions(-) diff --git a/kasa/modules/module.py b/kasa/modules/module.py index 7340d7e1..ff806a99 100644 --- a/kasa/modules/module.py +++ b/kasa/modules/module.py @@ -43,6 +43,15 @@ class Module(ABC): queries to the query that gets executed when Device.update() gets called. """ + @property + def estimated_query_response_size(self): + """Estimated maximum size of query response. + + The inheriting modules implement this to estimate how large a query response + will be so that queries can be split should an estimated response be too large + """ + return 256 # Estimate for modules that don't specify + @property def data(self): """Return the module specific raw data from the last update.""" diff --git a/kasa/modules/usage.py b/kasa/modules/usage.py index d1f96e7e..9877795d 100644 --- a/kasa/modules/usage.py +++ b/kasa/modules/usage.py @@ -21,6 +21,11 @@ class Usage(Module): return req + @property + def estimated_query_response_size(self): + """Estimated maximum query response size.""" + return 2048 + @property def daily_data(self): """Return statistics on daily basis.""" diff --git a/kasa/smartbulb.py b/kasa/smartbulb.py index 2d2f28ca..a09487d2 100644 --- a/kasa/smartbulb.py +++ b/kasa/smartbulb.py @@ -543,3 +543,8 @@ class SmartBulb(SmartDevice): return await self._query_helper( self.LIGHT_SERVICE, "set_preferred_state", preset.dict(exclude_none=True) ) + + @property + def max_device_response_size(self) -> int: + """Returns the maximum response size the device can safely construct.""" + return 4096 diff --git a/kasa/smartdevice.py b/kasa/smartdevice.py index 5c24c943..81ea135b 100755 --- a/kasa/smartdevice.py +++ b/kasa/smartdevice.py @@ -28,9 +28,6 @@ from .protocol import TPLinkSmartHomeProtocol _LOGGER = logging.getLogger(__name__) -# Certain module queries will crash devices; this list skips those queries -MODEL_MODULE_SKIPLIST = {"KL125(US)": ["cloud"]} # Issue #345 - class DeviceType(Enum): """Device type enum.""" @@ -337,20 +334,32 @@ class SmartDevice: ) self.add_module("emeter", Emeter(self, self.emeter_type)) + request_list = [] + est_response_size = 1024 if "system" in req else 0 for module_name, module in self.modules.items(): if not module.is_supported: _LOGGER.debug("Module %s not supported, skipping" % module) continue - modules_to_skip = MODEL_MODULE_SKIPLIST.get(self.model, []) - if module_name in modules_to_skip: - _LOGGER.debug(f"Module {module} is excluded for {self.model}, skipping") - continue + + est_response_size += module.estimated_query_response_size + if est_response_size > self.max_device_response_size: + request_list.append(req) + req = {} + est_response_size = module.estimated_query_response_size q = module.query() _LOGGER.debug("Adding query for %s: %s", module, q) req = merge(req, q) + request_list.append(req) - self._last_update = await self.protocol.query(req) + responses = [ + await self.protocol.query(request) for request in request_list if request + ] + + update: Dict = {} + for response in responses: + update = {**update, **response} + self._last_update = update def update_from_discover_info(self, info): """Update state from info from the discover call.""" @@ -658,6 +667,11 @@ class SmartDevice: ) return self.children[index] + @property + def max_device_response_size(self) -> int: + """Returns the maximum response size the device can safely construct.""" + return 16 * 1024 + @property def device_type(self) -> DeviceType: """Return the device type.""" diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py index 0839bc06..c70f544b 100644 --- a/kasa/tests/test_smartdevice.py +++ b/kasa/tests/test_smartdevice.py @@ -39,7 +39,9 @@ async def test_initial_update_emeter(dev, mocker): dev._last_update = None spy = mocker.spy(dev.protocol, "query") await dev.update() - assert spy.call_count == 2 + len(dev.children) + # Devices with small buffers may require 3 queries + expected_queries = 2 if dev.max_device_response_size > 4096 else 3 + assert spy.call_count == expected_queries + len(dev.children) @no_emeter @@ -164,6 +166,17 @@ async def test_features(dev): assert dev.features == set() +async def test_max_device_response_size(dev): + """Make sure every device return has a set max response size.""" + assert dev.max_device_response_size > 0 + + +async def test_estimated_response_sizes(dev): + """Make sure every module has an estimated response size set.""" + for mod in dev.modules.values(): + assert mod.estimated_query_response_size > 0 + + @pytest.mark.parametrize("device_class", smart_device_classes) def test_device_class_ctors(device_class): """Make sure constructor api not broken for new and existing SmartDevices."""