Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/aiohttp/http_writer.py: 26%

1"""Http related parsers and protocol.""" 

2 

3import asyncio 

4import re 

5import sys 

6from typing import ( # noqa 

7 TYPE_CHECKING, 

8 Any, 

9 Awaitable, 

10 Callable, 

11 Iterable, 

12 List, 

13 NamedTuple, 

14 Optional, 

15) 

16 

17from multidict import CIMultiDict 

18 

19from .abc import AbstractStreamWriter 

20from .base_protocol import BaseProtocol 

21from .client_exceptions import ClientConnectionResetError 

22from .compression_utils import ZLibCompressor 

23from .helpers import NO_EXTENSIONS 

24 

25__all__ = ("StreamWriter", "HttpVersion", "HttpVersion10", "HttpVersion11") 

26 

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

28 from collections.abc import Buffer 

29else: 

30 from typing import Union 

31 

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

33 

34 

35MIN_PAYLOAD_FOR_WRITELINES = 2048 

36IS_PY313_BEFORE_313_2 = (3, 13, 0) <= sys.version_info < (3, 13, 2) 

37IS_PY_BEFORE_312_9 = sys.version_info < (3, 12, 9) 

38SKIP_WRITELINES = IS_PY313_BEFORE_313_2 or IS_PY_BEFORE_312_9 

39# writelines is not safe for use 

40# on Python 3.12+ until 3.12.9 

41# on Python 3.13+ until 3.13.2 

42# and on older versions it not any faster than write 

43# CVE-2024-12254: https://github.com/python/cpython/pull/127656 

44 

45 

46class HttpVersion(NamedTuple): 

47 major: int 

48 minor: int 

49 

50 

51HttpVersion10 = HttpVersion(1, 0) 

52HttpVersion11 = HttpVersion(1, 1) 

53 

54 

55_T_OnChunkSent = Optional[Callable[[Buffer], Awaitable[None]]] 

56_T_OnHeadersSent = Optional[Callable[["CIMultiDict[str]"], Awaitable[None]]] 

57 

58 

59class StreamWriter(AbstractStreamWriter): 

60 

61 length: int | None = None 

62 chunked: bool = False 

63 _eof: bool = False 

64 _compress: ZLibCompressor | None = None 

65 

66 def __init__( 

67 self, 

68 protocol: BaseProtocol, 

69 loop: asyncio.AbstractEventLoop, 

70 on_chunk_sent: _T_OnChunkSent = None, 

71 on_headers_sent: _T_OnHeadersSent = None, 

72 ) -> None: 

73 self._protocol = protocol 

74 self.loop = loop 

75 self._on_chunk_sent: _T_OnChunkSent = on_chunk_sent 

76 self._on_headers_sent: _T_OnHeadersSent = on_headers_sent 

77 self._headers_buf: bytes | None = None 

78 self._headers_written: bool = False 

79 self.on_body_write = None 

80 

81 @property 

82 def transport(self) -> asyncio.Transport | None: 

83 return self._protocol.transport 

84 

85 @property 

86 def protocol(self) -> BaseProtocol: 

87 return self._protocol 

88 

89 def enable_chunking(self) -> None: 

90 self.chunked = True 

91 

92 def enable_compression( 

93 self, encoding: str = "deflate", strategy: int | None = None 

94 ) -> None: 

95 self._compress = ZLibCompressor(encoding=encoding, strategy=strategy) 

96 

97 def _write(self, chunk: Buffer) -> None: 

98 size = len(chunk) 

99 self.buffer_size += size 

100 self.output_size += size 

101 transport = self._protocol.transport 

102 if transport is None or transport.is_closing(): 

103 raise ClientConnectionResetError("Cannot write to closing transport") 

104 transport.write(chunk) 

105 

106 def _writelines(self, chunks: Iterable[Buffer]) -> None: 

107 size = 0 

108 for chunk in chunks: 

109 size += len(chunk) 

110 self.buffer_size += size 

111 self.output_size += size 

112 transport = self._protocol.transport 

113 if transport is None or transport.is_closing(): 

114 raise ClientConnectionResetError("Cannot write to closing transport") 

115 if SKIP_WRITELINES or size < MIN_PAYLOAD_FOR_WRITELINES: 

116 transport.write(b"".join(chunks)) 

117 else: 

118 transport.writelines(chunks) 

119 

120 def _write_chunked_payload(self, chunk: Buffer) -> None: 

121 """Write a chunk with proper chunked encoding.""" 

122 chunk_len_pre = f"{len(chunk):x}\r\n".encode("ascii") 

123 self._writelines((chunk_len_pre, chunk, b"\r\n")) 

124 

