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

1import asyncio 

2from contextlib import suppress 

3from typing import Any, Callable 

4 

5from .base_protocol import BaseProtocol 

6from .client_exceptions import ( 

7 ClientConnectionError, 

8 ClientOSError, 

9 ClientPayloadError, 

10 ServerDisconnectedError, 

11 SocketTimeoutError, 

12) 

13from .helpers import ( 

14 _EXC_SENTINEL, 

15 DEFAULT_CHUNK_SIZE, 

16 EMPTY_BODY_STATUS_CODES, 

17 BaseTimerContext, 

18 set_exception, 

19 set_result, 

20) 

21from .http import HttpResponseParser, RawResponseMessage 

22from .http_exceptions import HttpProcessingError 

23from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader 

24 

25 

26class ResponseHandler(BaseProtocol, DataQueue[tuple[RawResponseMessage, StreamReader]]): 

27 """Helper class to adapt between Protocol and StreamReader.""" 

28 

29 def __init__(self, loop: asyncio.AbstractEventLoop) -> None: 

30 BaseProtocol.__init__(self, loop=loop, parser=None) 

31 DataQueue.__init__(self, loop) 

32 

33 self._should_close = False 

34 

35 self._payload: StreamReader | None = None 

36 self._skip_payload = False 

37 self._payload_parser = None 

38 self._data_received_cb: Callable[[], None] | None = None 

39 

40 self._timer = None 

41 self._tail = b"" 

42 self._payload_parser_failed = False 

43 self._read_bufsize = DEFAULT_CHUNK_SIZE 

44 

45 self._read_timeout: float | None = None 

46 self._read_timeout_handle: asyncio.TimerHandle | None = None 

47 

48 self._timeout_ceil_threshold: float | None = 5 

49 

50 self._closed: None | asyncio.Future[None] = None 

51 self._connection_lost_called = False 

52 

53 @property 

54 def closed(self) -> None | asyncio.Future[None]: 

55 """Future that is set when the connection is closed. 

56 

57 This property returns a Future that will be completed when the connection 

58 is closed. The Future is created lazily on first access to avoid creating 

59 futures that will never be awaited. 

60 

61 Returns: 

62 - A Future[None] if the connection is still open or was closed after 

63 this property was accessed 

64 - None if connection_lost() was already called before this property 

65 was ever accessed (indicating no one is waiting for the closure) 

66 """ 

67 if self._closed is None and not self._connection_lost_called: 

68 self._closed = self._loop.create_future() 

69 return self._closed 

70 

71 @property 

72 def upgraded(self) -> bool: 

73 return self._upgraded 

74 

75 @property 

76 def should_close(self) -> bool: 

77 return bool( 

78 self._should_close 

79 or (self._payload is not None and not self._payload.is_eof()) 

80 or self._upgraded 

81 or self._exception is not None 

82 or self._payload_parser is not None 

83 or self._buffer 

84 or self._tail 

85 ) 

86 

87 def force_close(self) -> None: 

88 self._should_close = True 

89 

90 def close(self) -> None: 

91 self._exception = None # Break cyclic references 

92 transport = self.transport 

93 if transport is not None: 

94 transport.close() 

95 self.transport = None 

96 self._payload = None 

97 self._drop_timeout() 

98 

99 def abort(self) -> None: 

100 self._exception = None # Break cyclic references 

101 transport = self.transport 

102 if transport is not None: 

103 transport.abort() 

104 self.transport = None 

105 self._payload = None 

106 self._drop_timeout() 

107 

108 def is_connected(self) -> bool: 

109 return self.transport is not None and not self.transport.is_closing() 

110 

111 def connection_lost(self, exc: BaseException | None) -> None: 

112 self._connection_lost_called = True 

113 self._drop_timeout() 

114 

115 original_connection_error = exc 

116 reraised_exc = original_connection_error 

117 

118 connection_closed_cleanly = original_connection_error is None 

119 

120 if self._closed is not None: 

121 # If someone is waiting for the closed future, 

122 # we should set it to None or an exception. If 

123 # self._closed is None, it means that 

124 # connection_lost() was called already 

125 # or nobody is waiting for it. 

126 if connection_closed_cleanly: 

127 set_result(self._closed, None) 

128 else: 

129 assert original_connection_error is not None 

130 set_exception( 

131 self._closed, 

132 ClientConnectionError( 

133 f"Connection lost: {original_connection_error !s}", 

134 ), 

135 original_connection_error, 

136 ) 

137 

138 if self._payload_parser is not None: 

139 with suppress(Exception): # FIXME: log this somehow? 

140 self._payload_parser.feed_eof() 

141 

142 uncompleted = None 

143 if self._parser is not None: 

