Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/websockets/frames.py: 47%

1from __future__ import annotations 

2 

3import dataclasses 

4import enum 

5import io 

6import os 

7import re 

8import secrets 

9import struct 

10from collections.abc import Generator, Sequence 

11from typing import Callable, Self 

12 

13from .exceptions import PayloadTooBig, ProtocolError 

14from .typing import BytesLike 

15 

16 

17try: 

18 from .speedups import apply_mask 

19except ImportError: 

20 from .utils import apply_mask 

21 

22 

23__all__ = [ 

24 "Opcode", 

25 "CloseCode", 

26 "Frame", 

27 "Close", 

28] 

29 

30is_space_or_printable_ascii = re.compile(rb"[\x09-\x0D\x20-\x7E]*").fullmatch 

31 

32 

33class Opcode(enum.IntEnum): 

34 """Opcode values for WebSocket frames.""" 

35 

36 CONT, TEXT, BINARY = 0x00, 0x01, 0x02 

37 CLOSE, PING, PONG = 0x08, 0x09, 0x0A 

38 

39 

40CONT = Opcode.CONT 

41TEXT = Opcode.TEXT 

42BINARY = Opcode.BINARY 

43CLOSE = Opcode.CLOSE 

44PING = Opcode.PING 

45PONG = Opcode.PONG 

46 

47DATA_OPCODES = CONT, TEXT, BINARY 

48CTRL_OPCODES = CLOSE, PING, PONG 

49 

50 

51class CloseCode(enum.IntEnum): 

52 """Close code values for WebSocket close frames.""" 

53 

54 NORMAL_CLOSURE = 1000 

55 GOING_AWAY = 1001 

56 PROTOCOL_ERROR = 1002 

57 UNSUPPORTED_DATA = 1003 

58 # 1004 is reserved 

59 NO_STATUS_RCVD = 1005 

60 ABNORMAL_CLOSURE = 1006 

61 INVALID_DATA = 1007 

62 POLICY_VIOLATION = 1008 

63 MESSAGE_TOO_BIG = 1009 

64 MANDATORY_EXTENSION = 1010 

65 INTERNAL_ERROR = 1011 

66 SERVICE_RESTART = 1012 

67 TRY_AGAIN_LATER = 1013 

68 BAD_GATEWAY = 1014 

69 TLS_HANDSHAKE = 1015 

70 

71 

72# See https://www.iana.org/assignments/websocket/websocket.xhtml 

73CLOSE_CODE_EXPLANATIONS: dict[int, str] = { 

74 CloseCode.NORMAL_CLOSURE: "OK", 

75 CloseCode.GOING_AWAY: "going away", 

76 CloseCode.PROTOCOL_ERROR: "protocol error", 

77 CloseCode.UNSUPPORTED_DATA: "unsupported data", 

78 CloseCode.NO_STATUS_RCVD: "no status received [internal]", 

79 CloseCode.ABNORMAL_CLOSURE: "abnormal closure [internal]", 

80 CloseCode.INVALID_DATA: "invalid frame payload data", 

81 CloseCode.POLICY_VIOLATION: "policy violation", 

82 CloseCode.MESSAGE_TOO_BIG: "message too big", 

83 CloseCode.MANDATORY_EXTENSION: "mandatory extension", 

84 CloseCode.INTERNAL_ERROR: "internal error", 

85 CloseCode.SERVICE_RESTART: "service restart", 

86 CloseCode.TRY_AGAIN_LATER: "try again later", 

87 CloseCode.BAD_GATEWAY: "bad gateway", 

88 CloseCode.TLS_HANDSHAKE: "TLS handshake failure [internal]", 

89} 

90 

91 

92# Close code that are allowed in a close frame. 

93# Using a set optimizes `code in EXTERNAL_CLOSE_CODES`. 