125 def _send_headers_with_payload(self, chunk: Buffer, is_eof: bool) -> None: 

126 """Send buffered headers with payload, coalescing into single write.""" 

127 # Mark headers as written 

128 self._headers_written = True 

129 headers_buf = self._headers_buf 

130 self._headers_buf = None 

131 

132 if TYPE_CHECKING: 

133 # Safe because callers (write() and write_eof()) only invoke this method 

134 # after checking that self._headers_buf is truthy 

135 assert headers_buf is not None 

136 

137 if not self.chunked: 

138 # Non-chunked: coalesce headers with body 

139 if chunk: 

140 self._writelines((headers_buf, chunk)) 

141 else: 

142 self._write(headers_buf) 

143 return 

144 

145 # Coalesce headers with chunked data 

146 if chunk: 

147 chunk_len_pre = f"{len(chunk):x}\r\n".encode("ascii") 

148 if is_eof: 

149 self._writelines((headers_buf, chunk_len_pre, chunk, b"\r\n0\r\n\r\n")) 

150 else: 

151 self._writelines((headers_buf, chunk_len_pre, chunk, b"\r\n")) 

152 elif is_eof: 

153 self._writelines((headers_buf, b"0\r\n\r\n")) 

154 else: 

155 self._write(headers_buf) 

156 

157 async def write( 

158 self, chunk: Buffer, *, drain: bool = True, LIMIT: int = 0x10000 

159 ) -> None: 

160 """ 

161 Writes chunk of data to a stream. 

162 

163 write_eof() indicates end of stream. 

164 writer can't be used after write_eof() method being called. 

165 write() return drain future. 

166 """ 

167 if self._on_chunk_sent is not None: 

168 await self._on_chunk_sent(chunk) 

169 

170 notify = self.on_body_write 

171 if isinstance(chunk, memoryview): 

172 body_size = chunk.nbytes 

173 if body_size != len(chunk): 

174 # just reshape it 

175 chunk = chunk.cast("c") 

176 else: 

177 body_size = len(chunk) if notify is not None else 0 

178 

179 if self._compress is not None: 

180 chunk = await self._compress.compress(chunk) 

181 if not chunk: 

182 if notify is not None and body_size: 

183 notify(body_size, self.output_size) 

184 return 

185 

186 if self.length is not None: 

187 chunk_len = len(chunk) 

188 if self.length >= chunk_len: 

189 self.length = self.length - chunk_len 

190 else: 

191 # Bytes clipped at the declared length are discarded, not sent. 

192 body_size = self.length 

193 chunk = chunk[: self.length] 

194 self.length = 0 

195 if not chunk: 

196 return 

197 

198 # Handle buffered headers for small payload optimization 

199 if self._headers_buf and not self._headers_written: 

200 self._send_headers_with_payload(chunk, False) 

201 elif chunk: 

202 if self.chunked: 

203 self._write_chunked_payload(chunk) 

204 else: 

205 self._write(chunk) 

206 else: 

207 return 

208 

209 if notify is not None and body_size: 

210 notify(body_size, self.output_size) 

211 if drain and self.buffer_size > LIMIT: 

212 self.buffer_size = 0 

213 await self.drain() 

214 

215 async def write_headers( 

216 self, status_line: str, headers: "CIMultiDict[str]" 

217 ) -> None: 

218 """Write headers to the stream.""" 

219 if self._on_headers_sent is not None: 

220 await self._on_headers_sent(headers) 

221 # status + headers 

222 buf = _serialize_headers(status_line, headers) 

223 self._headers_written = False 

224 self._headers_buf = buf 

225 

226 def send_headers(self) -> None: 

227 """Force sending buffered headers if not already sent.""" 

228 if not self._headers_buf or self._headers_written: 

229 return 

230 

231 self._headers_written = True 

232 headers_buf = self._headers_buf 

233 self._headers_buf = None 

234 

235 if TYPE_CHECKING: 

236 # Safe because we only enter this block when self._headers_buf is truthy 

237 assert headers_buf is not None 

238 

239 self._write(headers_buf) 

240 

241 def set_eof(self) -> None: 

242 """Indicate that the message is complete.""" 

243 if self._eof: 

244 return 

245 

246 # If headers haven't been sent yet, send them now 

247 # This handles the case where there's no body at all 

248 if self._headers_buf and not self._headers_written: 

249 self._headers_written = True 

250 headers_buf = self._headers_buf 

251 self._headers_buf = None 

252 

253 if TYPE_CHECKING: 

254 # Safe because we only enter this block when self._headers_buf is truthy 

255 assert headers_buf is not None 

256 

257 # Combine headers and chunked EOF marker in a single write 

