Ensure response is fully read to prevent premature connection closure in rest command (#148532)

This commit is contained in:
Jan-Philipp Benecke 2025-07-11 09:06:18 +02:00 committed by GitHub
parent 32121a073c
commit cd73824e3e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 30 additions and 1 deletions

View File

@ -178,6 +178,11 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
) )
if not service.return_response: if not service.return_response:
# always read the response to avoid closing the connection
# before the server has finished sending it, while avoiding excessive memory usage
async for _ in response.content.iter_chunked(1024):
pass
return None return None
_content = None _content = None

View File

@ -328,7 +328,7 @@ async def test_rest_command_get_response_malformed_json(
aioclient_mock.get( aioclient_mock.get(
TEST_URL, TEST_URL,
content='{"status": "failure", 42', content=b'{"status": "failure", 42',
headers={"content-type": "application/json"}, headers={"content-type": "application/json"},
) )
@ -381,3 +381,27 @@ async def test_rest_command_get_response_none(
) )
assert not response assert not response
async def test_rest_command_response_iter_chunked(
hass: HomeAssistant,
setup_component: ComponentSetup,
aioclient_mock: AiohttpClientMocker,
) -> None:
"""Ensure response is consumed when return_response is False."""
await setup_component()
png = base64.decodebytes(
b"iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAIAAAB7QOjdAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQ"
b"UAAAAJcEhZcwAAFiUAABYlAUlSJPAAAAAPSURBVBhXY/h/ku////8AECAE1JZPvDAAAAAASUVORK5CYII="
)
aioclient_mock.get(TEST_URL, content=png)
with patch("aiohttp.StreamReader.iter_chunked", autospec=True) as mock_iter_chunked:
response = await hass.services.async_call(DOMAIN, "get_test", {}, blocking=True)
# Ensure the response is not returned
assert response is None
# Verify iter_chunked was called with a chunk size
assert mock_iter_chunked.called