@@ -49,6 +49,15 @@ SAFE_BROWSER_HEADERS = {
4949}
5050ZSTD_CONTENTSIZE_UNKNOWN = (1 << 64 ) - 1
5151ZSTD_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
5261MAX_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
552613def _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+
611719def _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