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

1from __future__ import annotations 

2 

3import dataclasses 

4import enum 

5import io 

6import os 

7import secrets 

8import struct 

9from collections.abc import Generator, Sequence 

10from typing import Callable, Self 

11 

12from .exceptions import PayloadTooBig, ProtocolError 

13from .typing import BytesLike 

14 

15 

16try: 

17 from .speedups import apply_mask 

18except ImportError: 

19 from .utils import apply_mask 

20 

21 

22__all__ = [ 

23 "Opcode", 

24 "CloseCode", 

25 "Frame", 

26 "Close", 

27] 

28 

29 

30class Opcode(enum.IntEnum): 

31 """Opcode values for WebSocket frames.""" 

32 

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

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

35 

36 

37CONT = Opcode.CONT 

38TEXT = Opcode.TEXT 

39BINARY = Opcode.BINARY 

40CLOSE = Opcode.CLOSE 

41PING = Opcode.PING 

42PONG = Opcode.PONG 

43 

44DATA_OPCODES = CONT, TEXT, BINARY 

45CTRL_OPCODES = CLOSE, PING, PONG 

46 

47 

48class CloseCode(enum.IntEnum): 

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

50 

51 NORMAL_CLOSURE = 1000 

52 GOING_AWAY = 1001 

53 PROTOCOL_ERROR = 1002 

54 UNSUPPORTED_DATA = 1003 

55 # 1004 is reserved 

56 NO_STATUS_RCVD = 1005 

57 ABNORMAL_CLOSURE = 1006 

58 INVALID_DATA = 1007 

59 POLICY_VIOLATION = 1008 

60 MESSAGE_TOO_BIG = 1009 

61 MANDATORY_EXTENSION = 1010 

62 INTERNAL_ERROR = 1011 

63 SERVICE_RESTART = 1012 

64 TRY_AGAIN_LATER = 1013 

65 BAD_GATEWAY = 1014 

66 TLS_HANDSHAKE = 1015 

67 

68 

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

70CLOSE_CODE_EXPLANATIONS: dict[int, str] = { 

71 CloseCode.NORMAL_CLOSURE: "OK", 

72 CloseCode.GOING_AWAY: "going away", 

73 CloseCode.PROTOCOL_ERROR: "protocol error", 

74 CloseCode.UNSUPPORTED_DATA: "unsupported data", 

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

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

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

78 CloseCode.POLICY_VIOLATION: "policy violation", 

79 CloseCode.MESSAGE_TOO_BIG: "message too big", 

80 CloseCode.MANDATORY_EXTENSION: "mandatory extension", 

81 CloseCode.INTERNAL_ERROR: "internal error", 

82 CloseCode.SERVICE_RESTART: "service restart", 

83 CloseCode.TRY_AGAIN_LATER: "try again later", 

84 CloseCode.BAD_GATEWAY: "bad gateway", 

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

86} 

87 

88 

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

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

91EXTERNAL_CLOSE_CODES = { 

92 CloseCode.NORMAL_CLOSURE, 

93 CloseCode.GOING_AWAY, 

94 CloseCode.PROTOCOL_ERROR, 

95 CloseCode.UNSUPPORTED_DATA, 

96 CloseCode.INVALID_DATA, 

97 CloseCode.POLICY_VIOLATION, 

98 CloseCode.MESSAGE_TOO_BIG, 

99 CloseCode.MANDATORY_EXTENSION, 

100 CloseCode.INTERNAL_ERROR, 

101 CloseCode.SERVICE_RESTART, 

102 CloseCode.TRY_AGAIN_LATER, 

103 CloseCode.BAD_GATEWAY, 

104} 

105 

106 

107OK_CLOSE_CODES = { 

108 CloseCode.NORMAL_CLOSURE, 

109 CloseCode.GOING_AWAY, 

110 CloseCode.NO_STATUS_RCVD, 

111} 

112 

113 

114@dataclasses.dataclass 

115class Frame: 

116 """ 

117 WebSocket frame. 

118 

119 Attributes: 

120 opcode: Opcode. 

121 data: Payload data. 

122 fin: FIN bit. 

123 rsv1: RSV1 bit. 

124 rsv2: RSV2 bit. 

125 rsv3: RSV3 bit. 

126 

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

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

129 

130 """ 

131 

132 opcode: Opcode 

133 data: BytesLike 

134 fin: bool = True 

135 rsv1: bool = False 

136 rsv2: bool = False 

137 rsv3: bool = False 

138 

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

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

141 

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

143 

144 def __str__(self) -> str: 

