diff --git a/websocket/_handshake.py b/websocket/_handshake.py index 4d20fed..179781d 100644 --- a/websocket/_handshake.py +++ b/websocket/_handshake.py @@ -28,6 +28,8 @@ else: import uuid import hashlib +import hmac +import os from ._logging import * from ._url import * @@ -143,7 +145,7 @@ def _validate(headers, key, subprotocols): value = (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8') hashed = base64encode(hashlib.sha1(value).digest()).strip().lower() - success = (hashed == result) + success = hmac.compare_digest(hashed, result) if success: return True, subproto else: @@ -151,5 +153,5 @@ def _validate(headers, key, subprotocols): def _create_sec_websocket_key(): - uid = uuid.uuid4() - return base64encode(uid.bytes).decode('utf-8').strip() + randomness = os.urandom(16) + return base64encode(randomness).decode('utf-8').strip() diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 6accaef..4573bf7 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -464,12 +464,12 @@ class WebSocketTest(unittest.TestCase): self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello") self.assertRaises(ws.WebSocketConnectionClosedException, s.recv) - def testUUID4(self): - """ WebSocket key should be a UUID4. + def testNonce(self): + """ WebSocket key should be a random 16-byte nonce. """ key = _create_sec_websocket_key() - u = uuid.UUID(bytes=base64decode(key.encode("utf-8"))) - self.assertEqual(4, u.version) + nonce = base64decode(key.encode("utf-8")) + self.assertEqual(16, len(nonce)) class WebSocketAppTest(unittest.TestCase):