Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/compression_utils.py: 63%

1import asyncio 

2import sys 

3import zlib 

4from abc import ABC, abstractmethod 

5from concurrent.futures import Executor 

6from typing import Any, Final, Generic, Protocol, TypedDict, TypeVar, cast 

7 

8if sys.version_info >= (3, 12): 

9 from collections.abc import Buffer 

10else: 

11 from typing import Union 

12 

13 Buffer = Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"] 

14 

15try: 

16 try: 

17 import brotlicffi as brotli 

18 except ImportError: 

19 import brotli 

20 

21 HAS_BROTLI = True 

22except ImportError: 

23 HAS_BROTLI = False 

24 

25try: 

26 if sys.version_info >= (3, 14): 

27 from compression.zstd import ZstdDecompressor # noqa: I900 

28 else: # TODO(PY314): Remove mentions of backports.zstd across codebase 

29 from backports.zstd import ZstdDecompressor 

30 

31 HAS_ZSTD = True 

32except ImportError: 

33 HAS_ZSTD = False 

34 

35 

36MAX_SYNC_CHUNK_SIZE = 4096 

37 

38# Unlimited decompression constants - different libraries use different conventions 

39ZLIB_MAX_LENGTH_UNLIMITED = 0 # zlib uses 0 to mean unlimited 

40ZSTD_MAX_LENGTH_UNLIMITED = -1 # zstd uses -1 to mean unlimited 

41 

42# Concatenated members are decoded through a window that starts small and 

43# doubles. A fresh decompressor copies everything past the member it decodes 

44# into unused_data, so handing it the whole remaining buffer at every boundary 

45# is quadratic over a stream of small members. 

46MEMBER_WINDOW_MIN = 64 

47MEMBER_WINDOW_MAX = 65536 

48 

49# Cap on concatenated members decoded in one call. Real payloads are unlikely 

50# to have more than a few members. 

51MAX_DECOMPRESS_MEMBERS = 1024 

52 

53 

54class TooManyMembersError(ValueError): 

55 """A stream concatenated more members than the caller allows.""" 

56 

57 

58class ZLibCompressObjProtocol(Protocol): 

59 def compress(self, data: Buffer) -> bytes: ... 

60 def flush(self, mode: int = ..., /) -> bytes: ... 

61 

62 

63class ZLibDecompressObjProtocol(Protocol): 

64 def decompress(self, data: Buffer, max_length: int = ...) -> bytes: ... 

65 def flush(self, length: int = ..., /) -> bytes: ... 

66 

67 @property 

68 def eof(self) -> bool: ... 

69 

70 @property 

71 def unconsumed_tail(self) -> bytes: ... 

72 

73 @property 

74 def unused_data(self) -> bytes: ... 

75 

76 

77class ZLibBackendProtocol(Protocol): 

78 MAX_WBITS: int 

79 Z_FULL_FLUSH: int 

80 Z_SYNC_FLUSH: int 

81 Z_BEST_SPEED: int 

82 Z_FINISH: int 

83 

84 def compressobj( 

85 self, 

86 level: int = ..., 

87 method: int = ..., 

88 wbits: int = ..., 

89 memLevel: int = ..., 

90 strategy: int = ..., 

91 zdict: Buffer | None = ..., 

92 ) -> ZLibCompressObjProtocol: ... 

93 def decompressobj( 

94 self, wbits: int = ..., zdict: Buffer = ... 

95 ) -> ZLibDecompressObjProtocol: ... 

96 

97 def compress( 

98 self, data: Buffer, /, level: int = ..., wbits: int = ... 

99 ) -> bytes: ... 

100 def decompress( 

101 self, data: Buffer, /, wbits: int = ..., bufsize: int = ... 

102 ) -> bytes: ... 

103 

104 

105class CompressObjArgs(TypedDict, total=False): 

106 wbits: int 

107 strategy: int 

108 level: int 

109 

110 

111class ZLibBackendWrapper: 

112 def __init__(self, _zlib_backend: ZLibBackendProtocol): 

113 self._zlib_backend: ZLibBackendProtocol = _zlib_backend 