145 """ 

146 Return a human-readable representation of a frame. 

147 

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

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

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

151 to incorrect data. 

152 

153 It attempts to decode UTF-8 payloads whenever possible, even for binary 

154 frames and control frames, because those frequently contain UTF-8 data. 

155 It applies the same logic to continuation frames, because we don't know 

156 if they continue a text frame or a binary frame. 

157 

158 """ 

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

160 data_repr, is_text = self._data_repr() 

161 

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

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

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

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

166 

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

168 

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

170 """ 

171 Return a human-readable representation of the payload. 

172 

173 Also returns whether the payload is text. 

174 

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

176 

177 This is a helper for the __str__ method. 

178 

179 """ 

180 if not self.data: 

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

182 

183 # Special case for close frames: parse close code and reason. 

184 # Fall back to the standard case if the payload is malformed. 

185 

186 if self.opcode is CLOSE: 

187 try: 

188 return str(Close.parse(self.data)), True 

189 except (ProtocolError, UnicodeDecodeError): 

190 pass 

191 

192 # Guess whether the payload is UTF-8 or binary, regardless of opcode, to 

193 # display UTF-8 text in binary frames nicely and generally to be helpful 

194 # and robust. Also support frames fragmented within UTF-8 sequences. 

195 

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

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

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

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

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

201 is_text = is_utf8_fragment( 

202 data_start, 

203 must_start_clean=self.opcode != CONT, 

204 ) and is_utf8_fragment( 

205 data_end, 

206 must_end_clean=self.fin, 

207 ) 

208 if is_text: 

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

210 

211 else: 

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

213 data = bytes(self.data) 

214 is_text = is_utf8_fragment( 

215 data, 

216 must_start_clean=self.opcode != CONT, 

217 must_end_clean=self.fin, 

218 ) 

219 if is_text: 

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

221 

222 # When the payload is text (except perhaps for boundaries), we decoded 

223 # enough in ``data_repr``. Now, do the same when the payload is binary. 

224 

225 if not is_text: 

226 binary = self.data 

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

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

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

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

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

232 

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

234 

235 if len(data_repr) > self.MAX_LOG_SIZE: 

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

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

238 

239 return data_repr, is_text 

240 

241 @classmethod 

