Enable open protection in the event loop (#117289)

pull/105424/head^2
J. Nick Koston 2024-05-13 08:50:31 +09:00 committed by GitHub
parent d06932bbc2
commit 11f49280c9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 55 additions and 4 deletions

View File

@ -1,5 +1,6 @@
"""Block blocking calls being done in asyncio."""
import builtins
from contextlib import suppress
from http.client import HTTPConnection
import importlib
@ -13,12 +14,21 @@ from .util.loop import protect_loop
_IN_TESTS = "unittest" in sys.modules
ALLOWED_FILE_PREFIXES = ("/proc",)
def _check_import_call_allowed(mapped_args: dict[str, Any]) -> bool:
# If the module is already imported, we can ignore it.
return bool((args := mapped_args.get("args")) and args[0] in sys.modules)
def _check_file_allowed(mapped_args: dict[str, Any]) -> bool:
# If the file is in /proc we can ignore it.
args = mapped_args["args"]
path = args[0] if type(args[0]) is str else str(args[0]) # noqa: E721
return path.startswith(ALLOWED_FILE_PREFIXES)
def _check_sleep_call_allowed(mapped_args: dict[str, Any]) -> bool:
#
# Avoid extracting the stack unless we need to since it
@ -50,11 +60,15 @@ def enable() -> None:
loop_thread_id=loop_thread_id,
)
# Currently disabled. pytz doing I/O when getting timezone.
# Prevent files being opened inside the event loop
# builtins.open = protect_loop(builtins.open)
if not _IN_TESTS:
# Prevent files being opened inside the event loop
builtins.open = protect_loop( # type: ignore[assignment]
builtins.open,
strict_core=False,
strict=False,
check_allowed=_check_file_allowed,
loop_thread_id=loop_thread_id,
)
# unittest uses `importlib.import_module` to do mocking
# so we cannot protect it if we are running tests
importlib.import_module = protect_loop(

View File

@ -1,7 +1,10 @@
"""Tests for async util methods from Python source."""
import contextlib
import importlib
from pathlib import Path, PurePosixPath
import time
from typing import Any
from unittest.mock import Mock, patch
import pytest
@ -198,3 +201,37 @@ async def test_protect_loop_importlib_import_module_in_integration(
"Detected blocking call to import_module inside the event loop by "
"integration 'hue' at homeassistant/components/hue/light.py, line 23"
) in caplog.text
async def test_protect_loop_open(caplog: pytest.LogCaptureFixture) -> None:
"""Test open of a file in /proc is not reported."""
block_async_io.enable()
with contextlib.suppress(FileNotFoundError):
open("/proc/does_not_exist").close()
assert "Detected blocking call to open with args" not in caplog.text
async def test_protect_open(caplog: pytest.LogCaptureFixture) -> None:
"""Test opening a file in the event loop logs."""
block_async_io.enable()
with contextlib.suppress(FileNotFoundError):
open("/config/data_not_exist").close()
assert "Detected blocking call to open with args" in caplog.text
@pytest.mark.parametrize(
"path",
[
"/config/data_not_exist",
Path("/config/data_not_exist"),
PurePosixPath("/config/data_not_exist"),
],
)
async def test_protect_open_path(path: Any, caplog: pytest.LogCaptureFixture) -> None:
"""Test opening a file by path in the event loop logs."""
block_async_io.enable()
with contextlib.suppress(FileNotFoundError):
open(path).close()
assert "Detected blocking call to open with args" in caplog.text