fix(ota_proxy): fix content-encoding header is not passed through in …

…direct proxy mode (#445)

This PR fixes a bug that prevents ota_proxy properly handling request of file being gzip by remote HTTP server. The cause of bug is ota_proxy doesn't include content-encoding header in response back to client when in direct proxy passthrough mode.

The bug is fixed by using multidict when parsing/preparing headers, so multidict<7.0,>=4.5 is added to the direct dep of otaclient. Note that multidict is used by aiohttp, so it is already a indirect dep of otaclient.

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
   "aiohttp>=3.10.11,<3.12",
   "cryptography>=43.0.1,<45",
   "grpcio>=1.53.2,<1.69",
+  "multidict<7.0,>=4.5",
   "protobuf>=4.21.12,<5.29",
   "pydantic<3,>=2.6",
   "pydantic-settings<3,>=2.3",

requirements.txt

@@ -4,6 +4,7 @@ aiofiles<25,>=24.1
 aiohttp>=3.10.11,<3.12
 cryptography>=43.0.1,<45
 grpcio>=1.53.2,<1.69
+multidict<7.0,>=4.5
 protobuf>=4.21.12,<5.29
 pydantic<3,>=2.6
 pydantic-settings<3,>=2.3

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"

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
 

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
@@ -409,7 +409,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.
@@ -458,7 +458,8 @@ async def _retrieve_file_by_cache(
 
     async def _retrieve_file_by_external_cache(
         self, client_cache_policy: OTAFileCacheControl
-    ) -> tuple[AsyncIterator[bytes], Mapping[str, str]] | None:
+    ) -> tuple[AsyncIterator[bytes], CIMultiDict[str]] | None:
+        # skip if not external cache or otaclient doesn't sent valid file_sha256
         if not self._external_cache_data_dir or not client_cache_policy.file_sha256:
             return
 
@@ -469,26 +470,31 @@ 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
 
     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 <raw_url>.
 
         This method retrieves a file descriptor for incoming client request.

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