94EXTERNAL_CLOSE_CODES = { 

95 CloseCode.NORMAL_CLOSURE, 

96 CloseCode.GOING_AWAY, 

97 CloseCode.PROTOCOL_ERROR, 

98 CloseCode.UNSUPPORTED_DATA, 

99 CloseCode.INVALID_DATA, 

100 CloseCode.POLICY_VIOLATION, 

101 CloseCode.MESSAGE_TOO_BIG, 

102 CloseCode.MANDATORY_EXTENSION, 

103 CloseCode.INTERNAL_ERROR, 

104 CloseCode.SERVICE_RESTART, 

105 CloseCode.TRY_AGAIN_LATER, 

106 CloseCode.BAD_GATEWAY, 

107} 

108 

109 

110OK_CLOSE_CODES = { 

111 CloseCode.NORMAL_CLOSURE, 

112 CloseCode.GOING_AWAY, 

113 CloseCode.NO_STATUS_RCVD, 

114} 

115 

116 

117@dataclasses.dataclass 

118class Frame: 

119 """ 

120 WebSocket frame. 

121 

122 Attributes: 

123 opcode: Opcode. 

124 data: Payload data. 

125 fin: FIN bit. 

126 rsv1: RSV1 bit. 

127 rsv2: RSV2 bit. 

128 rsv3: RSV3 bit. 

129 

130 Only these fields are needed. The MASK bit, payload length and masking-key 

131 are handled on the fly when parsing and serializing frames. 

132 

133 """ 

134 

135 opcode: Opcode 

136 data: BytesLike 

137 fin: bool = True 

138 rsv1: bool = False 

139 rsv2: bool = False 

140 rsv3: bool = False 

141 

142 # Configure if you want to see more in logs. Should be a multiple of 3. 

143 MAX_LOG_SIZE = int(os.environ.get("WEBSOCKETS_MAX_LOG_SIZE", "75")) 

144 

145 DEFAULT_IS_TEXT = {TEXT: True, BINARY: False, CLOSE: True} 

146 

147 def __str__(self) -> str: 

148 """ 

149 Return a human-readable representation of a frame. 

150 

151 This function is intended for logging and debugging. It doesn't aim to 

152 support round-tripping because payloads can be too long for displaying 

153 conveniently. Instead, it shows the beginning and the end. It's robust 

154 to incorrect data. 

155 

156 It attempts to decode UTF-8 in data frames, even in binary frames, 

157 because those frequently contain UTF-8 data. 

158 

159 """ 

160 expect_text = self.DEFAULT_IS_TEXT.get(self.opcode) 

161 data_repr, is_text = self._data_repr() 

162 

163 data_type = "" if expect_text == is_text else ("text" if is_text else "binary") 

164 length = f"{len(self.data)} byte{'' if len(self.data) == 1 else 's'}" 

165 non_final = "" if self.fin else "continued" 

166 metadata = ", ".join(filter(None, [data_type, length, non_final])) 

167 

168 return f"{self.opcode.name} {data_repr} [{metadata}]" 

169 

170 def _data_repr(self) -> tuple[str, bool | None]: 

171 """ 

172 Return a human-readable representation of the payload. 

173 

174 Also returns whether the payload is text. 

175 

176 The representation is elided to fit ``MAX_LOG_SIZE``. 

177 

178 This is a helper for the __str__ method. 

179 

180 """ 

181 data_repr: str = "" 

182 is_text: bool | None = None 

183 

184 if not self.data: 

185 return "''", self.DEFAULT_IS_TEXT.get(self.opcode) 

186 

187 # Close frames: parse close code and reason and display them as text. 

188 # Fall back to binary when the payload is malformed. 

189 

190 if self.opcode is CLOSE: 

191 try: 

192 data_repr = str(Close.parse(self.data)) 

193 is_text = True 

194 except (ProtocolError, UnicodeDecodeError): 

195 data_repr = " ".join(f"{byte:02x}" for byte in self.data) 

196 is_text = False 

197 

198 # Control frames: display printable ASCII payloads as text, else binary. 

199 # We could decode UTF-8 payloads, but this causes confusion when random 

200 # 4-bytes binary payloads are accidentally valid UTF-8 sequences. 

201 

202 elif self.opcode in CTRL_OPCODES: 

203 if is_space_or_printable_ascii(self.data): 

204 data_repr = repr(bytes(self.data).decode("ascii")) 

205 is_text = True 

206 else: 