114 

115 @property 

116 def name(self) -> str: 

117 return getattr(self._zlib_backend, "__name__", "undefined") 

118 

119 @property 

120 def MAX_WBITS(self) -> int: 

121 return self._zlib_backend.MAX_WBITS 

122 

123 @property 

124 def Z_FULL_FLUSH(self) -> int: 

125 return self._zlib_backend.Z_FULL_FLUSH 

126 

127 @property 

128 def Z_SYNC_FLUSH(self) -> int: 

129 return self._zlib_backend.Z_SYNC_FLUSH 

130 

131 @property 

132 def Z_BEST_SPEED(self) -> int: 

133 return self._zlib_backend.Z_BEST_SPEED 

134 

135 @property 

136 def Z_FINISH(self) -> int: 

137 return self._zlib_backend.Z_FINISH 

138 

139 def compressobj(self, *args: Any, **kwargs: Any) -> ZLibCompressObjProtocol: 

140 return self._zlib_backend.compressobj(*args, **kwargs) 

141 

142 def decompressobj(self, *args: Any, **kwargs: Any) -> ZLibDecompressObjProtocol: 

143 return self._zlib_backend.decompressobj(*args, **kwargs) 

144 

145 def compress(self, data: Buffer, *args: Any, **kwargs: Any) -> bytes: 

146 return self._zlib_backend.compress(data, *args, **kwargs) 

147 

148 def decompress(self, data: Buffer, *args: Any, **kwargs: Any) -> bytes: 

149 return self._zlib_backend.decompress(data, *args, **kwargs) 

150 

151 # Everything not explicitly listed in the Protocol we just pass through 

152 def __getattr__(self, attrname: str) -> Any: 

153 return getattr(self._zlib_backend, attrname) 

154 

155 

156ZLibBackend: ZLibBackendWrapper = ZLibBackendWrapper(zlib) 

157 

158 

159def set_zlib_backend(new_zlib_backend: ZLibBackendProtocol) -> None: 

160 ZLibBackend._zlib_backend = new_zlib_backend 

161 

162 

163def encoding_to_mode( 

164 encoding: str | None = None, 

165 suppress_deflate_header: bool = False, 

166) -> int: 

167 if encoding == "gzip": 

168 return 16 + ZLibBackend.MAX_WBITS 

169 

170 return -ZLibBackend.MAX_WBITS if suppress_deflate_header else ZLibBackend.MAX_WBITS 

171 

172 

173class MemberDecompressObjProtocol(Protocol): 

174 def decompress(self, data: Buffer, max_length: int = ...) -> bytes: ... 

175 

176 @property 

177 def eof(self) -> bool: ... 

178 

179 @property 

180 def unused_data(self) -> bytes: ... 

181 

182 

183_DecompressObjT = TypeVar("_DecompressObjT", bound=MemberDecompressObjProtocol) 

184 

185 

186class DecompressionBaseHandler(ABC): 

187 def __init__( 

188 self, 

189 executor: Executor | None = None, 

190 max_sync_chunk_size: int | None = MAX_SYNC_CHUNK_SIZE, 

191 ): 

192 """Base class for decompression handlers.""" 

193 self._executor = executor 

194 self._max_sync_chunk_size = max_sync_chunk_size 

195 

196 @abstractmethod 

197 def decompress_sync( 

198 self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED 

199 ) -> bytes: 

200 """Decompress the given data.""" 

201 

202 async def decompress( 

203 self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED 

204 ) -> bytes: 

205 """Decompress the given data.""" 

206 if ( 

207 self._max_sync_chunk_size is not None 

208 and len(data) > self._max_sync_chunk_size 

209 ): 

210 return await asyncio.get_event_loop().run_in_executor( 

211 self._executor, self.decompress_sync, data, max_length 

212 ) 

213 return self.decompress_sync(data, max_length) 

214 

215 @property 

216 @abstractmethod 

217 def data_available(self) -> bool: 

218 """Return True if more output is available by passing b"".""" 

219 

220 

221class ConcatDecompressionHandler(DecompressionBaseHandler, Generic[_DecompressObjT]): 

