Skip to content

Commit 8227c5c

Browse files
SudipSinhaclaude
andcommitted
fix: two-layer gzip decompression for all inbound routes
Layer 1 (transport): Change GzipRequestMiddleware.DEFAULT_PATHS from ("/data/upload",) to ("*",) so gzip decompression covers all HTTP routes including /consumer/kserve/v2. Content-type gating and decompression bomb protection remain intact. Layer 2 (application): Add decompress_if_gzip() for the CloudEvent endpoint (POST /) where Knative Eventing strips the Content-Encoding header while leaving the body gzip-compressed. Detects gzip by magic bytes (0x1F 0x8B) and decompresses before JSON parsing, matching the Java service's CloudEventConsumer.decompressIfGzip() fix. Extract process_cloud_event() from the HTTP endpoint so the upload endpoint's internal forwarding path (data_upload.py) can call it directly without going through HTTP request parsing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 715269c commit 8227c5c

7 files changed

Lines changed: 245 additions & 24 deletions

File tree

src/endpoints/consumer/consumer_endpoint.py

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,17 @@
1010
from typing import Annotated, Never
1111

1212
import numpy as np
13-
from fastapi import APIRouter, Header, HTTPException
13+
from fastapi import APIRouter, Header, HTTPException, Request
1414
from numpy import ndarray
15+
from pydantic import TypeAdapter, ValidationError
1516

1617
from src.endpoints.consumer import (
1718
InferencePartialPayload,
1819
KServeData,
1920
KServeInferenceRequest,
2021
KServeInferenceResponse,
2122
)
23+
from src.endpoints.consumer.gzip_utils import decompress_if_gzip
2224
from src.exceptions import ReconciliationError
2325
from src.service.data.datasources.data_source import DataSource
2426

