Skip to content

Commit bf43a3c

Browse files
api-clients-generation-pipeline[bot]ci.datadog-api-spec
andauthored
Regenerate client from commit d3abe21 of spec repo (#4067)
Co-authored-by: ci.datadog-api-spec <packages@datadoghq.com>
1 parent 486c268 commit bf43a3c

1 file changed

Lines changed: 119 additions & 11 deletions

File tree

‎tests/generated-test/test-server‎

Lines changed: 119 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,15 @@ SAFE_BROWSER_HEADERS = {
4949
}
5050
ZSTD_CONTENTSIZE_UNKNOWN = (1 << 64) - 1
5151
ZSTD_CONTENTSIZE_ERROR = (1 << 64) - 2
52+
ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
53+
ZSTD_SINGLE_SEGMENT_FLAG = 0x20
54+
ZSTD_DESCRIPTOR_LOW_BITS_MASK = 0x3F
55+
ZSTD_MIN_RAW_FRAME_SIZE = 6
56+
ZSTD_FCS_TWO_BYTE_OFFSET = 0x100
57+
UINT8_MAX = (1 << 8) - 1
58+
UINT16_WITH_OFFSET_MAX = (1 << 16) - 1 + ZSTD_FCS_TWO_BYTE_OFFSET
59+
UINT32_MAX = (1 << 32) - 1
60+
ZSTD_RAW_BLOCK_SIZE = 128 * 1024
5261
MAX_DECOMPRESSED_BODY_SIZE = 256 * 1024 * 1024
5362

5463

@@ -136,7 +145,7 @@ class RecordingDatabase:
136145
compression = (
137146
request_plan.get("compression")
138147
if request_plan and _request_matches_plan(expected["request"], request_plan)
139-
else None
148+
else expected["request"].get("compression")
140149
)
141150
return expected["response"], compression
142151

@@ -152,7 +161,7 @@ class RecordingDatabase:
152161
compression = (
153162
request_plan.get("compression")
154163
if request_plan and _request_matches_plan(interaction["request"], request_plan)
155-
else None
164+
else interaction["request"].get("compression")
156165
)
157166
return interaction["response"], compression
158167
raise LookupError("No unconsumed interaction matches this request")
@@ -387,17 +396,30 @@ class TestRequestHandler(BaseHTTPRequestHandler):
387396
else:
388397
body = body_data.get("value", "").encode("utf-8")
389398
status = response["status"]
399+
response_headers = response.get("headers", {})
400+
recorded_encoding = next(
401+
(str(value) for key, value in response_headers.items() if key.lower() == "content-encoding"),
402+
"",
403+
)
404+
recorded_compression = recorded_encoding.strip().casefold()
390405
if compression and _status_allows_message_content(status):
406+
if recorded_compression:
407+
decoded = _decompress_body(body, recorded_compression)
408+
if decoded is None:
409+
raise ValueError(f"Unsupported or invalid recorded response compression: {recorded_encoding}")
410+
body = decoded
391411
body = _compress_body(body, compression)
392412
self.send_response(status, response.get("reason"))
393-
for key, value in response.get("headers", {}).items():
413+
for key, value in response_headers.items():
394414
if (
395415
key.lower()
396416
not in HOP_BY_HOP_HEADERS | {"content-encoding", "content-length"} | SAFE_BROWSER_HEADERS.keys()
397417
):
398418
self.send_header(key, value)
399419
if compression and _status_allows_message_content(status):
400420
self.send_header("content-encoding", compression)
421+
elif recorded_encoding and _status_allows_message_content(status):
422+
self.send_header("content-encoding", recorded_encoding)
401423
for key, value in SAFE_BROWSER_HEADERS.items():
402424
self.send_header(key, value)
403425
if _status_allows_message_content(status):
@@ -513,17 +535,25 @@ def _compress_body(body: bytes, compression: str) -> bytes:
513535
if compression == "deflate":
514536
return zlib.compress(body)
515537
if compression == "zstd1":
516-
compressed = _compress_zstd(body)
517-
if compressed is not None:
518-
return compressed
538+
return _compress_zstd(body)
519539
raise ValueError(f"Unsupported response compression: {compression}")
520540

521541

522-
def _compress_zstd(body: bytes) -> bytes | None:
542+
def _decompress_body(body: bytes, compression: str) -> bytes | None:
543+
if compression == "gzip":
544+
return _decompress_zlib(body, wbits=zlib.MAX_WBITS | 16)
545+
if compression == "deflate":
546+
return _decompress_zlib(body, wbits=zlib.MAX_WBITS)
547+
if compression == "zstd1":
548+
return _decompress_zstd(body)
549+
return None
550+
551+
552+
def _compress_zstd(body: bytes) -> bytes:
523553
"""Compress one Zstandard frame with the platform libzstd."""
524554
library = _load_zstd()
525555
if library is None:
526-
return None
556+
return _compress_zstd_raw_frame(body)
527557
try:
528558
library.ZSTD_isError.argtypes = [ctypes.c_size_t]
529559
library.ZSTD_isError.restype = ctypes.c_uint
@@ -543,10 +573,41 @@ def _compress_zstd(body: bytes) -> bytes | None:
543573
destination = ctypes.create_string_buffer(max(1, capacity))
544574
compressed_size = library.ZSTD_compress(destination, capacity, source, len(body), 3)
545575
if library.ZSTD_isError(compressed_size):
546-
return None
576+
return _compress_zstd_raw_frame(body)
547577
return destination.raw[:compressed_size]
548578
except (AttributeError, OSError, OverflowError, TypeError):
549-
return None
579+
return _compress_zstd_raw_frame(body)
580+
581+
582+
def _compress_zstd_raw_frame(body: bytes) -> bytes:
583+
"""Create a valid Zstandard frame made only of dependency-free raw blocks."""
584+
size = len(body)
585+
if size <= UINT8_MAX:
586+
descriptor = 0x20
587+
content_size = size.to_bytes(1, "little")
588+
elif size <= UINT16_WITH_OFFSET_MAX:
589+
descriptor = 0x60
590+
content_size = (size - ZSTD_FCS_TWO_BYTE_OFFSET).to_bytes(2, "little")
591+
elif size <= UINT32_MAX:
592+
descriptor = 0xA0
593+
content_size = size.to_bytes(4, "little")
594+
else:
595+
descriptor = 0xE0
596+
content_size = size.to_bytes(8, "little")
597+
598+
frame = bytearray(ZSTD_MAGIC)
599+
frame.append(descriptor)
600+
frame.extend(content_size)
601+
offset = 0
602+
while offset < size:
603+
chunk = body[offset : offset + ZSTD_RAW_BLOCK_SIZE]
604+
offset += len(chunk)
605+
header = (len(chunk) << 3) | int(offset == size)
606+
frame.extend(header.to_bytes(3, "little"))
607+
frame.extend(chunk)
608+
if not body:
609+
frame.extend(b"\x01\x00\x00")
610+
return bytes(frame)
550611

551612

552613
def _load_zstd() -> Any | None:
@@ -571,7 +632,7 @@ def _decompress_zstd(body: bytes) -> bytes | None:
571632
"""Decompress one complete Zstandard frame with the platform libzstd."""
572633
library = _load_zstd()
573634
if library is None:
574-
return None
635+
return _decompress_zstd_raw_frame(body)
575636

576637
try:
577638
library.ZSTD_isError.argtypes = [ctypes.c_size_t]
@@ -608,6 +669,53 @@ def _decompress_zstd(body: bytes) -> bytes | None:
608669
return None
609670

610671

672+
def _decompress_zstd_raw_frame(body: bytes) -> bytes | None:
673+
"""Decode the raw-block Zstandard frames emitted by the dependency-free fallback."""
674+
if len(body) < ZSTD_MIN_RAW_FRAME_SIZE or body[:4] != ZSTD_MAGIC:
675+
return None
676+
descriptor = body[4]
677+
if descriptor & ZSTD_DESCRIPTOR_LOW_BITS_MASK != ZSTD_SINGLE_SEGMENT_FLAG:
678+
return None
679+
size_flag = descriptor >> 6
680+
size_bytes = (1, 2, 4, 8)[size_flag]
681+
cursor = 5
682+
if len(body) < cursor + size_bytes:
683+
return None
684+
expected_size = int.from_bytes(body[cursor : cursor + size_bytes], "little")
685+
if size_flag == 1:
686+
expected_size += ZSTD_FCS_TWO_BYTE_OFFSET
687+
if expected_size > MAX_DECOMPRESSED_BODY_SIZE:
688+
return None
689+
cursor += size_bytes
690+
output = bytearray()
691+
last_block = False
692+
while not last_block:
693+
if len(body) < cursor + 3:
694+
return None
695+
header = int.from_bytes(body[cursor : cursor + 3], "little")
696+
cursor += 3
697+
last_block = bool(header & 1)
698+
block_type = (header >> 1) & 0x3
699+
block_size = header >> 3
700+
if block_type == 0:
701+
if len(body) < cursor + block_size:
702+
return None
703+
output.extend(body[cursor : cursor + block_size])
704+
cursor += block_size
705+
elif block_type == 1:
706+
if cursor >= len(body):
707+
return None
708+
output.extend(body[cursor : cursor + 1] * block_size)
709+
cursor += 1
710+
else:
711+
return None
712+
if len(output) > MAX_DECOMPRESSED_BODY_SIZE:
713+
return None
714+
if cursor != len(body) or len(output) != expected_size:
715+
return None
716+
return bytes(output)
717+
718+
611719
def _normalise_json(value: Any) -> Any:
612720
if isinstance(value, dict):
613721
return {key: _normalise_json(item) for key, item in value.items()}

0 commit comments

Comments
 (0)