222 """Handler for a codec whose streams may concatenate independent members. 

223 

224 Concatenated gzip/deflate members and multi-frame zstd 

225 (https://datatracker.ietf.org/doc/html/rfc8878#section-3.1.1) decode the 

226 same way: a decompressor handles one member, then flags eof and leaves the 

227 rest of the input in unused_data, so every member after it needs a fresh 

228 one. 

229 """ 

230 

231 # Sentinel this codec's decompress() takes to mean "no output limit". 

232 _unlimited: int 

233 _decompressor: _DecompressObjT 

234 # Input a max_length-capped walk stopped short of, fed back on the next call. 

235 _pending_unused_data: bytes | None = None 

236 

237 @abstractmethod 

238 def _new_decompressor(self) -> _DecompressObjT: 

239 """Return a decompressor for the next member.""" 

240 

241 def _decompress_members(self, first: bytes, max_length: int) -> bytes: 

242 """Decode the members following the one ``first`` came from.""" 

243 remaining = memoryview(self._decompressor.unused_data) 

244 parts = [first] 

245 produced = len(first) 

246 pos = 0 

247 window = MEMBER_WINDOW_MIN 

248 budget = max_length 

249 members = 1 

250 

251 while pos < len(remaining): 

252 if self._decompressor.eof: 

253 members += 1 

254 if members > MAX_DECOMPRESS_MEMBERS: 

255 raise TooManyMembersError( 

256 f"Compressed stream has more than " 

257 f"{MAX_DECOMPRESS_MEMBERS} members" 

258 ) 

259 # Replace the spent decompressor before the budget check below 

260 # can break out of the loop: it still lists these bytes in its 

261 # unused_data and would hand them back on the next call. 

262 self._decompressor = self._new_decompressor() 

263 window = MEMBER_WINDOW_MIN 

264 if max_length != self._unlimited: 

265 budget = max_length - produced 

266 if budget <= 0: 

267 self._pending_unused_data = bytes(remaining[pos:]) 

268 break 

269 

270 end = min(pos + window, len(remaining)) 

271 chunk = self._decompressor.decompress(remaining[pos:end], budget) 

272 if chunk: 

273 parts.append(chunk) 

274 produced += len(chunk) 

275 

276 if self._decompressor.eof: 

277 pos = end - len(self._decompressor.unused_data) 

278 else: 

279 pos = end 

280 # Doubling the window on each iteration avoids too many calls 

281 # when a large member is present, while protecting us from 

282 # quadratic usage when members are of window+1 length. 

283 window = min(window * 2, MEMBER_WINDOW_MAX) 

284 

285 return b"".join(parts) 

286 

287 

288class ZLibCompressor: 

289 def __init__( 

290 self, 

291 encoding: str | None = None, 

292 suppress_deflate_header: bool = False, 

293 level: int | None = None, 

294 wbits: int | None = None, 

295 strategy: int | None = None, 

296 executor: Executor | None = None, 

297 max_sync_chunk_size: int | None = MAX_SYNC_CHUNK_SIZE, 

298 ): 

299 self._executor = executor 

300 self._max_sync_chunk_size = max_sync_chunk_size 

301 self._mode = ( 

302 encoding_to_mode(encoding, suppress_deflate_header) 

303 if wbits is None 

304 else wbits 

305 ) 

306 self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend) 

307 

308 kwargs: CompressObjArgs = {} 

309 kwargs["wbits"] = self._mode 

310 if strategy is not None: 

311 kwargs["strategy"] = strategy 

312 if level is not None: 

313 kwargs["level"] = level 

314 self._compressor = self._zlib_backend.compressobj(**kwargs) 

315 

316 def compress_sync(self, data: Buffer) -> bytes: 

317 return self._compressor.compress(data) 

318 

319 async def compress(self, data: Buffer) -> bytes: 

