diff --git a/pyproject.toml b/pyproject.toml index 719857c61..0e7dd7e1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "aiohttp>=3.9.5,<3.11", "cryptography==42.0.8", "grpcio<1.54,>=1.53.2", + "multidict<7.0,>=4.5", "protobuf<4.22,>=4.21.12", "pydantic<3,>=2.6", "pydantic-settings<3,>=2.3", diff --git a/src/ota_proxy/_consts.py b/src/ota_proxy/_consts.py index d07769f03..6dcea5315 100644 --- a/src/ota_proxy/_consts.py +++ b/src/ota_proxy/_consts.py @@ -16,7 +16,9 @@ HEADER_AUTHORIZATION = istr("authorization") HEADER_COOKIE = istr("cookie") HEADER_CONTENT_ENCODING = istr("content-encoding") +HEADER_CONTENT_TYPE = istr("content-type") BHEADER_OTA_FILE_CACHE_CONTROL = OTAFileCacheControl.HEADER_LOWERCASE.encode("utf-8") BHEADER_AUTHORIZATION = b"authorization" BHEADER_COOKIE = b"cookie" BHEADER_CONTENT_ENCODING = b"content-encoding" +BHEADER_CONTENT_TYPE = b"content-type" diff --git a/src/ota_proxy/config.py b/src/ota_proxy/config.py index c6b571adb..e963ff1f3 100644 --- a/src/ota_proxy/config.py +++ b/src/ota_proxy/config.py @@ -16,8 +16,8 @@ class Config: BASE_DIR = "/ota-cache" CHUNK_SIZE = 1 * 1024 * 1024 # 4MB - DISK_USE_LIMIT_SOFT_P = 90 # in p% - DISK_USE_LIMIT_HARD_P = 92 # in p% + DISK_USE_LIMIT_SOFT_P = 80 # in p% + DISK_USE_LIMIT_HARD_P = 85 # in p% DISK_USE_PULL_INTERVAL = 2 # in seconds # value is the largest numbers of files that # might need to be deleted for the bucket to hold a new entry diff --git a/src/ota_proxy/db.py b/src/ota_proxy/db.py index 239c9bd51..1c5b1b08e 100644 --- a/src/ota_proxy/db.py +++ b/src/ota_proxy/db.py @@ -20,6 +20,7 @@ from pathlib import Path from typing import Optional +from multidict import CIMultiDict from pydantic import SkipValidation from simple_sqlite3_orm import ( ConstrainRepr, @@ -99,12 +100,12 @@ class CacheMeta(TableSpec): def __hash__(self) -> int: return hash(tuple(getattr(self, attrn) for attrn in self.model_fields)) - def export_headers_to_client(self) -> dict[str, str]: + def export_headers_to_client(self) -> CIMultiDict[str]: """Export required headers for client. Currently includes content-type, content-encoding and ota-file-cache-control headers. """ - res = {} + res = CIMultiDict() if self.content_encoding: res[HEADER_CONTENT_ENCODING] = self.content_encoding diff --git a/src/ota_proxy/ota_cache.py b/src/ota_proxy/ota_cache.py index 531525cd9..641f381c8 100644 --- a/src/ota_proxy/ota_cache.py +++ b/src/ota_proxy/ota_cache.py @@ -26,7 +26,7 @@ from urllib.parse import SplitResult, quote, urlsplit import aiohttp -from multidict import CIMultiDictProxy +from multidict import CIMultiDict, CIMultiDictProxy from otaclient_common.common import get_backoff from otaclient_common.typing import StrOrPath @@ -394,7 +394,7 @@ async def _do_request() -> AsyncIterator[bytes]: async def _retrieve_file_by_cache( self, cache_identifier: str, *, retry_cache: bool - ) -> Optional[Tuple[AsyncIterator[bytes], Mapping[str, str]]]: + ) -> tuple[AsyncIterator[bytes], CIMultiDict[str]] | None: """ Returns: A tuple of bytes iterator and headers dict for back to client. @@ -443,7 +443,7 @@ async def _retrieve_file_by_cache( async def _retrieve_file_by_external_cache( self, client_cache_policy: OTAFileCacheControl - ) -> Optional[Tuple[AsyncIterator[bytes], Mapping[str, str]]]: + ) -> tuple[AsyncIterator[bytes], CIMultiDict[str]] | None: # skip if not external cache or otaclient doesn't sent valid file_sha256 if not self._external_cache or not client_cache_policy.file_sha256: return @@ -455,18 +455,23 @@ async def _retrieve_file_by_external_cache( ) if cache_file_zst.is_file(): - return read_file(cache_file_zst, executor=self._executor), { - HEADER_OTA_FILE_CACHE_CONTROL: OTAFileCacheControl.export_kwargs_as_header( + _header = CIMultiDict() + _header[HEADER_OTA_FILE_CACHE_CONTROL] = ( + OTAFileCacheControl.export_kwargs_as_header( file_sha256=cache_identifier, file_compression_alg=cfg.EXTERNAL_CACHE_STORAGE_COMPRESS_ALG, ) - } - elif cache_file.is_file(): - return read_file(cache_file, executor=self._executor), { - HEADER_OTA_FILE_CACHE_CONTROL: OTAFileCacheControl.export_kwargs_as_header( + ) + return read_file(cache_file_zst, executor=self._executor), _header + + if cache_file.is_file(): + _header = CIMultiDict() + _header[HEADER_OTA_FILE_CACHE_CONTROL] = ( + OTAFileCacheControl.export_kwargs_as_header( file_sha256=cache_identifier ) - } + ) + return read_file(cache_file, executor=self._executor), _header # exposed API @@ -474,7 +479,7 @@ async def retrieve_file( self, raw_url: str, headers_from_client: Dict[str, str], - ) -> Optional[Tuple[AsyncIterator[bytes], Mapping[str, str]]]: + ) -> tuple[AsyncIterator[bytes], CIMultiDict[str] | CIMultiDictProxy[str]] | None: """Retrieve a file descriptor for the requested . This method retrieves a file descriptor for incoming client request. diff --git a/src/ota_proxy/server_app.py b/src/ota_proxy/server_app.py index d1f815b30..bd1999a00 100644 --- a/src/ota_proxy/server_app.py +++ b/src/ota_proxy/server_app.py @@ -17,20 +17,23 @@ import logging from contextlib import asynccontextmanager from http import HTTPStatus -from typing import Dict, List, Mapping, Tuple, Union +from typing import Dict, List, Tuple, Union from urllib.parse import urlparse import aiohttp +from multidict import CIMultiDict, CIMultiDictProxy from otaclient_common.logging import BurstSuppressFilter from ._consts import ( BHEADER_AUTHORIZATION, BHEADER_CONTENT_ENCODING, + BHEADER_CONTENT_TYPE, BHEADER_COOKIE, BHEADER_OTA_FILE_CACHE_CONTROL, HEADER_AUTHORIZATION, HEADER_CONTENT_ENCODING, + HEADER_CONTENT_TYPE, HEADER_COOKIE, HEADER_OTA_FILE_CACHE_CONTROL, METHOD_GET, @@ -88,21 +91,26 @@ def parse_raw_headers(raw_headers: List[Tuple[bytes, bytes]]) -> Dict[str, str]: return headers -def encode_headers(headers: Mapping[str, str]) -> List[Tuple[bytes, bytes]]: +def encode_headers( + headers: Union[CIMultiDict[str], CIMultiDictProxy[str]] +) -> List[Tuple[bytes, bytes]]: """Encode headers dict to list of bytes tuples for sending back to client. Uvicorn requests application to pre-process headers to bytes. - Currently we only need to send content-encoding and ota-file-cache-control header - back to client. + Currently we send the following headers back to client: + 1. content-encoding + 2. content-type + 3. ota-file-cache-control header """ bytes_headers: List[Tuple[bytes, bytes]] = [] - for name, value in headers.items(): - if name == HEADER_CONTENT_ENCODING and value: - bytes_headers.append((BHEADER_CONTENT_ENCODING, value.encode("utf-8"))) - elif name == HEADER_OTA_FILE_CACHE_CONTROL and value: - bytes_headers.append( - (BHEADER_OTA_FILE_CACHE_CONTROL, value.encode("utf-8")) - ) + if _encoding := headers.get(HEADER_CONTENT_ENCODING): + bytes_headers.append((BHEADER_CONTENT_ENCODING, _encoding.encode("utf-8"))) + if _type := headers.get(HEADER_CONTENT_TYPE): + bytes_headers.append((BHEADER_CONTENT_TYPE, _type.encode("utf-8"))) + if _cache_control := headers.get(HEADER_OTA_FILE_CACHE_CONTROL): + bytes_headers.append( + (BHEADER_OTA_FILE_CACHE_CONTROL, _cache_control.encode("utf-8")) + ) return bytes_headers