207 data_repr = " ".join(f"{byte:02x}" for byte in self.data) 

208 is_text = False 

209 

210 # Data frames: check whether the payload is UTF-8, regardless of opcode, 

211 # in order to display nicely UTF-8 text in binary frames, and be robust. 

212 # Also support frames fragmented within UTF-8 sequences. 

213 

214 else: 

215 if len(self.data) > 4 * self.MAX_LOG_SIZE: 

216 # Process only the start and end, as the middle will be elided. 

217 # Cast to bytes because self.data could be a memoryview. 

218 data_start = bytes(self.data[: 8 * self.MAX_LOG_SIZE // 3]) 

219 data_end = bytes(self.data[-4 * self.MAX_LOG_SIZE // 3 :]) 

220 is_text = is_utf8_fragment( 

221 data_start, 

222 must_start_clean=self.opcode != CONT, 

223 ) and is_utf8_fragment( 

224 data_end, 

225 must_end_clean=self.fin, 

226 ) 

227 if is_text: 

228 data_repr = repr((data_start + data_end).decode(errors="replace")) 

229 

230 else: 

231 # Cast to bytes because self.data could be a memoryview. 

232 data = bytes(self.data) 

233 is_text = is_utf8_fragment( 

234 data, 

235 must_start_clean=self.opcode != CONT, 

236 must_end_clean=self.fin, 

237 ) 

238 if is_text: 

239 data_repr = repr(data.decode(errors="replace")) 

240 

241 # When the payload is text (except perhaps for boundaries), we have 

242 # enough in ``data_repr``. Do the same when the payload is binary. 

243 

244 if not is_text: 

245 binary = self.data 

246 if len(binary) > self.MAX_LOG_SIZE // 3: 

247 cut = (self.MAX_LOG_SIZE // 3 - 1) // 3 # by default cut = 8 

248 # Encode two dummy bytes to force eliding and adding an ellipsis. 

249 binary = b"".join([binary[: 2 * cut], b"\x00\x00", binary[-cut:]]) 

250 data_repr = " ".join(f"{byte:02x}" for byte in binary) 

251 

252 # Elide the middle of the representation to fit the maximum log size. 

253 

254 if len(data_repr) > self.MAX_LOG_SIZE: 

255 cut = self.MAX_LOG_SIZE // 3 - 1 # by default cut = 24 

256 data_repr = data_repr[: 2 * cut] + "..." + data_repr[-cut:] 

257 

258 return data_repr, is_text 

259 

260 @classmethod 

261 def parse( 

262 cls, 

263 read_exact: Callable[[int], Generator[None, None, bytes | bytearray]], 

264 *, 

265 mask: bool, 

266 max_size: int | None = None, 

267 extensions: Sequence[extensions.Extension] | None = None, 

268 ) -> Generator[None, None, Frame]: 

269 """ 

270 Parse a WebSocket frame. 

271 

272 This is a generator-based coroutine. 

273 

274 Args: 

275 read_exact: Generator-based coroutine that reads the requested 

276 bytes or raises an exception if there isn't enough data. 

277 mask: Whether the frame should be masked i.e. whether the read 

278 happens on the server side. 

279 max_size: Maximum payload size in bytes. 

280 extensions: List of extensions, applied in reverse order. 

281 

282 Raises: 

283 EOFError: If the connection is closed without a full WebSocket frame. 

284 PayloadTooBig: If the frame's payload size exceeds ``max_size``. 

285 ProtocolError: If the frame contains incorrect values. 

286 

287 """ 

288 # Read the header. 

289 data = yield from read_exact(2) 

290 head1, head2 = struct.unpack("!BB", data) 

291 

292 # While not Pythonic, this is marginally faster than calling bool(). 

293 fin = True if head1 & 0b10000000 else False 

294 rsv1 = True if head1 & 0b01000000 else False 

295 rsv2 = True if head1 & 0b00100000 else False 

296 rsv3 = True if head1 & 0b00010000 else False 

297 

298 try: 

299 opcode = Opcode(head1 & 0b00001111) 

300 except ValueError as exc: 

301 raise ProtocolError("invalid opcode") from exc 

302 

303 if (True if head2 & 0b10000000 else False) != mask: 

304 raise ProtocolError("incorrect masking") 

305 

306 length = head2 & 0b01111111 

307 if length == 126: 

308 data = yield from read_exact(2) 

309 (length,) = struct.unpack("!H", data) 

310 elif length == 127: 

311 data = yield from read_exact(8) 

312 (length,) = struct.unpack("!Q", data) 

313 if max_size is not None and length > max_size: 

314 raise PayloadTooBig(length, max_size) 

315 if mask: 

316 mask_bytes = yield from read_exact(4) 

317 

318 # Read the data. 

319 data = yield from read_exact(length) 

320 if mask: 

321 data = apply_mask(data, mask_bytes) 

322 

323 frame = cls(opcode, data, fin, rsv1, rsv2, rsv3) 

324 

325 if extensions is None: 

326 extensions = [] 

327 for extension in reversed(extensions): 

328 frame = extension.decode(frame, max_size=max_size) 

329 

330 frame.check() 

331 

332 return frame 

333 

334 def serialize( 

335 self, 

336 *, 

337 mask: bool, 

338 extensions: Sequence[extensions.Extension] | None = None, 

339 ) -> bytes: 

340 """ 

341 Serialize a WebSocket frame. 

342 

343 Args: 

344 mask: Whether the frame should be masked i.e. whether the write 

345 happens on the client side. 

346 extensions: List of extensions, applied in order. 

347 

348 Raises: 

349 ProtocolError: If the frame contains incorrect values. 

350 

351 """ 

352 self.check() 

353 

354 if extensions is None: 

355 extensions = [] 

356 for extension in extensions: 

357 self = extension.encode(self) 

358 

359 output = io.BytesIO() 

360 

361 # Prepare the header. 

362 head1 = ( 

363 (0b10000000 if self.fin else 0) 

364 | (0b01000000 if self.rsv1 else 0) 

365 | (0b00100000 if self.rsv2 else 0) 

366 | (0b00010000 if self.rsv3 else 0) 

367 | self.opcode 

368 ) 

369 

370 head2 = 0b10000000 if mask else 0 

371 

372 length = len(self.data) 

373 if length < 126: 

374 output.write(struct.pack("!BB", head1, head2 | length)) 

375 elif length < 65536: 

376 output.write(struct.pack("!BBH", head1, head2 | 126, length)) 

377 else: 

378 output.write(struct.pack("!BBQ", head1, head2 | 127, length)) 

379 

380 if mask: 

381 mask_bytes = secrets.token_bytes(4) 

382 output.write(mask_bytes) 

383 

384 # Prepare the data. 

385 data: BytesLike 

386 if mask: 

387 data = apply_mask(self.data, mask_bytes) 

388 else: 

389 data = self.data 

390 output.write(data) 

391 

392 return output.getvalue() 

393 

394 def check(self) -> None: 

395 """ 

396 Check that reserved bits and opcode have acceptable values. 

397 

398 Raises: 

399 ProtocolError: If a reserved bit or the opcode is invalid. 

400 

401 """ 

402 if self.rsv1 or self.rsv2 or self.rsv3: 

403 raise ProtocolError("reserved bits must be 0") 

404 

405 if self.opcode in CTRL_OPCODES: 

406 if len(self.data) > 125: 

407 raise ProtocolError("control frame too long") 

408 if not self.fin: 

409 raise ProtocolError("fragmented control frame") 

410 

411 

412@dataclasses.dataclass 

413class Close: 

414 """ 

415 Code and reason for WebSocket close frames. 

416 

417 Attributes: 

418 code: Close code. 

419 reason: Close reason. 

420 

421 """ 

422 

423 code: CloseCode | int 

424 reason: str 

425 

426 def __str__(self) -> str: 

427 """ 

428 Return a human-readable representation of a close code and reason. 

429 

430 """ 

431 if 3000 <= self.code < 4000: 

432 explanation = "registered" 

433 elif 4000 <= self.code < 5000: 

434 explanation = "private use" 

435 else: 

436 explanation = CLOSE_CODE_EXPLANATIONS.get(self.code, "unknown") 

437 result = f"{self.code} ({explanation})" 

438 

439 if self.reason: 

440 result = f"{result} {self.reason}" 

441 

442 return result 

443 

444 @classmethod 

445 def parse(cls, data: BytesLike) -> Self: 

446 """ 

447 Parse the payload of a close frame. 

448 

449 Args: 

450 data: Payload of the close frame. 

451 

452 Raises: 

453 ProtocolError: If data is ill-formed. 

454 UnicodeDecodeError: If the reason isn't valid UTF-8. 

455 

456 """ 

457 if isinstance(data, memoryview): 

458 raise AssertionError("only compressed outgoing frames use memoryview") 

459 if len(data) >= 2: 

460 (code,) = struct.unpack("!H", data[:2]) 

461 reason = data[2:].decode() 

462 close = cls(code, reason) 

463 close.check() 

464 return close 

465 elif len(data) == 0: 

466 return cls(CloseCode.NO_STATUS_RCVD, "") 

467 else: 

468 raise ProtocolError("close frame too short") 

469 

470 def serialize(self) -> bytes: 

471 """ 

472 Serialize the payload of a close frame. 

473 

474 """ 

475 self.check() 

476 return struct.pack("!H", self.code) + self.reason.encode() 

477 

478 def check(self) -> None: 

479 """ 

480 Check that the close code has a valid value for a close frame. 

481 

482 Raises: 

483 ProtocolError: If the close code is invalid. 

484 

485 """ 

486 if not (self.code in EXTERNAL_CLOSE_CODES or 3000 <= self.code < 5000): 

487 raise ProtocolError("invalid status code") 

488 

489 

490def is_utf8_fragment( 

491 data: bytes, 

492 must_start_clean: bool = False, 

493 must_end_clean: bool = False, 

494) -> bool: 

495 """Guess if data is a fragment of UTF-8 text.""" 

496 # Possible byte sequences for UTF-8 characters are: 

497 # 0xxxxxxx 

498 # 110xxxxx 10xxxxxx 

499 # 1110xxxx 10xxxxxx 10xxxxxx 

500 # 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx 

501 

502 # The algorithm determines ``start`` and ``end`` so that ``data[start:end]`` 

503 # must be a valid UTF-8 sequence for data to be a valid UTF-8 fragment. 

504 

505 start, end = 0, len(data) 

506 

507 if not must_start_clean: 

508 # Remove continuation bytes from the beginning. 

509 max_start = min(3, len(data)) 

510 while start < max_start: 

511 byte = data[start] 

512 

513 # Continuation byte 

514 if byte & 0b11000000 == 0b10000000: 

515 start += 1 

516 continue 

517 

518 break 

519 

520 if not must_end_clean: 

521 # Remove a partial multibyte sequence from the end. 

522 end -= 1 # index of the last byte 

523 min_end = max(len(data) - 4, start) 

524 while end >= min_end: 

525 byte = data[end] 

526 # Continuation byte 

527 if byte & 0b11000000 == 0b10000000: 

528 end -= 1 

529 continue 

530 

531 # ASCII byte 

532 if byte & 0b10000000 == 0b00000000: 

533 seq_len = 1 

534 # Leading byte of a 2-byte sequence 

535 elif byte & 0b11100000 == 0b11000000: 

536 seq_len = 2 

537 # Leading byte of a 3-byte sequence 

538 elif byte & 0b11110000 == 0b11100000: 

539 seq_len = 3 

540 # Leading byte of a 4-byte sequence 

541 elif byte & 0b11111000 == 0b11110000: 

542 seq_len = 4 

543 # Invalid byte 

544 else: 

545 seq_len = 0 

546 

547 # Cut only when there's an incomplete sequence at the end. 

548 if seq_len <= len(data) - end: 

549 end = len(data) 

550 

551 break 

552 

553 try: 

554 text = data[start:end].decode() 

555 except UnicodeDecodeError: 

556 return False 

557 else: 

558 # Non-printable characters signal binary data. 

559 return "\\x" not in repr(text) 

560 

561 

562# At the bottom to break import cycles created by type annotations. 

563from . import extensions # noqa: E402