Speed up hassio ingress (#95777)

This commit is contained in:
J. Nick Koston 2023-07-07 22:50:39 -10:00 committed by GitHub
parent bbf97fdf01
commit 2b4f6ffcd6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 161 additions and 62 deletions

View file

@ -85,6 +85,13 @@ NO_STORE = re.compile(
# pylint: enable=implicit-str-concat
# fmt: on
RESPONSE_HEADERS_FILTER = {
TRANSFER_ENCODING,
CONTENT_LENGTH,
CONTENT_TYPE,
CONTENT_ENCODING,
}
class HassIOView(HomeAssistantView):
"""Hass.io view to handle base part."""
@ -170,8 +177,9 @@ class HassIOView(HomeAssistantView):
)
response.content_type = client.content_type
response.enable_compression()
await response.prepare(request)
async for data in client.content.iter_chunked(4096):
async for data in client.content.iter_chunked(8192):
await response.write(data)
return response
@ -190,21 +198,13 @@ class HassIOView(HomeAssistantView):
def _response_header(response: aiohttp.ClientResponse, path: str) -> dict[str, str]:
"""Create response header."""
headers = {}
for name, value in response.headers.items():
if name in (
TRANSFER_ENCODING,
CONTENT_LENGTH,
CONTENT_TYPE,
CONTENT_ENCODING,
):
continue
headers[name] = value
headers = {
name: value
for name, value in response.headers.items()
if name not in RESPONSE_HEADERS_FILTER
}
if NO_STORE.match(path):
headers[CACHE_CONTROL] = "no-store, max-age=0"
return headers

View file

@ -17,11 +17,32 @@ from yarl import URL
from homeassistant.components.http import HomeAssistantView
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.typing import UNDEFINED
from .const import X_HASS_SOURCE, X_INGRESS_PATH
_LOGGER = logging.getLogger(__name__)
INIT_HEADERS_FILTER = {
hdrs.CONTENT_LENGTH,
hdrs.CONTENT_ENCODING,
hdrs.TRANSFER_ENCODING,
hdrs.ACCEPT_ENCODING, # Avoid local compression, as we will compress at the border
hdrs.SEC_WEBSOCKET_EXTENSIONS,
hdrs.SEC_WEBSOCKET_PROTOCOL,
hdrs.SEC_WEBSOCKET_VERSION,
hdrs.SEC_WEBSOCKET_KEY,
}
RESPONSE_HEADERS_FILTER = {
hdrs.TRANSFER_ENCODING,
hdrs.CONTENT_LENGTH,
hdrs.CONTENT_TYPE,
hdrs.CONTENT_ENCODING,
}
MIN_COMPRESSED_SIZE = 128
MAX_SIMPLE_RESPONSE_SIZE = 4194000
@callback
def async_setup_ingress_view(hass: HomeAssistant, host: str):
@ -145,28 +166,35 @@ class HassIOIngress(HomeAssistantView):
skip_auto_headers={hdrs.CONTENT_TYPE},
) as result:
headers = _response_header(result)
content_length_int = 0
content_length = result.headers.get(hdrs.CONTENT_LENGTH, UNDEFINED)
# Simple request
if (
hdrs.CONTENT_LENGTH in result.headers
and int(result.headers.get(hdrs.CONTENT_LENGTH, 0)) < 4194000
) or result.status in (204, 304):
if result.status in (204, 304) or (
content_length is not UNDEFINED
and (content_length_int := int(content_length or 0))
<= MAX_SIMPLE_RESPONSE_SIZE
):
# Return Response
body = await result.read()
return web.Response(
simple_response = web.Response(
headers=headers,
status=result.status,
content_type=result.content_type,
body=body,
)
if content_length_int > MIN_COMPRESSED_SIZE:
simple_response.enable_compression()
await simple_response.prepare(request)
return simple_response
# Stream response
response = web.StreamResponse(status=result.status, headers=headers)
response.content_type = result.content_type
try:
response.enable_compression()
await response.prepare(request)
async for data in result.content.iter_chunked(4096):
async for data in result.content.iter_chunked(8192):
await response.write(data)
except (
@ -179,24 +207,20 @@ class HassIOIngress(HomeAssistantView):
return response
@lru_cache(maxsize=32)
def _forwarded_for_header(forward_for: str | None, peer_name: str) -> str:
"""Create X-Forwarded-For header."""
connected_ip = ip_address(peer_name)
return f"{forward_for}, {connected_ip!s}" if forward_for else f"{connected_ip!s}"
def _init_header(request: web.Request, token: str) -> CIMultiDict | dict[str, str]:
"""Create initial header."""
headers = {}
# filter flags
for name, value in request.headers.items():
if name in (
hdrs.CONTENT_LENGTH,
hdrs.CONTENT_ENCODING,
hdrs.TRANSFER_ENCODING,
hdrs.SEC_WEBSOCKET_EXTENSIONS,
hdrs.SEC_WEBSOCKET_PROTOCOL,
hdrs.SEC_WEBSOCKET_VERSION,
hdrs.SEC_WEBSOCKET_KEY,
):
continue
headers[name] = value
headers = {
name: value
for name, value in request.headers.items()
if name not in INIT_HEADERS_FILTER
}
# Ingress information
headers[X_HASS_SOURCE] = "core.ingress"
headers[X_INGRESS_PATH] = f"/api/hassio_ingress/{token}"
@ -208,12 +232,7 @@ def _init_header(request: web.Request, token: str) -> CIMultiDict | dict[str, st
_LOGGER.error("Can't set forward_for header, missing peername")
raise HTTPBadRequest()
connected_ip = ip_address(peername[0])
if forward_for:
forward_for = f"{forward_for}, {connected_ip!s}"
else:
forward_for = f"{connected_ip!s}"
headers[hdrs.X_FORWARDED_FOR] = forward_for
headers[hdrs.X_FORWARDED_FOR] = _forwarded_for_header(forward_for, peername[0])
# Set X-Forwarded-Host
if not (forward_host := request.headers.get(hdrs.X_FORWARDED_HOST)):
@ -223,7 +242,7 @@ def _init_header(request: web.Request, token: str) -> CIMultiDict | dict[str, st
# Set X-Forwarded-Proto
forward_proto = request.headers.get(hdrs.X_FORWARDED_PROTO)
if not forward_proto:
forward_proto = request.url.scheme
forward_proto = request.scheme
headers[hdrs.X_FORWARDED_PROTO] = forward_proto
return headers
@ -231,31 +250,20 @@ def _init_header(request: web.Request, token: str) -> CIMultiDict | dict[str, st
def _response_header(response: aiohttp.ClientResponse) -> dict[str, str]:
"""Create response header."""
headers = {}
for name, value in response.headers.items():
if name in (
hdrs.TRANSFER_ENCODING,
hdrs.CONTENT_LENGTH,
hdrs.CONTENT_TYPE,
hdrs.CONTENT_ENCODING,
):
continue
headers[name] = value
return headers
return {
name: value
for name, value in response.headers.items()
if name not in RESPONSE_HEADERS_FILTER
}
def _is_websocket(request: web.Request) -> bool:
"""Return True if request is a websocket."""
headers = request.headers
if (
return bool(
"upgrade" in headers.get(hdrs.CONNECTION, "").lower()
and headers.get(hdrs.UPGRADE, "").lower() == "websocket"
):
return True
return False
)
async def _websocket_forward(ws_from, ws_to):

View file

@ -348,3 +348,94 @@ async def test_forwarding_paths_as_requested(
"/api/hassio_ingress/mock-token/hello/%252e./world",
)
assert await resp.text() == "test"
@pytest.mark.parametrize(
"build_type",
[
("a3_vl", "test/beer/ping?index=1"),
("core", "index.html"),
("local", "panel/config"),
("jk_921", "editor.php?idx=3&ping=5"),
("fsadjf10312", ""),
],
)
async def test_ingress_request_get_compressed(
hassio_noauth_client, build_type, aioclient_mock: AiohttpClientMocker
) -> None:
"""Test ingress compressed."""
body = "this_is_long_enough_to_be_compressed" * 100
aioclient_mock.get(
f"http://127.0.0.1/ingress/{build_type[0]}/{build_type[1]}",
text=body,
headers={"Content-Length": len(body)},
)
resp = await hassio_noauth_client.get(
f"/api/hassio_ingress/{build_type[0]}/{build_type[1]}",
headers={"X-Test-Header": "beer", "Accept-Encoding": "gzip, deflate"},
)
# Check we got right response
assert resp.status == HTTPStatus.OK
body = await resp.text()
assert body == body
assert resp.headers["Content-Encoding"] == "deflate"
# Check we forwarded command
assert len(aioclient_mock.mock_calls) == 1
assert X_AUTH_TOKEN not in aioclient_mock.mock_calls[-1][3]
assert aioclient_mock.mock_calls[-1][3]["X-Hass-Source"] == "core.ingress"
assert (
aioclient_mock.mock_calls[-1][3]["X-Ingress-Path"]
== f"/api/hassio_ingress/{build_type[0]}"
)
assert aioclient_mock.mock_calls[-1][3]["X-Test-Header"] == "beer"
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_FOR]
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_HOST]
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_PROTO]
@pytest.mark.parametrize(
"build_type",
[
("a3_vl", "test/beer/ping?index=1"),
("core", "index.html"),
("local", "panel/config"),
("jk_921", "editor.php?idx=3&ping=5"),
("fsadjf10312", ""),
],
)
async def test_ingress_request_get_not_changed(
hassio_noauth_client, build_type, aioclient_mock: AiohttpClientMocker
) -> None:
"""Test ingress compressed and not modified."""
aioclient_mock.get(
f"http://127.0.0.1/ingress/{build_type[0]}/{build_type[1]}",
text="test",
status=HTTPStatus.NOT_MODIFIED,
)
resp = await hassio_noauth_client.get(
f"/api/hassio_ingress/{build_type[0]}/{build_type[1]}",
headers={"X-Test-Header": "beer", "Accept-Encoding": "gzip, deflate"},
)
# Check we got right response
assert resp.status == HTTPStatus.NOT_MODIFIED
body = await resp.text()
assert body == ""
assert "Content-Encoding" not in resp.headers # too small to compress
# Check we forwarded command
assert len(aioclient_mock.mock_calls) == 1
assert X_AUTH_TOKEN not in aioclient_mock.mock_calls[-1][3]
assert aioclient_mock.mock_calls[-1][3]["X-Hass-Source"] == "core.ingress"
assert (
aioclient_mock.mock_calls[-1][3]["X-Ingress-Path"]
== f"/api/hassio_ingress/{build_type[0]}"
)
assert aioclient_mock.mock_calls[-1][3]["X-Test-Header"] == "beer"
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_FOR]
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_HOST]
assert aioclient_mock.mock_calls[-1][3][X_FORWARDED_PROTO]