Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/http_writer.py: 25%

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

2 

3import asyncio 

4import re 

5import sys 

6from collections.abc import Awaitable, Callable, Iterable 

7from typing import TYPE_CHECKING, NamedTuple 

8 

9from multidict import CIMultiDict 

10 

11from .abc import AbstractStreamWriter 

12from .base_protocol import BaseProtocol 

13from .client_exceptions import ClientConnectionResetError 

14from .compression_utils import ZLibCompressor 

15from .helpers import NO_EXTENSIONS 

16 

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

18 

19 

20MIN_PAYLOAD_FOR_WRITELINES = 2048 

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

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

23SKIP_WRITELINES = IS_PY313_BEFORE_313_2 or IS_PY_BEFORE_312_9 

24# writelines is not safe for use 

25# on Python 3.12+ until 3.12.9 

26# on Python 3.13+ until 3.13.2 

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

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

29 

30 

31class HttpVersion(NamedTuple): 

32 major: int 

33 minor: int 

34 

35 

36HttpVersion10 = HttpVersion(1, 0) 

37HttpVersion11 = HttpVersion(1, 1) 

38 

39 

40_T_OnChunkSent = ( 

41 Callable[ 

42 ["bytes | bytearray | memoryview[int] | memoryview[bytes]"], 

43 Awaitable[None], 

44 ] 

45 | None 

46) 

47_T_OnHeadersSent = Callable[["CIMultiDict[str]"], Awaitable[None]] | None 

48 

49 

50class StreamWriter(AbstractStreamWriter): 

51 

52 length: int | None = None 

53 chunked: bool = False 

54 _eof: bool = False 

55 _compress: ZLibCompressor | None = None 

56 

57 def __init__( 

58 self, 

59 protocol: BaseProtocol, 

60 loop: asyncio.AbstractEventLoop, 

61 on_chunk_sent: _T_OnChunkSent = None, 

62 on_headers_sent: _T_OnHeadersSent = None, 

63 ) -> None: 

64 self._protocol = protocol 

65 self.loop = loop 

66 self._on_chunk_sent: _T_OnChunkSent = on_chunk_sent 

67 self._on_headers_sent: _T_OnHeadersSent = on_headers_sent 

68 self._headers_buf: bytes | None = None 

69 self._headers_written: bool = False 

70 self.on_body_write = None 

71 

72 @property 

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

74 return self._protocol.transport 

75 

76 @property 

77 def protocol(self) -> BaseProtocol: 

78 return self._protocol 

79 

80 def enable_chunking(self) -> None: 

81 self.chunked = True 

82 

83 def enable_compression( 

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

85 ) -> None: 

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

87 

88 def _write( 

89 self, chunk: "bytes | bytearray | memoryview[int] | memoryview[bytes]" 

90 ) -> None: 

91 size = len(chunk) 

92 self.buffer_size += size 

93 self.output_size += size 

94 transport = self._protocol.transport 

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

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

97 transport.write(chunk) 

98 

99 def _writelines( 

100 self, 

101 chunks: Iterable["bytes | bytearray | memoryview[int] | memoryview[bytes]"], 

102 ) -> None: 

103 size = 0 

104 for chunk in chunks: 

105 size += len(chunk) 

106 self.buffer_size += size 

107 self.output_size += size 

108 transport = self._protocol.transport 

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

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

111 if SKIP_WRITELINES or size < MIN_PAYLOAD_FOR_WRITELINES: 

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

113 else: 

114 transport.writelines(chunks) 

115 

116 def _write_chunked_payload( 

117 self, chunk: "bytes | bytearray | memoryview[int] | memoryview[bytes]" 

118 ) -> None: 

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

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

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

122 

123 def _send_headers_with_payload( 

124 self, 

125 chunk: "bytes | bytearray | memoryview[int] | memoryview[bytes]", 

126 is_eof: bool, 

127 ) -> None: 

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

129 # Mark headers as written 

130 self._headers_written = True 

131 headers_buf = self._headers_buf 

132 self._headers_buf = None 

133 

134 if TYPE_CHECKING: 

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

136 # after checking that self._headers_buf is truthy 

137 assert headers_buf is not None 

138 

139 if not self.chunked: 

140 # Non-chunked: coalesce headers with body 

141 if chunk: 

142 self._writelines((headers_buf, chunk)) 

143 else: 

144 self._write(headers_buf) 

145 return 

146 

147 # Coalesce headers with chunked data 

148 if chunk: 

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

150 if is_eof: 

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

152 else: 

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

154 elif is_eof: 

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

156 else: 

157 self._write(headers_buf) 

158 

159 async def write( 

160 self, 

161 chunk: "bytes | bytearray | memoryview[int] | memoryview[bytes]", 

162 *, 

163 drain: bool = True, 

164 LIMIT: int = 0x10000, 

165 ) -> None: 

166 """ 

167 Writes chunk of data to a stream. 

168 

169 write_eof() indicates end of stream. 

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

171 write() return drain future. 

172 """ 

173 if self._on_chunk_sent is not None: 