144 try: 

145 uncompleted = self._parser.feed_eof() 

146 except Exception as underlying_exc: 

147 if self._payload is not None: 

148 client_payload_exc_msg = ( 

149 f"Response payload is not completed: {underlying_exc !r}" 

150 ) 

151 if not connection_closed_cleanly: 

152 client_payload_exc_msg = ( 

153 f"{client_payload_exc_msg !s}. " 

154 f"{original_connection_error !r}" 

155 ) 

156 set_exception( 

157 self._payload, 

158 ClientPayloadError(client_payload_exc_msg), 

159 underlying_exc, 

160 ) 

161 

162 if not self.is_eof(): 

163 if isinstance(original_connection_error, OSError): 

164 reraised_exc = ClientOSError(*original_connection_error.args) 

165 if connection_closed_cleanly: 

166 reraised_exc = ServerDisconnectedError(uncompleted) 

167 # assigns self._should_close to True as side effect, 

168 # we do it anyway below 

169 underlying_non_eof_exc = ( 

170 _EXC_SENTINEL 

171 if connection_closed_cleanly 

172 else original_connection_error 

173 ) 

174 assert underlying_non_eof_exc is not None 

175 assert reraised_exc is not None 

176 self.set_exception(reraised_exc, underlying_non_eof_exc) 

177 

178 self._should_close = True 

179 self._parser = None 

180 self._payload = None 

181 self._payload_parser = None 

182 self._payload_parser_failed = False 

183 self._reading_paused = False 

184 self._buffer_paused = False 

185 

186 super().connection_lost(reraised_exc) 

187 

188 def eof_received(self) -> None: 

189 # should call parser.feed_eof() most likely 

190 self._drop_timeout() 

191 

192 def pause_reading(self) -> None: 

193 super().pause_reading() 

194 self._drop_timeout() 

195 

196 def resume_reading(self, resume_parser: bool = True) -> None: 

197 was_paused = self._reading_paused 

198 super().resume_reading(resume_parser) 

199 # sock_read measures the peer, so it stays stopped while the tail 

200 # bound is still holding the transport. 

201 if was_paused and not self._buffer_paused: 

202 self._reschedule_timeout() 

203 

204 def _drain_tail(self) -> None: 

205 if self._tail: 

206 data, self._tail = self._tail, b"" 

207 self.data_received(data) 

208 # Tested empty-first so a read_bufsize of 0 cannot wedge the connection. 

209 still_bound = bool(self._tail) and len(self._tail) >= self._read_bufsize 

210 if self._buffer_paused and not still_bound: 

211 self._resume_reading_for_buffer() 

212 if self._reading_paused or self._buffer_paused: 

213 # The drain restarted sock_read through data_received(); anything 

214 # still holding the transport means it stays stopped. 

215 self._drop_timeout() 

216 

217 def set_exception( 

218 self, 

219 exc: BaseException, 

220 exc_cause: BaseException = _EXC_SENTINEL, 

221 ) -> None: 

222 self._should_close = True 

223 self._drop_timeout() 

224 super().set_exception(exc, exc_cause) 

225 

226 def set_parser( 

227 self, 

228 parser: Any, 

229 payload: Any, 

230 data_received_cb: Callable[[], None] | None = None, 

231 ) -> None: 

232 # TODO: actual types are: 

233 # parser: WebSocketReader 

234 # payload: WebSocketDataQueue 

235 # but they are not generi enough 

236 # Need an ABC for both types 

237 self._payload = payload 

238 self._payload_parser = parser 

239 self._data_received_cb = data_received_cb 

240 

241 self._drop_timeout() 

242 

243 self._drain_tail() 

244 

245 def set_response_params( 

246 self, 

247 *, 

248 timer: BaseTimerContext | None = None, 

249 skip_payload: bool = False, 

250 read_until_eof: bool = False, 

251 auto_decompress: bool = True, 

252 read_timeout: float | None = None, 

253 read_bufsize: int = DEFAULT_CHUNK_SIZE, 

254 timeout_ceil_threshold: float = 5, 

255 max_line_size: int = 8190, 

256 max_field_size: int = 8190, 

257 max_headers: int = 128, 

258 ) -> None: 

259 self._skip_payload = skip_payload 

260 

261 self._read_bufsize = read_bufsize 

262 self._read_timeout = read_timeout 

263 

264 self._timeout_ceil_threshold = timeout_ceil_threshold 

265 

266 self._parser = HttpResponseParser( 

267 self, 

268 self._loop, 

269 read_bufsize, 

270 timer=timer, 

271 payload_exception=ClientPayloadError, 

272 response_with_body=not skip_payload, 

273 read_until_eof=read_until_eof, 

274 auto_decompress=auto_decompress, 

275 max_line_size=max_line_size, 

276 max_field_size=max_field_size, 

277 max_headers=max_headers, 

278 ) 

