From dbcdb16d6ef82dc310dc03841008c54bd4b61d59 Mon Sep 17 00:00:00 2001 From: "patchback[bot]" <45432694+patchback[bot]@users.noreply.github.com> Date: Fri, 9 Aug 2024 09:11:19 -0500 Subject: [PATCH] [PR #8667/406cd2c7 backport][3.10] Improve performance of generating random WebSocket mask (#8668) Co-authored-by: J. Nick Koston --- CHANGES/8667.misc.rst | 1 + aiohttp/http_websocket.py | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) create mode 100644 CHANGES/8667.misc.rst diff --git a/CHANGES/8667.misc.rst b/CHANGES/8667.misc.rst new file mode 100644 index 00000000000..1c43b6e069a --- /dev/null +++ b/CHANGES/8667.misc.rst @@ -0,0 +1 @@ +Improved performance of generating random WebSocket mask -- by :user:`bdraco`. diff --git a/aiohttp/http_websocket.py b/aiohttp/http_websocket.py index 39f2e4a5c15..b513a45ebdc 100644 --- a/aiohttp/http_websocket.py +++ b/aiohttp/http_websocket.py @@ -8,6 +8,7 @@ import sys import zlib from enum import IntEnum +from functools import partial from struct import Struct from typing import ( Any, @@ -103,6 +104,7 @@ class WSMsgType(IntEnum): PACK_LEN2 = Struct("!BBH").pack PACK_LEN3 = Struct("!BBQ").pack PACK_CLOSE_CODE = Struct("!H").pack +PACK_RANDBITS = Struct("!L").pack MSG_SIZE: Final[int] = 2**14 DEFAULT_LIMIT: Final[int] = 2**16 @@ -612,7 +614,7 @@ def __init__( self.protocol = protocol self.transport = transport self.use_mask = use_mask - self.randrange = random.randrange + self.get_random_bits = partial(random.getrandbits, 32) self.compress = compress self.notakeover = notakeover self._closing = False @@ -668,8 +670,7 @@ async def _send_frame( else: header = PACK_LEN3(0x80 | rsv | opcode, 127 | mask_bit, msg_length) if use_mask: - mask_int = self.randrange(0, 0xFFFFFFFF) - mask = mask_int.to_bytes(4, "big") + mask = PACK_RANDBITS(self.get_random_bits()) message = bytearray(message) _websocket_mask(mask, message) self._write(header + mask + message)