Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/client_proto.py: 20%

1import asyncio 

2from collections.abc import Callable 

3from contextlib import suppress 

4from typing import Protocol 

5 

6from ._websocket.reader import WebSocketDataQueue 

7from .base_protocol import BaseProtocol 

8from .client_exceptions import ( 

9 ClientConnectionError, 

10 ClientOSError, 

11 ClientPayloadError, 

12 ServerDisconnectedError, 

13 SocketTimeoutError, 

14) 

15from .helpers import ( 

16 _EXC_SENTINEL, 

17 DEFAULT_CHUNK_SIZE, 

18 EMPTY_BODY_STATUS_CODES, 

19 BaseTimerContext, 

20 ErrorableProtocol, 

21 set_exception, 

22 set_result, 

23) 

24from .http import HttpResponseParser, RawResponseMessage, WebSocketReader 

25from .http_exceptions import HttpProcessingError 

26from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader 

27 

28 

29class _Payload(ErrorableProtocol, Protocol): 

30 def is_eof(self) -> bool: ... 

31 

32 

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

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

35 

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

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

38 DataQueue.__init__(self, loop) 

39 

40 self._should_close = False 

41 

42 self._payload: _Payload | None = None 

43 self._skip_payload = False 

44 self._payload_parser: WebSocketReader | None = None 

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

46 

47 self._timer = None 

48 self._tail = b"" 

49 self._payload_parser_failed = False 

50 self._read_bufsize = DEFAULT_CHUNK_SIZE 

51 

52 self._read_timeout: float | None = None 

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

54 

55 self._timeout_ceil_threshold: float | None = 5 

56 

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

58 self._connection_lost_called = False 

59 

60 @property 

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

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

63 

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

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

66 futures that will never be awaited. 

67 

68 Returns: 

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

70 this property was accessed 

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

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

