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

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 

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 "OP_CONT", 

25 "OP_TEXT", 

26 "OP_BINARY", 

27 "OP_CLOSE", 

28 "OP_PING", 

29 "OP_PONG", 

30 "DATA_OPCODES", 

31 "CTRL_OPCODES", 

32 "CloseCode", 

33 "Frame", 

34 "Close", 

35] 

36 

37 

38class Opcode(enum.IntEnum): 

39 """Opcode values for WebSocket frames.""" 

40 

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

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

43 

44 

45OP_CONT = Opcode.CONT 

46OP_TEXT = Opcode.TEXT 

47OP_BINARY = Opcode.BINARY 

48OP_CLOSE = Opcode.CLOSE 

49OP_PING = Opcode.PING 

50OP_PONG = Opcode.PONG 

51 

52DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY 

53CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG 

54 

55 

56class CloseCode(enum.IntEnum): 

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

58 

59 NORMAL_CLOSURE = 1000 

60 GOING_AWAY = 1001 

61 PROTOCOL_ERROR = 1002 

62 UNSUPPORTED_DATA = 1003 

63 # 1004 is reserved 

64 NO_STATUS_RCVD = 1005 

65 ABNORMAL_CLOSURE = 1006 

66 INVALID_DATA = 1007 

67 POLICY_VIOLATION = 1008 

68 MESSAGE_TOO_BIG = 1009 

69 MANDATORY_EXTENSION = 1010 

70 INTERNAL_ERROR = 1011 

71 SERVICE_RESTART = 1012 

72 TRY_AGAIN_LATER = 1013 

73 BAD_GATEWAY = 1014 

74 TLS_HANDSHAKE = 1015 

75 

76 

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

78CLOSE_CODE_EXPLANATIONS: dict[int, str] = { 

79 CloseCode.NORMAL_CLOSURE: "OK", 

80 CloseCode.GOING_AWAY: "going away", 

81 CloseCode.PROTOCOL_ERROR: "protocol error", 

82 CloseCode.UNSUPPORTED_DATA: "unsupported data", 

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

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

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

86 CloseCode.POLICY_VIOLATION: "policy violation", 

87 CloseCode.MESSAGE_TOO_BIG: "message too big", 

88 CloseCode.MANDATORY_EXTENSION: "mandatory extension", 

89 CloseCode.INTERNAL_ERROR: "internal error", 

90 CloseCode.SERVICE_RESTART: "service restart", 

91 CloseCode.TRY_AGAIN_LATER: "try again later", 

92 CloseCode.BAD_GATEWAY: "bad gateway", 

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

94} 

95 

96 

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

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

99EXTERNAL_CLOSE_CODES = { 

100 CloseCode.NORMAL_CLOSURE, 

101 CloseCode.GOING_AWAY, 

102 CloseCode.PROTOCOL_ERROR, 

103 CloseCode.UNSUPPORTED_DATA, 

104 CloseCode.INVALID_DATA, 

105 CloseCode.POLICY_VIOLATION, 

106 CloseCode.MESSAGE_TOO_BIG, 

107 CloseCode.MANDATORY_EXTENSION, 

108 CloseCode.INTERNAL_ERROR, 

109 CloseCode.SERVICE_RESTART, 

110 CloseCode.TRY_AGAIN_LATER, 

111 CloseCode.BAD_GATEWAY, 

112} 

113 

114 

115OK_CLOSE_CODES = { 

116 CloseCode.NORMAL_CLOSURE, 

117 CloseCode.GOING_AWAY, 

118 CloseCode.NO_STATUS_RCVD, 

119} 

120 

121 

122@dataclasses.dataclass 

123class Frame: 

124 """ 

125 WebSocket frame. 

126 

127 Attributes: 

128 opcode: Opcode. 

129 data: Payload data. 

130 fin: FIN bit. 

131 rsv1: RSV1 bit. 

132 rsv2: RSV2 bit. 

133 rsv3: RSV3 bit. 

134 

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

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

137 

138 """ 

139 

140 opcode: Opcode 

141 data: BytesLike 

142 fin: bool = True 

143 rsv1: bool = False 

144 rsv2: bool = False 

145 rsv3: bool = False 

146 

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

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

149 

150 def __str__(self) -> str: 

151 """ 

152 Return a human-readable representation of a frame. 

153 

154 """ 

155 coding = None 

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

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

158 

159 if self.opcode is OP_TEXT: 

160 # Decoding only the beginning and the end is needlessly hard. 

161 # Decode the entire payload then elide later if necessary. 

162 data = repr(bytes(self.data).decode()) 

163 elif self.opcode is OP_BINARY: 

164 # We'll show at most the first 16 bytes and the last 8 bytes. 

165 # Encode just what we need, plus two dummy bytes to elide later. 

166 binary = self.data 

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

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

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

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

171 elif self.opcode is OP_CLOSE: 

172 data = str(Close.parse(self.data)) 

173 elif self.data: 

174 # We don't know if a Continuation frame contains text or binary. 

175 # Ping and Pong frames could contain UTF-8. 

176 # Attempt to decode as UTF-8 and display it as text; fallback to 

177 # binary. If self.data is a memoryview, it has no decode() method, 

178 # which raises AttributeError. 

179 try: 

180 data = repr(bytes(self.data).decode()) 

181 coding = "text" 

182 except (UnicodeDecodeError, AttributeError): 

183 binary = self.data 

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

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

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

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

188 coding = "binary" 

189 else: 

190 data = "''" 

191 