@@ -456,23 +458,24 @@ def process_payload(
456458
return np.array(kserve_data.data), column_names
457459

458460

459-
@router.post("/")
460-
async def consume_cloud_event(
461+
_kserve_payload_adapter = TypeAdapter(KServeInferenceRequest | KServeInferenceResponse)
462+
463+
464+
async def process_cloud_event(
461465
payload: KServeInferenceRequest | KServeInferenceResponse,
462-
ce_id: Annotated[str | None, Header()] = None,
466+
ce_id: str | None = None,
463467
tag: str | None = None,
464468
) -> dict[str, str]:
465-
"""Consume KServe v2 payloads from cloud events.
469+
"""Process a KServe payload from a cloud event or internal call.
466470
467-
This endpoint accepts both input (request) and output (response) payloads
468-
from ModelMesh-served models and stores them for reconciliation.
471+
This is the core logic shared by the HTTP endpoint and the upload
472+
endpoint's internal forwarding path.
469473
470-
:param payload: KServe inference request or response
471-
:param ce_id: Cloud event ID from header
474+
:param payload: Parsed KServe inference request or response
475+
:param ce_id: Cloud event ID from header (overrides payload.id)
472476
:param tag: Optional tag to associate with the data
473477
:raises HTTPException: If payload processing fails
474478
"""
475-
# set payload id from cloud event header if present
476479
if ce_id is not None:
477480
payload.id = ce_id
478481

@@ -482,7 +485,6 @@ async def consume_cloud_event(
482485
detail="Payload requires 'id' field or 'ce-id' header",
483486
)
484487

485-
# get global storage interface
486488
storage_interface = get_global_storage_interface()
487489

488490
try:
@@ -494,13 +496,11 @@ async def consume_cloud_event(
494496
)
495497
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail=msg)
496498
logger.info("KServe Inference Input %s received.", payload.id)
497-
# if a match is found, the payload is auto-deleted from data
498499
partial_output = await storage_interface.get_partial_payload(
499500
payload.id, is_input=False, is_modelmesh=False
500501
)
501502
if partial_output is not None:
502503
if not isinstance(partial_output, KServeInferenceResponse):
503-
# This should never happen - indicates storage interface error
504504
raise HTTPException(
505505
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
506506
detail="Invalid payload type from storage",
@@ -532,7 +532,6 @@ async def consume_cloud_event(
532532
)
533533
if partial_input is not None:
534534
if not isinstance(partial_input, KServeInferenceRequest):
535-
# This should never happen - indicates storage interface error
536535
raise HTTPException(
537536
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
538537
detail="Invalid payload type from storage",
@@ -548,8 +547,6 @@ async def consume_cloud_event(
548547
"message": f"Output payload {payload.id} processed successfully",
549548
}
550549

551-
# Defensive programming: this should never happen due to type annotation
552-
# but adding explicit fallback for type safety
553550
raise HTTPException(
554551
status_code=HTTPStatus.BAD_REQUEST,
555552
detail="Payload must be either KServeInferenceRequest or KServeInferenceResponse",
@@ -558,3 +555,34 @@ async def consume_cloud_event(
558555
except ReconciliationError as e:
559556
logger.exception("Reconciliation failed for payload %s", payload.id)
560557
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail=str(e)) from e
558+
559+
560+
@router.post("/")
561+
async def consume_cloud_event(
562+
http_request: Request,
563+
ce_id: Annotated[str | None, Header()] = None,
564+
tag: str | None = None,
565+
) -> dict[str, str]:
566+
"""Consume KServe v2 payloads from cloud events.
567+
568+
Knative Eventing may strip the Content-Encoding header while leaving the
569+
body gzip-compressed, so this endpoint detects gzip by magic bytes and
570+
decompresses before JSON parsing.
571+
572+
:param http_request: Raw HTTP request (body may be gzip-compressed without header)
573+
:param ce_id: Cloud event ID from header
574+
:param tag: Optional tag to associate with the data
575+
:raises HTTPException: If payload processing fails
576+
"""
577+
raw_body = await http_request.body()
578+
body = decompress_if_gzip(raw_body)
579+
580+
try:
581+
payload = _kserve_payload_adapter.validate_json(body)
582+
except ValidationError as e:
583+
raise HTTPException(
584+
status_code=HTTPStatus.BAD_REQUEST,
585+
detail=f"Invalid payload: {e}",
586+
) from e
587+
588+
return await process_cloud_event(payload, ce_id=ce_id, tag=tag)
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
"""Gzip magic-byte detection for CloudEvent payloads.
2+
3+
Knative Eventing reconstructs HTTP requests via the CloudEvents SDK,
4+
which drops transport headers like Content-Encoding while leaving the
5+
body gzip-compressed. This module detects gzip by magic bytes (0x1F 0x8B)
6+
and decompresses at the application layer — the same fix applied to the
7+
Java service's CloudEventConsumer.decompressIfGzip().
8+
"""
9+
10+
import gzip
11+
import logging
12+
from io import BytesIO
13+
14+
from src.middleware.gzip_middleware import GzipRequestMiddleware
15+
16+
logger = logging.getLogger(__name__)
17+
18+
_GZIP_MAGIC = b"\x1f\x8b"
19+
_GZIP_MAGIC_LEN = len(_GZIP_MAGIC)
20+
_CHUNK_SIZE = 64 * 1024 # 64KB streaming chunks
21+
DEFAULT_MAX_DECOMPRESSED_SIZE = GzipRequestMiddleware.DEFAULT_MAX_SIZE
22+
23+
24+
def decompress_if_gzip(
25+
data: bytes,
26+
max_size: int = DEFAULT_MAX_DECOMPRESSED_SIZE,
27+
) -> bytes:
28+
"""Decompress data if it starts with gzip magic bytes.
29+
30+
Returns the original data unchanged if it is not gzip-compressed
31+
or if decompression fails.
32+
33+
:param data: Raw bytes to check and potentially decompress
34+
:param max_size: Maximum allowed decompressed size in bytes
35+
:return: Decompressed bytes, or original data if not gzip
36+
"""
37+
if len(data) < _GZIP_MAGIC_LEN or data[:_GZIP_MAGIC_LEN] != _GZIP_MAGIC:
38+
return data
39+
40+
try:
41+
decompressed = bytearray()
42+
with BytesIO(data) as bio, gzip.GzipFile(fileobj=bio) as gz:
43+
while True:
44+
chunk = gz.read(_CHUNK_SIZE)
45+
if not chunk:
46+
break
47+
if len(decompressed) + len(chunk) > max_size:
48+
msg = f"Decompressed CloudEvent payload exceeds {max_size} bytes"
49+
raise ValueError(msg)
50+
decompressed.extend(chunk)
51+
52+
logger.debug("Decompressed gzip CloudEvent payload")
53+
return bytes(decompressed)
54+
55+
except (gzip.BadGzipFile, OSError):
56+
logger.warning(
57+
"CloudEvent payload starts with gzip magic bytes but failed to decompress, using raw bytes",
58+
)
59+
return data

src/endpoints/data/data_upload.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from pydantic import BaseModel
99

1010
from src.endpoints.consumer import KServeInferenceRequest, KServeInferenceResponse
11-
from src.endpoints.consumer.consumer_endpoint import consume_cloud_event
11+
from src.endpoints.consumer.consumer_endpoint import process_cloud_event
1212
from src.exceptions import ReconciliationError
1313
from src.service.constants import TRUSTYAI_TAG_PREFIX
1414
from src.service.data.model_data import ModelData
@@ -81,8 +81,8 @@ async def upload(payload: UploadPayload) -> dict[str, str]:
8181
else:
8282
previous_data_points = 0
8383

84-
await consume_cloud_event(payload.response, req_id)
85-
await consume_cloud_event(payload.request, req_id, tag=payload.data_tag)
84+
await process_cloud_event(payload.response, req_id)
85+
await process_cloud_event(payload.request, req_id, tag=payload.data_tag)
8686

8787
model_data = ModelData(payload.model_name)
8888
new_data_points = (await model_data.row_counts())[0]

src/middleware/gzip_middleware.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,11 @@ class GzipRequestMiddleware:
3636
removing the Content-Encoding header, and updating Content-Length.
3737
Includes protection against decompression bombs via max_size limit.
3838
39-
Defaults: paths=["/data/upload"], max_size=16MB, fail_on_error=True
39+
Defaults: paths=["*"] (all paths), max_size=16MB, fail_on_error=True
4040
"""
4141

4242
# Default configuration constants
43-
DEFAULT_PATHS = ("/data/upload",) # Tuple to avoid mutable default
43+
DEFAULT_PATHS = ("*",) # All paths: Content-Encoding is a transport-level concern
4444
DEFAULT_ALLOWED_CONTENT_TYPES = (
4545
"application/json",
4646
"application/cloudevents+json",
@@ -86,7 +86,7 @@ def __init__(
8686
8787
Args:
8888
app: ASGI application
89-
paths: Path patterns to apply (supports wildcards, default: ["/data/upload"])
89+
paths: Path patterns to apply (supports wildcards, default: ["*"] = all paths)
9090
max_size: Max decompressed bytes (default: 16MB)
9191
fail_on_error: Return error on failure vs pass through (default: True)
9292
allowed_content_types: Eligible content types
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
"""Tests for consumer endpoints."""
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
"""Tests for Layer 2 gzip decompression in the CloudEvent consumer endpoint.
2+
3+
Knative Eventing strips Content-Encoding headers while leaving the body
4+
gzip-compressed. These tests verify that the CloudEvent endpoint detects
5+
gzip by magic bytes and decompresses before JSON parsing.
6+
"""
7+
8+
import gzip
9+
import json
10+
11+
import pytest
12+
13+
from src.endpoints.consumer.gzip_utils import decompress_if_gzip
14+
15+
16+
class TestDecompressIfGzip:
17+
"""Unit tests for the decompress_if_gzip utility."""
18+
19+
def test_decompresses_gzip_data(self) -> None:
20+
"""Valid gzip data is decompressed."""
21+
original = b'{"inputs": [{"name": "x", "shape": [1], "datatype": "FP32", "data": [1.0]}]}'
22+
compressed = gzip.compress(original)
23+
24+
result = decompress_if_gzip(compressed)
25+
26+
assert result == original
27+
28+
def test_returns_non_gzip_unchanged(self) -> None:
29+
"""Non-gzip data is returned unchanged."""
30+
data = b'{"inputs": [{"name": "x", "shape": [1], "datatype": "FP32", "data": [1.0]}]}'
31+
32+
result = decompress_if_gzip(data)
33+
34+
assert result is data
35+
36+
def test_returns_empty_bytes_unchanged(self) -> None:
37+
"""Empty bytes are returned unchanged."""
38+
result = decompress_if_gzip(b"")
39+
40+
assert result == b""
41+
42+
def test_returns_single_byte_unchanged(self) -> None:
43+
"""Single byte (too short for magic check) is returned unchanged."""
44+
result = decompress_if_gzip(b"\x1f")
45+
46+
assert result == b"\x1f"
47+
48+
def test_invalid_gzip_with_magic_bytes_returns_original(self) -> None:
49+
"""Data starting with gzip magic but not valid gzip returns original."""
50+
fake_gzip = b"\x1f\x8b\x00\x00invalid"
51+
52+
result = decompress_if_gzip(fake_gzip)
53+
54+
assert result == fake_gzip
55+
56+
def test_size_limit_raises_on_decompression_bomb(self) -> None:
57+
"""Exceeding max_size raises ValueError."""
58+
large_data = b"x" * 10_000
59+
compressed = gzip.compress(large_data)
60+
61+
with pytest.raises(ValueError, match="exceeds"):
62+
decompress_if_gzip(compressed, max_size=100)
63+
64+
def test_size_limit_within_bounds_succeeds(self) -> None:
65+
"""Data within max_size decompresses successfully."""
66+
data = b'{"test": true}'
67+
compressed = gzip.compress(data)
68+
69+
result = decompress_if_gzip(compressed, max_size=1024)
70+
71+
assert result == data
72+
73+
def test_preserves_json_fidelity(self) -> None:
74+
"""Decompressed JSON round-trips correctly."""
75+
payload = {
76+
"model_name": "example",
77+
"id": "req-001",
78+
"outputs": [
79+
{
80+
"name": "predict",
81+
"shape": [2, 1],
82+
"datatype": "FP64",
83+
"data": [[0.1], [0.9]],
84+
},
85+
],
86+
}
87+
original = json.dumps(payload).encode()
88+
compressed = gzip.compress(original)
89+
90+
result = decompress_if_gzip(compressed)
91+
92+
assert json.loads(result) == payload

0 commit comments

Comments
 (0)