242 def parse( 

243 cls, 

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

245 *, 

246 mask: bool, 

247 max_size: int | None = None, 

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

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

250 """ 

251 Parse a WebSocket frame. 

252 

253 This is a generator-based coroutine. 

254 

255 Args: 

256 read_exact: Generator-based coroutine that reads the requested 

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

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

259 happens on the server side. 

260 max_size: Maximum payload size in bytes. 

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

262 

263 Raises: 

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

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

266 ProtocolError: If the frame contains incorrect values. 

267 

268 """ 

269 # Read the header. 

270 data = yield from read_exact(2) 

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

272 

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

274 fin = True if head1 & 0b10000000 else False 

275 rsv1 = True if head1 & 0b01000000 else False 

276 rsv2 = True if head1 & 0b00100000 else False 

277 rsv3 = True if head1 & 0b00010000 else False 

278 

279 try: 

280 opcode = Opcode(head1 & 0b00001111) 

281 except ValueError as exc: 

282 raise ProtocolError("invalid opcode") from exc 

283 

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

285 raise ProtocolError("incorrect masking") 

286 

287 length = head2 & 0b01111111 

288 if length == 126: 

289 data = yield from read_exact(2) 

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

291 elif length == 127: 

292 data = yield from read_exact(8) 

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

294 if max_size is not None and length > max_size: 

295 raise PayloadTooBig(length, max_size) 

296 if mask: 

297 mask_bytes = yield from read_exact(4) 

298 

299 # Read the data. 

300 data = yield from read_exact(length) 

301 if mask: 

302 data = apply_mask(data, mask_bytes) 

303 

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

305 

306 if extensions is None: 

307 extensions = [] 

308 for extension in reversed(extensions): 

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

310 

311 frame.check() 

312 

313 return frame 

314 

315 def serialize( 

316 self, 

317 *, 

318 mask: bool, 

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

320 ) -> bytes: 

321 """ 

322 Serialize a WebSocket frame. 

323 

324 Args: 

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

326 happens on the client side. 

327 extensions: List of extensions, applied in order. 

328 

329 Raises: 

330 ProtocolError: If the frame contains incorrect values. 

331 

332 """ 

333 self.check() 

334 

335 if extensions is None: 

336 extensions = [] 

337 for extension in extensions: 

338 self = extension.encode(self) 

339 

340 output = io.BytesIO() 

341 

342 # Prepare the header. 

343 head1 = ( 

344 (0b10000000 if self.fin else 0) 

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

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

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

348 | self.opcode 

349 ) 

350 

351 head2 = 0b10000000 if mask else 0 

352 

353 length = len(self.data) 

354 if length < 126: 

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

356 elif length < 65536: 

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

358 else: 

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

360 

361 if mask: 

362 mask_bytes = secrets.token_bytes(4) 

363 output.write(mask_bytes) 

364 

365 # Prepare the data. 

366 data: BytesLike 

367 if mask: 

368 data = apply_mask(self.data, mask_bytes) 

369 else: 

370 data = self.data 

371 output.write(data) 

372 

373 return output.getvalue() 

374 

375 def check(self) -> None: 

376 """ 

377 Check that reserved bits and opcode have acceptable values. 

378 

379 Raises: 

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

381 

382 """ 

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

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

385 

386 if self.opcode in CTRL_OPCODES: 

387 if len(self.data) > 125: 

388 raise ProtocolError("control frame too long") 

389 if not self.fin: 

390 raise ProtocolError("fragmented control frame") 

391 

392 

393@dataclasses.dataclass 

394class Close: 

395 """ 

396 Code and reason for WebSocket close frames. 

397 

398 Attributes: 

399 code: Close code. 

400 reason: Close reason. 

401 

402 """ 

403 

404 code: CloseCode | int 

405 reason: str 

406 

407 def __str__(self) -> str: 

408 """ 

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

410 

411 """ 

412 if 3000 <= self.code < 4000: 

413 explanation = "registered" 

414 elif 4000 <= self.code < 5000: 

415 explanation = "private use" 

416 else: 

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

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

419 

420 if self.reason: 

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

422 

423 return result 

424 

425 @classmethod 

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

427 """ 

428 Parse the payload of a close frame. 

429 

430 Args: 

431 data: Payload of the close frame. 

432 

433 Raises: 

434 ProtocolError: If data is ill-formed. 

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

436 

437 """ 

438 if isinstance(data, memoryview): 

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

440 if len(data) >= 2: 

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

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

443 close = cls(code, reason) 

444 close.check() 

445 return close 

446 elif len(data) == 0: 

447 return cls(CloseCode.NO_STATUS_RCVD, "") 

448 else: 

449 raise ProtocolError("close frame too short") 

450 

451 def serialize(self) -> bytes: 

452 """ 

453 Serialize the payload of a close frame. 

454 

455 """ 

456 self.check() 

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

458 

459 def check(self) -> None: 

460 """ 

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

462 

463 Raises: 

464 ProtocolError: If the close code is invalid. 

465 

466 """ 

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

468 raise ProtocolError("invalid status code") 

469 

470 

471def is_utf8_fragment( 

472 data: bytes, 

473 must_start_clean: bool = False, 

474 must_end_clean: bool = False, 

475) -> bool: 

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

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

478 # 0xxxxxxx 

479 # 110xxxxx 10xxxxxx 

480 # 1110xxxx 10xxxxxx 10xxxxxx 

481 # 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx 

482 

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

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

485 

486 start, end = 0, len(data) 

487 

488 if not must_start_clean: 

489 # Remove continuation bytes from the beginning. 

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

491 while start < max_start: 

492 byte = data[start] 

493 

494 # Continuation byte 

495 if byte & 0b11000000 == 0b10000000: 

496 start += 1 

497 continue 

498 

499 break 

500 

501 if not must_end_clean: 

502 # Remove a partial multibyte sequence from the end. 

503 end -= 1 # index of the last byte 

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

505 while end >= min_end: 

506 byte = data[end] 

507 # Continuation byte 

508 if byte & 0b11000000 == 0b10000000: 

509 end -= 1 

510 continue 

511 

512 # ASCII byte 

513 if byte & 0b10000000 == 0b00000000: 

514 seq_len = 1 

515 # Leading byte of a 2-byte sequence 

516 elif byte & 0b11100000 == 0b11000000: 

517 seq_len = 2 

518 # Leading byte of a 3-byte sequence 

519 elif byte & 0b11110000 == 0b11100000: 

520 seq_len = 3 

521 # Leading byte of a 4-byte sequence 

522 elif byte & 0b11111000 == 0b11110000: 

523 seq_len = 4 

524 # Invalid byte 

525 else: 

526 seq_len = 0 

527 

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

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

530 end = len(data) 

531 

532 break 

533 

534 try: 

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

536 except UnicodeDecodeError: 

537 return False 

538 else: 

539 # Non-printable characters signal binary data. 

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

541 

542 

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

544from . import extensions # noqa: E402