258 if self.chunked: 

259 self._writelines((headers_buf, b"0\r\n\r\n")) 

260 else: 

261 self._write(headers_buf) 

262 elif self.chunked and self._headers_written: 

263 # Headers already sent, just send the final chunk marker 

264 self._write(b"0\r\n\r\n") 

265 

266 self._eof = True 

267 

268 async def write_eof(self, chunk: bytes = b"") -> None: 

269 if self._eof: 

270 return 

271 

272 if chunk and self._on_chunk_sent is not None: 

273 await self._on_chunk_sent(chunk) 

274 

275 # Handle body/compression 

276 if self._compress: 

277 chunks: list[bytes] = [] 

278 chunks_len = 0 

279 if chunk and (compressed_chunk := await self._compress.compress(chunk)): 

280 chunks_len = len(compressed_chunk) 

281 chunks.append(compressed_chunk) 

282 

283 flush_chunk = self._compress.flush() 

284 chunks_len += len(flush_chunk) 

285 chunks.append(flush_chunk) 

286 assert chunks_len 

287 

288 # Send buffered headers with compressed data if not yet sent 

289 if self._headers_buf and not self._headers_written: 

290 self._headers_written = True 

291 headers_buf = self._headers_buf 

292 self._headers_buf = None 

293 

294 if self.chunked: 

295 # Coalesce headers with compressed chunked data 

296 chunk_len_pre = f"{chunks_len:x}\r\n".encode("ascii") 

297 self._writelines( 

298 (headers_buf, chunk_len_pre, *chunks, b"\r\n0\r\n\r\n") 

299 ) 

300 else: 

301 # Coalesce headers with compressed data 

302 self._writelines((headers_buf, *chunks)) 

303 # Headers already sent, just write compressed data 

304 elif self.chunked: 

305 chunk_len_pre = f"{chunks_len:x}\r\n".encode("ascii") 

306 self._writelines((chunk_len_pre, *chunks, b"\r\n0\r\n\r\n")) 

307 elif len(chunks) > 1: 

308 self._writelines(chunks) 

309 else: 

310 self._write(chunks[0]) 

311 # No compression - send buffered headers if not yet sent 

312 elif self._headers_buf and not self._headers_written: 

313 # Use helper to send headers with payload 

314 self._send_headers_with_payload(chunk, True) 

315 # Handle remaining body 

316 elif self.chunked: 

317 if chunk: 

318 # Write final chunk with EOF marker 

319 self._writelines( 

320 (f"{len(chunk):x}\r\n".encode("ascii"), chunk, b"\r\n0\r\n\r\n") 

321 ) 

322 else: 

323 self._write(b"0\r\n\r\n") 

324 elif chunk: 

325 self._write(chunk) 

326 else: 

327 self._eof = True 

328 return 

329 

330 if chunk and (notify := self.on_body_write) is not None: 

331 notify(len(chunk), self.output_size) 

332 await self.drain() 

333 self._eof = True 

334 

335 async def drain(self) -> None: 

336 """Flush the write buffer. 

337 

338 The intended use is to write 

339 

340 await w.write(data) 

341 await w.drain() 

342 """ 

343 protocol = self._protocol 

344 if protocol.transport is not None and protocol._paused: 

345 await protocol._drain_helper() 

346 

347 

348# https://www.rfc-editor.org/info/rfc9110/#section-5.5-5 

349# https://www.rfc-editor.org/info/rfc9112/#section-4-3 

350_FORBIDDEN_HEADER_CHARS_RE = re.compile(r"[\x00-\x08\x0a-\x1f\x7f]") 

351 

352 

353def _safe_header(string: str) -> str: 

354 if _FORBIDDEN_HEADER_CHARS_RE.search(string) is not None: 

355 raise ValueError( 

356 "Forbidden control character detected in headers. " 

357 "Potential header injection attack." 

358 ) 

359 return string 

360 

361 

362def _py_serialize_headers(status_line: str, headers: "CIMultiDict[str]") -> bytes: 

363 _safe_header(status_line) 

364 headers_gen = (_safe_header(k) + ": " + _safe_header(v) for k, v in headers.items()) 

365 line = status_line + "\r\n" + "\r\n".join(headers_gen) + "\r\n\r\n" 

366 return line.encode("utf-8") 

367 

368 

369_serialize_headers = _py_serialize_headers 

370 

371try: 

372 import aiohttp._http_writer as _http_writer # type: ignore[import-not-found] 

373 

374 _c_serialize_headers = _http_writer._serialize_headers 

375 if not NO_EXTENSIONS: 

376 _serialize_headers = _c_serialize_headers 

377except ImportError: 

378 pass