320 """Compress the data and returned the compressed bytes. 

321 

322 Note that flush() must be called after the last call to compress() 

323 

324 If the data size is large than the max_sync_chunk_size, the compression 

325 will be done in the executor. Otherwise, the compression will be done 

326 in the event loop. 

327 

328 **WARNING: This method is NOT cancellation-safe when used with flush().** 

329 If this operation is cancelled, the compressor state may be corrupted. 

330 The connection MUST be closed after cancellation to avoid data corruption 

331 in subsequent compress operations. 

332 

333 For cancellation-safe compression (e.g., WebSocket), the caller MUST wrap 

334 compress() + flush() + send operations in a shield and lock to ensure atomicity. 

335 """ 

336 # For large payloads, offload compression to executor to avoid blocking event loop 

337 should_use_executor = ( 

338 self._max_sync_chunk_size is not None 

339 and len(data) > self._max_sync_chunk_size 

340 ) 

341 if should_use_executor: 

342 return await asyncio.get_running_loop().run_in_executor( 

343 self._executor, self._compressor.compress, data 

344 ) 

345 return self.compress_sync(data) 

346 

347 def flush(self, mode: int | None = None) -> bytes: 

348 """Flush the compressor synchronously. 

349 

350 **WARNING: This method is NOT cancellation-safe when called after compress().** 

351 The flush() operation accesses shared compressor state. If compress() was 

352 cancelled, calling flush() may result in corrupted data. The connection MUST 

353 be closed after compress() cancellation. 

354 

355 For cancellation-safe compression (e.g., WebSocket), the caller MUST wrap 

356 compress() + flush() + send operations in a shield and lock to ensure atomicity. 

357 """ 

358 return self._compressor.flush( 

359 mode if mode is not None else self._zlib_backend.Z_FINISH 

360 ) 

361 

362 

363class ZLibDecompressor(ConcatDecompressionHandler[ZLibDecompressObjProtocol]): 

364 _unlimited = ZLIB_MAX_LENGTH_UNLIMITED 

365 

366 def __init__( 

367 self, 

368 encoding: str | None = None, 

369 suppress_deflate_header: bool = False, 

370 executor: Executor | None = None, 

371 max_sync_chunk_size: int | None = MAX_SYNC_CHUNK_SIZE, 

372 ): 

373 super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) 

374 self._mode = encoding_to_mode(encoding, suppress_deflate_header) 

375 self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend) 

376 self._decompressor = self._new_decompressor() 

377 self._last_empty = False 

378 

379 def _new_decompressor(self) -> ZLibDecompressObjProtocol: 

380 return self._zlib_backend.decompressobj(wbits=self._mode) 

381 

382 def decompress_sync( 

383 self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED 

384 ) -> bytes: 

385 if self._pending_unused_data is not None: 

386 data = self._pending_unused_data + bytes(data) 

387 self._pending_unused_data = None 

388 result = self._decompressor.decompress( 

389 self._decompressor.unconsumed_tail + data, max_length 

390 ) 

391 

392 # Concatenated gzip/deflate stream: decode the members after this one. 

393 if self._decompressor.eof and self._decompressor.unused_data: 

394 result = self._decompress_members(result, max_length) 

395 

396 # Only way to know that isal has no further data is checking we get no output 

397 self._last_empty = result == b"" 

398 

399 # Member ended exactly at chunk boundary — no unused_data, but the 

400 # next feed_data() call would fail on the spent decompressor. 

401 # Only reset for gzip; deflate's feed_eof() relies on eof=True to 

402 # confirm the stream is complete. 

403 if self._decompressor.eof and self._mode > self._zlib_backend.MAX_WBITS: 

404 self._decompressor = self._new_decompressor() 

405 

406 return result 

407 

408 def flush(self, length: int = 0) -> bytes: 

409 return ( 

410 self._decompressor.flush(length) 

411 if length > 0 

412 else self._decompressor.flush() 

413 ) 

414 

415 @property 

416 def data_available(self) -> bool: 

417 return ( 

418 bool(self._decompressor.unconsumed_tail) 

419 or not self._last_empty 

420 or self._pending_unused_data is not None 

421 ) 

422 

423 @property 

424 def eof(self) -> bool: 

425 return self._decompressor.eof 

426 

427 