192 if len(data) > self.MAX_LOG_SIZE: 

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

194 data = data[: 2 * cut] + "..." + data[-cut:] 

195 

196 metadata = ", ".join(filter(None, [coding, length, non_final])) 

197 

198 return f"{self.opcode.name} {data} [{metadata}]" 

199 

200 @classmethod 

201 def parse( 

202 cls, 

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

204 *, 

205 mask: bool, 

206 max_size: int | None = None, 

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

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

209 """ 

210 Parse a WebSocket frame. 

211 

212 This is a generator-based coroutine. 

213 

214 Args: 

215 read_exact: Generator-based coroutine that reads the requested 

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

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

218 happens on the server side. 

219 max_size: Maximum payload size in bytes. 

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

221 

222 Raises: 

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

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

225 ProtocolError: If the frame contains incorrect values. 

226 

227 """ 

228 # Read the header. 

229 data = yield from read_exact(2) 

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

231 

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

233 fin = True if head1 & 0b10000000 else False 

234 rsv1 = True if head1 & 0b01000000 else False 

235 rsv2 = True if head1 & 0b00100000 else False 

236 rsv3 = True if head1 & 0b00010000 else False 

237 

238 try: 

239 opcode = Opcode(head1 & 0b00001111) 

240 except ValueError as exc: 

241 raise ProtocolError("invalid opcode") from exc 

242 

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

244 raise ProtocolError("incorrect masking") 

245 

246 length = head2 & 0b01111111 

247 if length == 126: 

248 data = yield from read_exact(2) 

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

250 elif length == 127: 

251 data = yield from read_exact(8) 

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

253 if max_size is not None and length > max_size: 

254 raise PayloadTooBig(length, max_size) 

255 if mask: 

256 mask_bytes = yield from read_exact(4) 

257 

258 # Read the data. 

259 data = yield from read_exact(length) 

260 if mask: 

261 data = apply_mask(data, mask_bytes) 

262 

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

264 

265 if extensions is None: 

266 extensions = [] 

267 for extension in reversed(extensions): 

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

269 

270 frame.check() 

271 

272 return frame 

273 

274 def serialize( 

275 self, 

276 *, 

277 mask: bool, 

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

279 ) -> bytes: 

280 """ 

281 Serialize a WebSocket frame. 

282 

283 Args: 

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

285 happens on the client side. 

286 extensions: List of extensions, applied in order. 

287 

288 Raises: 

289 ProtocolError: If the frame contains incorrect values. 

290 

291 """ 

292 self.check() 

293 

294 if extensions is None: 

295 extensions = [] 

296 for extension in extensions: 

297 self = extension.encode(self) 

298 

299 output = io.BytesIO() 

300 

301 # Prepare the header. 

302 head1 = ( 

303 (0b10000000 if self.fin else 0) 

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

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

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

307 | self.opcode 

308 ) 

309 

310 head2 = 0b10000000 if mask else 0 

311 

312 length = len(self.data) 

313 if length < 126: 

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

315 elif length < 65536: 

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

317 else: 

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

319 

320 if mask: 

321 mask_bytes = secrets.token_bytes(4) 

322 output.write(mask_bytes) 

323 

324 # Prepare the data. 

325 data: BytesLike 

326 if mask: 

327 data = apply_mask(self.data, mask_bytes) 

328 else: 

329 data = self.data 

330 output.write(data) 

331 

332 return output.getvalue() 

333 

334 def check(self) -> None: 

335 """ 

336 Check that reserved bits and opcode have acceptable values. 

337 

338 Raises: 

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

340 

341 """ 

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

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

344 

345 if self.opcode in CTRL_OPCODES: 

346 if len(self.data) > 125: 

347 raise ProtocolError("control frame too long") 

348 if not self.fin: 

349 raise ProtocolError("fragmented control frame") 

350 

351 

352@dataclasses.dataclass 

353class Close: 

354 """ 

355 Code and reason for WebSocket close frames. 

356 

357 Attributes: 

358 code: Close code. 

359 reason: Close reason. 

360 

361 """ 

362 

363 code: CloseCode | int 

364 reason: str 

365 

366 def __str__(self) -> str: 

367 """ 

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

369 

370 """ 

371 if 3000 <= self.code < 4000: 

372 explanation = "registered" 

373 elif 4000 <= self.code < 5000: 

374 explanation = "private use" 

375 else: 

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

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

378 

379 if self.reason: 

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

381 

382 return result 

383 

384 @classmethod 

385 def parse(cls, data: BytesLike) -> Close: 

386 """ 

387 Parse the payload of a close frame. 

388 

389 Args: 

390 data: Payload of the close frame. 

391 

392 Raises: 

393 ProtocolError: If data is ill-formed. 

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

395 

396 """ 

397 if isinstance(data, memoryview): 

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

399 if len(data) >= 2: 

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

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

402 close = cls(code, reason) 

403 close.check() 

404 return close 

405 elif len(data) == 0: 

406 return cls(CloseCode.NO_STATUS_RCVD, "") 

407 else: 

408 raise ProtocolError("close frame too short") 

409 

410 def serialize(self) -> bytes: 

411 """ 

412 Serialize the payload of a close frame. 

413 

414 """ 

415 self.check() 

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

417 

418 def check(self) -> None: 

419 """ 

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

421 

422 Raises: 

423 ProtocolError: If the close code is invalid. 

424 

425 """ 

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

427 raise ProtocolError("invalid status code") 

428 

429 

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

431from . import extensions # noqa: E402