174 await self._on_chunk_sent(chunk) 

175 

176 notify = self.on_body_write 

177 if isinstance(chunk, memoryview): 

178 body_size = chunk.nbytes 

179 if body_size != len(chunk): 

180 # just reshape it 

181 chunk = chunk.cast("c") 

182 else: 

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

184 

185 if self._compress is not None: 

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

187 if not chunk: 

188 if notify is not None and body_size: 

189 notify(body_size, self.output_size) 

190 return 

191 

192 if self.length is not None: 

193 chunk_len = len(chunk) 

194 if self.length >= chunk_len: 

195 self.length = self.length - chunk_len 

196 else: 

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

198 body_size = self.length 

199 chunk = chunk[: self.length] 

200 self.length = 0 

201 if not chunk: 

202 return 

203 

204 # Handle buffered headers for small payload optimization 

205 if self._headers_buf and not self._headers_written: 

206 self._send_headers_with_payload(chunk, False) 

207 elif chunk: 

208 if self.chunked: 

209 self._write_chunked_payload(chunk) 

210 else: 

211 self._write(chunk) 

212 else: 

213 return 

214 

215 if notify is not None and body_size: 

216 notify(body_size, self.output_size) 

217 if drain and self.buffer_size > LIMIT: 

218 self.buffer_size = 0 

219 await self.drain() 

220 

221 async def write_headers( 

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

223 ) -> None: 

224 """Write headers to the stream.""" 

225 if self._on_headers_sent is not None: 

226 await self._on_headers_sent(headers) 

227 # status + headers 

228 buf = _serialize_headers(status_line, headers) 

229 self._headers_written = False 

230 self._headers_buf = buf 

231 

232 def send_headers(self) -> None: 

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

234 if not self._headers_buf or self._headers_written: 

235 return 

236 

237 self._headers_written = True 

238 headers_buf = self._headers_buf 

239 self._headers_buf = None 

240 

241 if TYPE_CHECKING: 

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

243 assert headers_buf is not None 

244 

245 self._write(headers_buf) 

246 

247 def set_eof(self) -> None: 

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

249 if self._eof: 

250 return 

251 

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

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

254 if self._headers_buf and not self._headers_written: 

255 self._headers_written = True 

256 headers_buf = self._headers_buf 

257 self._headers_buf = None 

258 

259 if TYPE_CHECKING: 

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

261 assert headers_buf is not None 

262 

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

264 if self.chunked: 

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

266 else: 

267 self._write(headers_buf) 

268 elif self.chunked and self._headers_written: 

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

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

271 

272 self._eof = True 

273 

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

275 if self._eof: 

276 return 

277 

278 if chunk and self._on_chunk_sent is not None: 

279 await self._on_chunk_sent(chunk) 

280 

281 # Handle body/compression 

282 if self._compress: 

283 chunks: list[bytes] = [] 

284 chunks_len = 0 

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

286 chunks_len = len(compressed_chunk) 

287 chunks.append(compressed_chunk) 

288 

289 flush_chunk = self._compress.flush() 

290 chunks_len += len(flush_chunk) 

291 chunks.append(flush_chunk) 

292 assert chunks_len 

293 

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

295 if self._headers_buf and not self._headers_written: 

296 self._headers_written = True 

297 headers_buf = self._headers_buf 

298 self._headers_buf = None 

299 

300 if self.chunked: 

301 # Coalesce headers with compressed chunked data 

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

303 self._writelines( 

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

305 ) 

306 else: 

307 # Coalesce headers with compressed data 

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

309 # Headers already sent, just write compressed data 

310 elif self.chunked: 

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

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

313 elif len(chunks) > 1: 

314 self._writelines(chunks) 

315 else: 

316 self._write(chunks[0]) 

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

318 elif self._headers_buf and not self._headers_written: 

319 # Use helper to send headers with payload 

320 self._send_headers_with_payload(chunk, True) 

321 # Handle remaining body 

322 elif self.chunked: 

323 if chunk: 

324 # Write final chunk with EOF marker 

325 self._writelines( 

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

327 ) 

328 else: 

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

330 elif chunk: 

331 self._write(chunk) 

332 else: 

333 self._eof = True 

334 return 

335 

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

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

338 await self.drain() 

339 self._eof = True 

340 

341 async def drain(self) -> None: 

342 """Flush the write buffer. 

343 

344 The intended use is to write 

345 

346 await w.write(data) 

347 await w.drain() 

348 """ 

349 protocol = self._protocol 

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

351 await protocol._drain_helper() 

352 

353 

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

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

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

357 

358 

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

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

361 raise ValueError( 

362 "Forbidden control character detected in headers. " 

363 "Potential header injection attack." 

364 ) 

365 return string 

366 

367 

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

369 _safe_header(status_line) 

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

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

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

373 

374 

375_serialize_headers = _py_serialize_headers 

376 

377try: 

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

379 

380 _c_serialize_headers = _http_writer._serialize_headers 

381 if not NO_EXTENSIONS: 

382 _serialize_headers = _c_serialize_headers 

383except ImportError: 

384 pass