428class BrotliDecompressor(DecompressionBaseHandler): 

429 # Supports both 'brotlipy' and 'Brotli' packages 

430 # since they share an import name. The top branches 

431 # are for 'brotlipy' and bottom branches for 'Brotli' 

432 def __init__( 

433 self, 

434 executor: Executor | None = None, 

435 max_sync_chunk_size: int | None = MAX_SYNC_CHUNK_SIZE, 

436 ) -> None: 

437 """Decompress data using the Brotli library.""" 

438 if not HAS_BROTLI: 

439 raise RuntimeError( 

440 "The brotli decompression is not available. " 

441 "Please install `Brotli` module" 

442 ) 

443 self._obj = brotli.Decompressor() 

444 self._last_empty = False 

445 super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) 

446 

447 def decompress_sync( 

448 self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED 

449 ) -> bytes: 

450 """Decompress the given data.""" 

451 if hasattr(self._obj, "decompress"): 

452 if max_length == ZLIB_MAX_LENGTH_UNLIMITED: 

453 result = cast(bytes, self._obj.decompress(data)) 

454 else: 

455 result = cast(bytes, self._obj.decompress(data, max_length)) 

456 else: 

457 if max_length == ZLIB_MAX_LENGTH_UNLIMITED: 

458 result = cast(bytes, self._obj.process(data)) 

459 else: 

460 result = cast(bytes, self._obj.process(data, max_length)) 

461 # Only way to know that brotli has no further data is checking we get no output 

462 self._last_empty = result == b"" 

463 return result 

464 

465 def flush(self) -> bytes: 

466 """Flush the decompressor.""" 

467 if hasattr(self._obj, "flush"): 

468 return cast(bytes, self._obj.flush()) 

469 return b"" 

470 

471 @property 

472 def data_available(self) -> bool: 

473 return not self._obj.is_finished() and not self._last_empty 

474 

475 

476class ZSTDDecompressor(ConcatDecompressionHandler["ZstdDecompressor"]): 

477 _unlimited = ZSTD_MAX_LENGTH_UNLIMITED 

478 

479 def __init__( 

480 self, 

481 executor: Executor | None = None, 

482 max_sync_chunk_size: int | None = MAX_SYNC_CHUNK_SIZE, 

483 ) -> None: 

484 if not HAS_ZSTD: 

485 raise RuntimeError( 

486 "The zstd decompression is not available. " 

487 "Please install `backports.zstd` module" 

488 ) 

489 super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size) 

490 self._decompressor = self._new_decompressor() 

491 

492 def _new_decompressor(self) -> "ZstdDecompressor": 

493 return ZstdDecompressor() 

494 

495 def decompress_sync( 

496 self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED 

497 ) -> bytes: 

498 # zstd uses -1 for unlimited, while zlib uses 0 for unlimited 

499 # Convert the zlib convention (0=unlimited) to zstd convention (-1=unlimited) 

500 zstd_max_length = ( 

501 ZSTD_MAX_LENGTH_UNLIMITED 

502 if max_length == ZLIB_MAX_LENGTH_UNLIMITED 

503 else max_length 

504 ) 

505 if self._pending_unused_data is not None: 

506 data = self._pending_unused_data + data 

507 self._pending_unused_data = None 

508 result = self._decompressor.decompress(data, zstd_max_length) 

509 

510 # Concatenated zstd stream: decode the frames after this one. 

511 if self._decompressor.eof and self._decompressor.unused_data: 

512 result = self._decompress_members(result, zstd_max_length) 

513 

514 # Frame ended exactly at chunk boundary — no unused_data, but the 

515 # next feed_data() call would fail on the spent decompressor. 

516 # Prepare a fresh one for the next chunk. 

517 if self._decompressor.eof: 

518 self._decompressor = self._new_decompressor() 

519 

520 return result 

521 

522 def flush(self) -> bytes: 

523 return b"" 

524 

525 @property 

526 def data_available(self) -> bool: 

527 return ( 

528 not self._decompressor.needs_input and not self._decompressor.eof 

529 ) or self._pending_unused_data is not None