73 """ 

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

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

76 return self._closed 

77 

78 @property 

79 def upgraded(self) -> bool: 

80 return self._upgraded 

81 

82 @property 

83 def should_close(self) -> bool: 

84 return bool( 

85 self._should_close 

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

87 or self._upgraded 

88 or self._exception is not None 

89 or self._payload_parser is not None 

90 or self._buffer 

91 or self._tail 

92 ) 

93 

94 def force_close(self) -> None: 

95 self._should_close = True 

96 

97 def close(self) -> None: 

98 self._exception = None # Break cyclic references 

99 transport = self.transport 

100 if transport is not None: 

101 transport.close() 

102 self.transport = None 

103 self._payload = None 

104 self._drop_timeout() 

105 

106 def abort(self) -> None: 

107 self._exception = None # Break cyclic references 

108 transport = self.transport 

109 if transport is not None: 

110 transport.abort() 

111 self.transport = None 

112 self._payload = None 

113 self._drop_timeout() 

114 

115 def is_connected(self) -> bool: 

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

117 

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

119 self._connection_lost_called = True 

120 self._drop_timeout() 

121 

122 original_connection_error = exc 

123 reraised_exc = original_connection_error 

124 

125 connection_closed_cleanly = original_connection_error is None 

126 

127 if self._closed is not None: 

128 # If someone is waiting for the closed future, 

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

130 # self._closed is None, it means that 

131 # connection_lost() was called already 

132 # or nobody is waiting for it. 

133 if connection_closed_cleanly: 

134 set_result(self._closed, None) 

135 else: 

136 assert original_connection_error is not None 

137 set_exception( 

138 self._closed, 

139 ClientConnectionError( 

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

141 ), 

142 original_connection_error, 

143 ) 

144 

145 if self._payload_parser is not None: 

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

147 self._payload_parser.feed_eof() 

148 

149 uncompleted = None 

150 if self._parser is not None: 

151 try: 

152 uncompleted = self._parser.feed_eof() 

153 except Exception as underlying_exc: 

154 if self._payload is not None: 

155 client_payload_exc_msg = ( 

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

157 ) 

158 if not connection_closed_cleanly: 

159 client_payload_exc_msg = ( 

160 f"{client_payload_exc_msg !s}. " 

161 f"{original_connection_error !r}" 

162 ) 

163 set_exception( 

164 self._payload, 

165 ClientPayloadError(client_payload_exc_msg), 

166 underlying_exc, 

167 ) 

168 

169 if not self.is_eof(): 

170 if isinstance(original_connection_error, OSError): 

171 reraised_exc = ClientOSError(*original_connection_error.args) 

172 if connection_closed_cleanly: 

173 reraised_exc = ServerDisconnectedError(uncompleted) 

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

175 # we do it anyway below 

176 underlying_non_eof_exc = ( 

177 _EXC_SENTINEL 

178 if connection_closed_cleanly 

179 else original_connection_error 

180 ) 

181 assert underlying_non_eof_exc is not None 

182 assert reraised_exc is not None 

183 self.set_exception(reraised_exc, underlying_non_eof_exc) 

184 

185 self._should_close = True 

186 self._parser = None 

187 self._payload = None 

188 self._payload_parser = None 

189 self._payload_parser_failed = False 

190 self._reading_paused = False 

191 self._buffer_paused = False 

192 

193 super().connection_lost(reraised_exc) 

194 

195 def eof_received(self) -> None: 

196 # should call parser.feed_eof() most likely 

197 self._drop_timeout() 

198 

199 def pause_reading(self) -> None: 

200 super().pause_reading() 

201 self._drop_timeout() 

202 

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

204 was_paused = self._reading_paused 

205 super().resume_reading(resume_parser) 

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

207 # bound is still holding the transport. 

208 if was_paused and not self._buffer_paused: 

209 self._reschedule_timeout() 

210 

211 def _drain_tail(self) -> None: 

212 if self._tail: 

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

214 self.data_received(data) 

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

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

217 if self._buffer_paused and not still_bound: 

218 self._resume_reading_for_buffer() 

219 if self._reading_paused or self._buffer_paused: 

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

221 # still holding the transport means it stays stopped. 

222 self._drop_timeout() 

223 

224 def set_exception( 

225 self, 

226 exc: type[BaseException] | BaseException, 

227 exc_cause: BaseException = _EXC_SENTINEL, 

228 ) -> None: 

229 self._should_close = True 

230 self._drop_timeout() 

231 super().set_exception(exc, exc_cause) 

232 

233 def set_parser( 

234 self, 

235 parser: WebSocketReader, 

236 payload: WebSocketDataQueue, 

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

238 ) -> None: 

239 self._payload = payload 

240 self._payload_parser = parser 

241 self._data_received_cb = data_received_cb 

242 

243 self._drop_timeout() 

244 

245 self._drain_tail() 

246 

247 def set_response_params( 

248 self, 

249 *, 

250 timer: BaseTimerContext | None = None, 

251 skip_payload: bool = False, 

252 read_until_eof: bool = False, 

253 auto_decompress: bool = True, 

254 read_timeout: float | None = None, 

255 read_bufsize: int = DEFAULT_CHUNK_SIZE, 

256 timeout_ceil_threshold: float = 5, 

257 max_line_size: int = 8190, 

258 max_field_size: int = 8190, 

259 max_headers: int = 128, 

260 ) -> None: 

261 self._skip_payload = skip_payload 

262 

263 self._read_bufsize = read_bufsize 

264 self._read_timeout = read_timeout 

265 

266 self._timeout_ceil_threshold = timeout_ceil_threshold 

267 

268 self._parser = HttpResponseParser( 

269 self, 

270 self._loop, 

271 read_bufsize, 

272 timer=timer, 

273 payload_exception=ClientPayloadError, 

274 response_with_body=not skip_payload, 

275 read_until_eof=read_until_eof, 

276 auto_decompress=auto_decompress, 

277 max_line_size=max_line_size, 

278 max_field_size=max_field_size, 

279 max_headers=max_headers, 

280 ) 

281 

282 self._drain_tail() 

283 

284 def _drop_timeout(self) -> None: 

285 if self._read_timeout_handle is not None: 

286 self._read_timeout_handle.cancel() 

287 self._read_timeout_handle = None 

288 

289 def _reschedule_timeout(self) -> None: 

290 timeout = self._read_timeout 

291 if self._read_timeout_handle is not None: 

292 self._read_timeout_handle.cancel() 

293 

294 if timeout: 

295 self._read_timeout_handle = self._loop.call_later( 

296 timeout, self._on_read_timeout 

297 ) 

298 else: 

299 self._read_timeout_handle = None 

300 

301 def start_timeout(self) -> None: 

302 self._reschedule_timeout() 

303 

304 @property 

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

306 return self._read_timeout 

307 

308 @read_timeout.setter 

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

310 self._read_timeout = read_timeout 

311 

312 def _on_read_timeout(self) -> None: 

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

314 self.set_exception(exc) 

315 if self._payload is not None: 

316 set_exception(self._payload, exc) 

317 

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

319 if self._payload_parser_failed: 

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

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

322 # so they must not hold off sock_read either. 

323 return 

324 

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

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

327 if data: 

328 self._reschedule_timeout() 

329 

330 # custom payload parser - currently always WebSocketReader 

331 if self._payload_parser is not None: 

332 if self._data_received_cb is not None: 

333 self._data_received_cb() 

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

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

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

337 if protocol_error: 

338 self._payload = None 

339 self._payload_parser = None 

340 self._payload_parser_failed = True 

341 return 

342 

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

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

345 self._tail += data 

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

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

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

349 if ( 

350 not self._buffer_paused 

351 and self._tail 

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

353 ): 

354 self._pause_reading_for_buffer() 

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

356 self._drop_timeout() 

357 return 

358 

359 # parse http messages 

360 try: 

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

362 except BaseException as underlying_exc: 

363 if self.transport is not None: 

364 # connection.release() could be called BEFORE 

365 # data_received(), the transport is already 

366 # closed in this case 

367 self.transport.close() 

368 if not isinstance(underlying_exc, Exception): 

369 raise 

370 # should_close is True after the call 

371 if isinstance(underlying_exc, HttpProcessingError): 

372 exc = HttpProcessingError( 

373 code=underlying_exc.code, 

374 message=underlying_exc.message, 

375 headers=underlying_exc.headers, 

376 ) 

377 else: 

378 exc = HttpProcessingError() 

379 self.set_exception(exc, underlying_exc) 

380 return 

381 

382 self._upgraded = upgraded 

383 

384 payload: StreamReader | None = None 

385 for message, payload in messages: 

386 if message.should_close: 

387 self._should_close = True 

388 

389 self._payload = payload 

390 

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

392 self.feed_data((message, EMPTY_PAYLOAD)) 

393 else: 

394 self.feed_data((message, payload)) 

395 

396 if payload is not None: 

397 # new message(s) was processed 

398 # register timeout handler unsubscribing 

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

400 # EMPTY_PAYLOAD 

401 if payload is not EMPTY_PAYLOAD: 

402 payload.on_eof(self._drop_timeout) 

403 else: 

404 self._drop_timeout() 

405 

406 if upgraded and tail: 

407 self.data_received(tail)