279 

280 self._drain_tail() 

281 

282 def _drop_timeout(self) -> None: 

283 if self._read_timeout_handle is not None: 

284 self._read_timeout_handle.cancel() 

285 self._read_timeout_handle = None 

286 

287 def _reschedule_timeout(self) -> None: 

288 timeout = self._read_timeout 

289 if self._read_timeout_handle is not None: 

290 self._read_timeout_handle.cancel() 

291 

292 if timeout: 

293 self._read_timeout_handle = self._loop.call_later( 

294 timeout, self._on_read_timeout 

295 ) 

296 else: 

297 self._read_timeout_handle = None 

298 

299 def start_timeout(self) -> None: 

300 self._reschedule_timeout() 

301 

302 @property 

303 def read_timeout(self) -> float | None: 

304 return self._read_timeout 

305 

306 @read_timeout.setter 

307 def read_timeout(self, read_timeout: float | None) -> None: 

308 self._read_timeout = read_timeout 

309 

310 def _on_read_timeout(self) -> None: 

311 exc = SocketTimeoutError("Timeout on reading data from socket") 

312 self.set_exception(exc) 

313 if self._payload is not None: 

314 set_exception(self._payload, exc) 

315 

316 def data_received(self, data: bytes) -> None: 

317 if self._payload_parser_failed: 

318 # Dropped, not closed on, so queued messages stay answerable and 

319 # the peer's FIN still arrives. Discarded bytes are not progress, 

320 # so they must not hold off sock_read either. 

321 return 

322 

323 # If no data, then we are resuming decompression. We haven't received 

324 # data from the socket, so we can avoid the reschedule overhead. 

325 if data: 

326 self._reschedule_timeout() 

327 

328 # custom payload parser - currently always WebSocketReader 

329 if self._payload_parser is not None: 

330 if self._data_received_cb is not None: 

331 self._data_received_cb() 

332 # WebSocketReader signals EOF only for a protocol error, never 

333 # for a clean close, and has already put the error on the queue. 

334 protocol_error, _ = self._payload_parser.feed_data(data) 

335 if protocol_error: 

336 self._payload = None 

337 self._payload_parser = None 

338 self._payload_parser_failed = True 

339 return 

340 

341 if self._upgraded or self._parser is None: 

342 # i.e. websocket connection, websocket parser is not set yet 

343 self._tail += data 

344 # Nothing drains this until set_parser() runs, and an await 

345 # after the 101 can hold that off; stop reading instead. 

346 # Tested non-empty so a read_bufsize of 0 cannot pause on nothing. 

347 if ( 

348 not self._buffer_paused 

349 and self._tail 

350 and len(self._tail) >= self._read_bufsize 

351 ): 

352 self._pause_reading_for_buffer() 

353 # sock_read measures the peer; this pause is ours. 

354 self._drop_timeout() 

355 return 

356 

357 # parse http messages 

358 try: 

359 messages, upgraded, tail = self._parser.feed_data(data) 

360 except BaseException as underlying_exc: 

361 if self.transport is not None: 

362 # connection.release() could be called BEFORE 

363 # data_received(), the transport is already 

364 # closed in this case 

365 self.transport.close() 

366 if not isinstance(underlying_exc, Exception): 

367 raise 

368 # should_close is True after the call 

369 if isinstance(underlying_exc, HttpProcessingError): 

370 exc = HttpProcessingError( 

371 code=underlying_exc.code, 

372 message=underlying_exc.message, 

373 headers=underlying_exc.headers, 

374 ) 

375 else: 

376 exc = HttpProcessingError() 

377 self.set_exception(exc, underlying_exc) 

378 return 

379 

380 self._upgraded = upgraded 

381 

382 payload: StreamReader | None = None 

383 for message, payload in messages: 

384 if message.should_close: 

385 self._should_close = True 

386 

387 self._payload = payload 

388 

389 if self._skip_payload or message.code in EMPTY_BODY_STATUS_CODES: 

390 self.feed_data((message, EMPTY_PAYLOAD), 0) 

391 else: 

392 self.feed_data((message, payload), 0) 

393 

394 if payload is not None: 

395 # new message(s) was processed 

396 # register timeout handler unsubscribing 

397 # either on end-of-stream or immediately for 

398 # EMPTY_PAYLOAD 

399 if payload is not EMPTY_PAYLOAD: 

400 payload.on_eof(self._drop_timeout) 

401 else: 

402 self._drop_timeout() 

403 

404 if upgraded and tail: 

405 self.data_received(tail)