def test_round_trip(self): request_headers = Headers() request_key = build_request(request_headers) response_key = check_request(request_headers) self.assertEqual(request_key, response_key) response_headers = Headers() build_response(response_headers, response_key) check_response(response_headers, request_key)
def __init__(self, ws_uri: Union[str, URL], user_id: Optional[int], password: Optional[str], *, state: andesite.StateArgumentType = False, max_connect_attempts: int = None) -> None: self.__ws_uri = str(ws_uri) self.__headers = Headers() if password is not None: self.__headers["Authorization"] = password if user_id is not None: self.user_id = user_id self.__last_connection_id = None self.max_connect_attempts = max_connect_attempts self.web_socket_client = None # can't create the lock here, because if the user uses # asyncio.run and creates the client outside of it, the loop # within the lock will not be the same as the loop used by # asyncio.run (as it creates a new loop every time) self.__connect_lock = None self.__closed = False self.__read_loop = None self._json_encoder = JSONEncoder() self._json_decoder = JSONDecoder() self.state = state
def generate_auth_header(self, username, password): client_digest = hmac.new(password.encode(), msg=b'silenttrinity', digestmod=sha512).hexdigest() header_value = b64encode( f"{username}:{client_digest}".encode()).decode() return Headers({'Authorization': header_value})
def assertValidRequestHeaders(self): """ Provide request headers for modification. Assert that the transformation kept them valid. """ headers = Headers() build_request(headers) yield headers check_request(headers)
def assertValidResponseHeaders(self, key="CSIRmL8dWYxeAdr/XpEHRw=="): """ Provide response headers for modification. Assert that the transformation kept them valid. """ headers = Headers() build_response(headers, key) yield headers check_response(headers, key)
def assertInvalidRequestHeaders(self, exc_type): """ Provide request headers for modification. Assert that the transformation made them invalid. """ headers = Headers() build_request(headers) yield headers assert issubclass(exc_type, InvalidHandshake) with self.assertRaises(exc_type): check_request(headers)
async def process_request(self, path, _): if path.startswith('/ws/'): return None logger.info('Serving %s', path) package, file = self.translate_path(path) if not package: return HTTPStatus.NOT_FOUND, [], b'' header = { 'content-type': guess_type(file), 'cache-control': 'max-age=60', 'X-Frame-Options': "deny" } # serve user configs if package == 'user': try: path = config_directory / file with open(path, 'rb') as f: content = f.read() return HTTPStatus.OK, Headers(header), content except FileNotFoundError: return HTTPStatus.NOT_FOUND, [], b'' # serve other resource files try: try: content = resources.read_binary(package, f'{file}.gz') header['content-encoding'] = 'gzip' except FileNotFoundError: content = resources.read_binary(package, file) return HTTPStatus.OK, Headers(header), content except FileNotFoundError: return HTTPStatus.NOT_FOUND, [], b''
def assertInvalidResponseHeaders(self, exc_type, key="CSIRmL8dWYxeAdr/XpEHRw=="): """ Provide response headers for modification. Assert that the transformation made them invalid. """ headers = Headers() build_response(headers, key) yield headers assert issubclass(exc_type, InvalidHandshake) with self.assertRaises(exc_type): check_response(headers, key)
async def connect(cls, client): """ :param client: An instance of tclib.TinychatClient. :type client: tclib.TinychatClient :return: :rtype: """ ws = None log.info('requesting gateway') gateway = await connect_details(client.room) log.debug(f'gateway: {gateway}') if len(gateway) == 2: try: ws = await asyncio.wait_for(websockets.connect( uri=gateway['endpoint'], origin='https://tinychat.com', subprotocols=['tc'], extra_headers=Headers(tc_headers()), loop=client.loop, klass=cls), timeout=60, loop=client.loop) # check for connect timeout error. except asyncio.TimeoutError: log.warning('timed out waiting for client connect.') # return await cls.from_client(resume=resume) # send the join message. await ws.join(client, gateway['token']) # assign the dispatcher. # Sync From Async ws._dispatch = client.dispatch try: # make sure the websocket is in a open state. await ws.ensure_open() except websockets.exceptions.ConnectionClosed as e: log.warning(f'connection was closed: {e.code} {e.reason}') else: log.info(f'websocket connected to: {gateway["endpoint"]}') else: log.warning(f'failed to get gateway: {gateway}') return ws
async def _listen_forever(self): while True: # outer loop restarted every time the connection fails try: headers = Headers() headers['Authorization'] = f'Token {self.settings.session_key}' ssl_args = dict() if not self.settings.verify_certificate and 'wss:' in self.uri: ssl_args['ssl'] = ssl.SSLContext() async with websockets.connect(self.uri, extra_headers=headers, **ssl_args) as ws: self.logger.info(f'Websocket established with {self.uri}') self.logger.debug( f'Ping response: {json.dumps(self.ping_data)}') await ws.send(str(AgoraWebsocketMessage(self.ping_data))) self.logger.info( f'App is running --> listening for server messages...') while True: # listener loop try: msg = await ws.recv() except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed) as e: try: self.logger.warning( f'Connection probably closed (sending ping): {str(e)}' ) pong = await ws.ping() await asyncio.wait_for( pong, timeout=self.ping_timeout) continue except: self.logger.warning( f'Connection closed (trying again in {self.sleep_time}s): {str(e)}' ) await asyncio.sleep(self.sleep_time) break # inner loop await self._process_message(msg, ws) except socket.gaierror as e: self.logger.warning(f'Connection closed: {str(e)}') self.logger.warning(f'Trying to reconnect') continue except ConnectionRefusedError as e: self.logger.warning(f'Connection refused: {str(e)}') self.logger.warning(f'Trying to reconnect') continue
def __init__(self, addr: str = None, key: str = None, close_timeout: float = IO_TIMEOUT): """ Create a web socket client to connect platform. Parameters ---------- addr: str address for a service, like ws://host:port key: str key represents API key, which is used for verifying identity. """ self.__addr = addr if addr is not None else self.ADDRESS self.__key = key self.header = Headers(Authorization="Bearer " + key) if key is not None else None self.close_timeout = close_timeout self.__builder = flatbuffers.Builder(1400) self.__data_handler = None self.__queue = asyncio.Queue(10)
async def listen(self, endpoint_uri: str): uri = f"{self.proto}://{self.host}/{self.api_version}/{endpoint_uri}" console.print(f"Trying to connect to {uri} now...") headers = Headers(authorization=f"{self.token}") headers.update(**self.user_agent) sim_time: Optional[int] = None async with websockets.connect(uri, extra_headers=headers) as websocket: console.print(f"Connected to {uri}") self._connected = True async for message in websocket: try: message = json.loads(message) content = message["data"] try: sim_time = content["time"] except (KeyError, TypeError): pass if sim_time is not None: message["sim_time"] = sim_time await self.queue.put(message) except websockets.exceptions.ConnectionClosedOK: self.do_listen = False console.print("Websocket connection closed")
class ClientServerTests(unittest.TestCase): secure = False def setUp(self): self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) def tearDown(self): self.loop.close() def run_loop_once(self): # Process callbacks scheduled with call_soon by appending a callback # to stop the event loop then running it until it hits that callback. self.loop.call_soon(self.loop.stop) self.loop.run_forever() def start_server(self, **kwds): # Disable compression by default in tests. kwds.setdefault('compression', None) # Disable pings by default in tests. kwds.setdefault('ping_interval', None) start_server = serve(handler, 'localhost', 0, **kwds) self.server = self.loop.run_until_complete(start_server) def start_client(self, resource_name='/', user_info=None, **kwds): # Disable compression by default in tests. kwds.setdefault('compression', None) # Disable pings by default in tests. kwds.setdefault('ping_interval', None) secure = kwds.get('ssl') is not None server_uri = get_server_uri(self.server, secure, resource_name, user_info) start_client = connect(server_uri, **kwds) self.client = self.loop.run_until_complete(start_client) def stop_client(self): try: self.loop.run_until_complete( asyncio.wait_for(self.client.close_connection_task, timeout=1)) except asyncio.TimeoutError: # pragma: no cover self.fail("Client failed to stop") def stop_server(self): self.server.close() try: self.loop.run_until_complete( asyncio.wait_for(self.server.wait_closed(), timeout=1)) except asyncio.TimeoutError: # pragma: no cover self.fail("Server failed to stop") @contextlib.contextmanager def temp_server(self, **kwds): with temp_test_server(self, **kwds): yield @contextlib.contextmanager def temp_client(self, *args, **kwds): with temp_test_client(self, *args, **kwds): yield @with_server() @with_client() def test_basic(self): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") def test_server_close_while_client_connected(self): with self.temp_server(loop=self.loop): # This endpoint waits just a bit when the connection is canceled # in order to test that wait_closed() really waits for completion. self.start_client('/slow_stop') with self.assertRaises(ConnectionClosed): self.loop.run_until_complete(self.client.recv()) # Connection ends with 1001 going away. self.assertEqual(self.client.close_code, 1001) def test_explicit_event_loop(self): with self.temp_server(loop=self.loop): with self.temp_client(loop=self.loop): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") # The way the legacy SSL implementation wraps sockets makes it extremely # hard to write a test for Python 3.4. @unittest.skipIf(sys.version_info[:2] <= (3, 4), 'this test requires Python 3.5+') @with_server() def test_explicit_socket(self): class TrackedSocket(socket.socket): def __init__(self, *args, **kwargs): self.used_for_read = False self.used_for_write = False super().__init__(*args, **kwargs) def recv(self, *args, **kwargs): self.used_for_read = True return super().recv(*args, **kwargs) def send(self, *args, **kwargs): self.used_for_write = True return super().send(*args, **kwargs) server_socket = [ sock for sock in self.server.sockets if sock.family == socket.AF_INET ][0] client_socket = TrackedSocket(socket.AF_INET, socket.SOCK_STREAM) client_socket.connect(server_socket.getsockname()) try: self.assertFalse(client_socket.used_for_read) self.assertFalse(client_socket.used_for_write) with self.temp_client( sock=client_socket, # "You must set server_hostname when using ssl without a host" server_hostname='localhost' if self.secure else None, ): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") self.assertTrue(client_socket.used_for_read) self.assertTrue(client_socket.used_for_write) finally: client_socket.close() @unittest.skipUnless(hasattr(socket, 'AF_UNIX'), 'this test requires Unix sockets') def test_unix_socket(self): with tempfile.TemporaryDirectory() as temp_dir: path = bytes(pathlib.Path(temp_dir) / 'websockets') # Like self.start_server() but with unix_serve(). unix_server = unix_serve(handler, path) self.server = self.loop.run_until_complete(unix_server) client_socket = socket.socket(socket.AF_UNIX) client_socket.connect(path) try: with self.temp_client(sock=client_socket): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") finally: client_socket.close() self.stop_server() @with_server() @with_client('/attributes') def test_protocol_attributes(self): # The test could be connecting with IPv6 or IPv4. expected_client_attrs = [ server_socket.getsockname()[:2] + (self.secure, ) for server_socket in self.server.sockets ] client_attrs = (self.client.host, self.client.port, self.client.secure) self.assertIn(client_attrs, expected_client_attrs) expected_server_attrs = ('localhost', 0, self.secure) server_attrs = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_attrs, repr(expected_server_attrs)) @with_server() @with_client('/path') def test_protocol_path(self): client_path = self.client.path self.assertEqual(client_path, '/path') server_path = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_path, '/path') @with_server() @with_client('/headers', user_info=('user', 'pass')) def test_protocol_basic_auth(self): self.assertEqual(self.client.request_headers['Authorization'], 'Basic dXNlcjpwYXNz') @with_server() @with_client('/headers') def test_protocol_headers(self): client_req = self.client.request_headers client_resp = self.client.response_headers self.assertEqual(client_req['User-Agent'], USER_AGENT) self.assertEqual(client_resp['Server'], USER_AGENT) server_req = self.loop.run_until_complete(self.client.recv()) server_resp = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_req, repr(client_req)) self.assertEqual(server_resp, repr(client_resp)) @with_server() @with_client('/headers', extra_headers=Headers({'X-Spam': 'Eggs'})) def test_protocol_custom_request_headers(self): req_headers = self.loop.run_until_complete(self.client.recv()) self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", req_headers) @with_server() @with_client('/headers', extra_headers={'X-Spam': 'Eggs'}) def test_protocol_custom_request_headers_dict(self): req_headers = self.loop.run_until_complete(self.client.recv()) self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", req_headers) @with_server() @with_client('/headers', extra_headers=[('X-Spam', 'Eggs')]) def test_protocol_custom_request_headers_list(self): req_headers = self.loop.run_until_complete(self.client.recv()) self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", req_headers) @with_server() @with_client('/headers', extra_headers=[('User-Agent', 'Eggs')]) def test_protocol_custom_request_user_agent(self): req_headers = self.loop.run_until_complete(self.client.recv()) self.loop.run_until_complete(self.client.recv()) self.assertEqual(req_headers.count("User-Agent"), 1) self.assertIn("('User-Agent', 'Eggs')", req_headers) @with_server(extra_headers=lambda p, r: Headers({'X-Spam': 'Eggs'})) @with_client('/headers') def test_protocol_custom_response_headers_callable(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers=lambda p, r: {'X-Spam': 'Eggs'}) @with_client('/headers') def test_protocol_custom_response_headers_callable_dict(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers=lambda p, r: [('X-Spam', 'Eggs')]) @with_client('/headers') def test_protocol_custom_response_headers_callable_list(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers=Headers({'X-Spam': 'Eggs'})) @with_client('/headers') def test_protocol_custom_response_headers(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers={'X-Spam': 'Eggs'}) @with_client('/headers') def test_protocol_custom_response_headers_dict(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers=[('X-Spam', 'Eggs')]) @with_client('/headers') def test_protocol_custom_response_headers_list(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertIn("('X-Spam', 'Eggs')", resp_headers) @with_server(extra_headers=[('Server', 'Eggs')]) @with_client('/headers') def test_protocol_custom_response_user_agent(self): self.loop.run_until_complete(self.client.recv()) resp_headers = self.loop.run_until_complete(self.client.recv()) self.assertEqual(resp_headers.count("Server"), 1) self.assertIn("('Server', 'Eggs')", resp_headers) def make_http_request(self, path='/'): # Set url to 'https?://<host>:<port><path>'. url = get_server_uri(self.server, resource_name=path, secure=self.secure) url = url.replace('ws', 'http') if self.secure: open_health_check = functools.partial(urllib.request.urlopen, url, context=self.client_context) else: open_health_check = functools.partial(urllib.request.urlopen, url) return self.loop.run_in_executor(None, open_health_check) @with_server(create_protocol=HealthCheckServerProtocol) def test_http_request_http_endpoint(self): # Making a HTTP request to a HTTP endpoint succeeds. response = self.loop.run_until_complete( self.make_http_request('/__health__/')) with contextlib.closing(response): self.assertEqual(response.code, 200) self.assertEqual(response.read(), b'status = green\n') @with_server(create_protocol=HealthCheckServerProtocol) def test_http_request_ws_endpoint(self): # Making a HTTP request to a WS endpoint fails. with self.assertRaises(urllib.error.HTTPError) as raised: self.loop.run_until_complete(self.make_http_request()) self.assertEqual(raised.exception.code, 426) self.assertEqual(raised.exception.headers['Upgrade'], 'websocket') @with_server(create_protocol=HealthCheckServerProtocol) def test_ws_connection_http_endpoint(self): # Making a WS connection to a HTTP endpoint fails. with self.assertRaises(InvalidStatusCode) as raised: self.start_client('/__health__/') self.assertEqual(raised.exception.status_code, 200) @with_server(create_protocol=HealthCheckServerProtocol) def test_ws_connection_ws_endpoint(self): # Making a WS connection to a WS endpoint succeeds. self.start_client() self.loop.run_until_complete(self.client.send("Hello!")) self.loop.run_until_complete(self.client.recv()) self.stop_client() def assert_client_raises_code(self, status_code): with self.assertRaises(InvalidStatusCode) as raised: self.start_client() self.assertEqual(raised.exception.status_code, status_code) @with_server(create_protocol=UnauthorizedServerProtocol) def test_server_create_protocol(self): self.assert_client_raises_code(401) @with_server(create_protocol=( lambda *args, **kwargs: UnauthorizedServerProtocol(*args, **kwargs))) def test_server_create_protocol_function(self): self.assert_client_raises_code(401) @with_server(klass=UnauthorizedServerProtocol) def test_server_klass_backwards_compatibility(self): self.assert_client_raises_code(401) @with_server(create_protocol=ForbiddenServerProtocol, klass=UnauthorizedServerProtocol) def test_server_create_protocol_over_klass(self): self.assert_client_raises_code(403) @with_server() @with_client('/path', create_protocol=FooClientProtocol) def test_client_create_protocol(self): self.assertIsInstance(self.client, FooClientProtocol) @with_server() @with_client( '/path', create_protocol=( lambda *args, **kwargs: FooClientProtocol(*args, **kwargs)), ) def test_client_create_protocol_function(self): self.assertIsInstance(self.client, FooClientProtocol) @with_server() @with_client('/path', klass=FooClientProtocol) def test_client_klass(self): self.assertIsInstance(self.client, FooClientProtocol) @with_server() @with_client('/path', create_protocol=BarClientProtocol, klass=FooClientProtocol) def test_client_create_protocol_over_klass(self): self.assertIsInstance(self.client, BarClientProtocol) @with_server(close_timeout=7) @with_client('/close_timeout') def test_server_close_timeout(self): close_timeout = self.loop.run_until_complete(self.client.recv()) self.assertEqual(eval(close_timeout), 7) @with_server(timeout=6) @with_client('/close_timeout') def test_server_timeout_backwards_compatibility(self): close_timeout = self.loop.run_until_complete(self.client.recv()) self.assertEqual(eval(close_timeout), 6) @with_server(close_timeout=7, timeout=6) @with_client('/close_timeout') def test_server_close_timeout_over_timeout(self): close_timeout = self.loop.run_until_complete(self.client.recv()) self.assertEqual(eval(close_timeout), 7) @with_server() @with_client('/close_timeout', close_timeout=7) def test_client_close_timeout(self): self.assertEqual(self.client.close_timeout, 7) @with_server() @with_client('/close_timeout', timeout=6) def test_client_timeout_backwards_compatibility(self): self.assertEqual(self.client.close_timeout, 6) @with_server() @with_client('/close_timeout', close_timeout=7, timeout=6) def test_client_close_timeout_over_timeout(self): self.assertEqual(self.client.close_timeout, 7) @with_server() @with_client('/extensions') def test_no_extension(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([])) self.assertEqual(repr(self.client.extensions), repr([])) @with_server(extensions=[ServerNoOpExtensionFactory()]) @with_client('/extensions', extensions=[ClientNoOpExtensionFactory()]) def test_extension(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([NoOpExtension()])) self.assertEqual(repr(self.client.extensions), repr([NoOpExtension()])) @with_server() @with_client('/extensions', extensions=[ClientNoOpExtensionFactory()]) def test_extension_not_accepted(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([])) self.assertEqual(repr(self.client.extensions), repr([])) @with_server(extensions=[ServerNoOpExtensionFactory()]) @with_client('/extensions') def test_extension_not_requested(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([])) self.assertEqual(repr(self.client.extensions), repr([])) @with_server(extensions=[ServerNoOpExtensionFactory([('foo', None)])]) def test_extension_client_rejection(self): with self.assertRaises(NegotiationError): self.start_client('/extensions', extensions=[ClientNoOpExtensionFactory()]) @with_server(extensions=[ # No match because the client doesn't send client_max_window_bits. ServerPerMessageDeflateFactory(client_max_window_bits=10), ServerPerMessageDeflateFactory(), ]) @with_client('/extensions', extensions=[ClientPerMessageDeflateFactory()]) def test_extension_no_match_then_match(self): # The order requested by the client has priority. server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([PerMessageDeflate(False, False, 15, 15)])) self.assertEqual( repr(self.client.extensions), repr([PerMessageDeflate(False, False, 15, 15)]), ) @with_server(extensions=[ServerPerMessageDeflateFactory()]) @with_client('/extensions', extensions=[ClientNoOpExtensionFactory()]) def test_extension_mismatch(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([])) self.assertEqual(repr(self.client.extensions), repr([])) @with_server(extensions=[ ServerNoOpExtensionFactory(), ServerPerMessageDeflateFactory() ]) @with_client( '/extensions', extensions=[ ClientPerMessageDeflateFactory(), ClientNoOpExtensionFactory() ], ) def test_extension_order(self): # The order requested by the client has priority. server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual( server_extensions, repr([PerMessageDeflate(False, False, 15, 15), NoOpExtension()]), ) self.assertEqual( repr(self.client.extensions), repr([PerMessageDeflate(False, False, 15, 15), NoOpExtension()]), ) @with_server(extensions=[ServerNoOpExtensionFactory()]) @unittest.mock.patch.object(WebSocketServerProtocol, 'process_extensions') def test_extensions_error(self, _process_extensions): _process_extensions.return_value = 'x-no-op', [NoOpExtension()] with self.assertRaises(NegotiationError): self.start_client('/extensions', extensions=[ClientPerMessageDeflateFactory()]) @with_server(extensions=[ServerNoOpExtensionFactory()]) @unittest.mock.patch.object(WebSocketServerProtocol, 'process_extensions') def test_extensions_error_no_extensions(self, _process_extensions): _process_extensions.return_value = 'x-no-op', [NoOpExtension()] with self.assertRaises(InvalidHandshake): self.start_client('/extensions') @with_server(compression='deflate') @with_client('/extensions', compression='deflate') def test_compression_deflate(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([PerMessageDeflate(False, False, 15, 15)])) self.assertEqual( repr(self.client.extensions), repr([PerMessageDeflate(False, False, 15, 15)]), ) @with_server( extensions=[ ServerPerMessageDeflateFactory(client_no_context_takeover=True, server_max_window_bits=10) ], compression='deflate', # overridden by explicit config ) @with_client( '/extensions', extensions=[ ClientPerMessageDeflateFactory(server_no_context_takeover=True, client_max_window_bits=12) ], compression='deflate', # overridden by explicit config ) def test_compression_deflate_and_explicit_config(self): server_extensions = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_extensions, repr([PerMessageDeflate(True, True, 12, 10)])) self.assertEqual(repr(self.client.extensions), repr([PerMessageDeflate(True, True, 10, 12)])) def test_compression_unsupported_server(self): with self.assertRaises(ValueError): self.loop.run_until_complete(self.start_server(compression='xz')) @with_server() def test_compression_unsupported_client(self): with self.assertRaises(ValueError): self.loop.run_until_complete(self.start_client(compression='xz')) @with_server() @with_client('/subprotocol') def test_no_subprotocol(self): server_subprotocol = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_subprotocol, repr(None)) self.assertEqual(self.client.subprotocol, None) @with_server(subprotocols=['superchat', 'chat']) @with_client('/subprotocol', subprotocols=['otherchat', 'chat']) def test_subprotocol(self): server_subprotocol = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_subprotocol, repr('chat')) self.assertEqual(self.client.subprotocol, 'chat') @with_server(subprotocols=['superchat']) @with_client('/subprotocol', subprotocols=['otherchat']) def test_subprotocol_not_accepted(self): server_subprotocol = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_subprotocol, repr(None)) self.assertEqual(self.client.subprotocol, None) @with_server() @with_client('/subprotocol', subprotocols=['otherchat', 'chat']) def test_subprotocol_not_offered(self): server_subprotocol = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_subprotocol, repr(None)) self.assertEqual(self.client.subprotocol, None) @with_server(subprotocols=['superchat', 'chat']) @with_client('/subprotocol') def test_subprotocol_not_requested(self): server_subprotocol = self.loop.run_until_complete(self.client.recv()) self.assertEqual(server_subprotocol, repr(None)) self.assertEqual(self.client.subprotocol, None) @with_server(subprotocols=['superchat']) @unittest.mock.patch.object(WebSocketServerProtocol, 'process_subprotocol') def test_subprotocol_error(self, _process_subprotocol): _process_subprotocol.return_value = 'superchat' with self.assertRaises(NegotiationError): self.start_client('/subprotocol', subprotocols=['otherchat']) self.run_loop_once() @with_server(subprotocols=['superchat']) @unittest.mock.patch.object(WebSocketServerProtocol, 'process_subprotocol') def test_subprotocol_error_no_subprotocols(self, _process_subprotocol): _process_subprotocol.return_value = 'superchat' with self.assertRaises(InvalidHandshake): self.start_client('/subprotocol') self.run_loop_once() @with_server(subprotocols=['superchat', 'chat']) @unittest.mock.patch.object(WebSocketServerProtocol, 'process_subprotocol') def test_subprotocol_error_two_subprotocols(self, _process_subprotocol): _process_subprotocol.return_value = 'superchat, chat' with self.assertRaises(InvalidHandshake): self.start_client('/subprotocol', subprotocols=['superchat', 'chat']) self.run_loop_once() @with_server() @unittest.mock.patch('websockets.server.read_request') def test_server_receives_malformed_request(self, _read_request): _read_request.side_effect = ValueError("read_request failed") with self.assertRaises(InvalidHandshake): self.start_client() @with_server() @unittest.mock.patch('websockets.client.read_response') def test_client_receives_malformed_response(self, _read_response): _read_response.side_effect = ValueError("read_response failed") with self.assertRaises(InvalidHandshake): self.start_client() self.run_loop_once() @with_server() @unittest.mock.patch('websockets.client.build_request') def test_client_sends_invalid_handshake_request(self, _build_request): def wrong_build_request(headers): return '42' _build_request.side_effect = wrong_build_request with self.assertRaises(InvalidHandshake): self.start_client() @with_server() @unittest.mock.patch('websockets.server.build_response') def test_server_sends_invalid_handshake_response(self, _build_response): def wrong_build_response(headers, key): return build_response(headers, '42') _build_response.side_effect = wrong_build_response with self.assertRaises(InvalidHandshake): self.start_client() @with_server() @unittest.mock.patch('websockets.client.read_response') def test_server_does_not_switch_protocols(self, _read_response): @asyncio.coroutine def wrong_read_response(stream): status_code, headers = yield from read_response(stream) return 400, headers _read_response.side_effect = wrong_read_response with self.assertRaises(InvalidStatusCode): self.start_client() self.run_loop_once() @with_server() @unittest.mock.patch( 'websockets.server.WebSocketServerProtocol.process_request') def test_server_error_in_handshake(self, _process_request): _process_request.side_effect = Exception("process_request crashed") with self.assertRaises(InvalidHandshake): self.start_client() @with_server() @unittest.mock.patch('websockets.server.WebSocketServerProtocol.send') def test_server_handler_crashes(self, send): send.side_effect = ValueError("send failed") with self.temp_client(): self.loop.run_until_complete(self.client.send("Hello!")) with self.assertRaises(ConnectionClosed): self.loop.run_until_complete(self.client.recv()) # Connection ends with an unexpected error. self.assertEqual(self.client.close_code, 1011) @with_server() @unittest.mock.patch('websockets.server.WebSocketServerProtocol.close') def test_server_close_crashes(self, close): close.side_effect = ValueError("close failed") with self.temp_client(): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") # Connection ends with an abnormal closure. self.assertEqual(self.client.close_code, 1006) @with_server() @with_client() @unittest.mock.patch.object(WebSocketClientProtocol, 'handshake') def test_client_closes_connection_before_handshake(self, handshake): # We have mocked the handshake() method to prevent the client from # performing the opening handshake. Force it to close the connection. self.client.writer.close() # The server should stop properly anyway. It used to hang because the # task handling the connection was waiting for the opening handshake. @with_server() @unittest.mock.patch('websockets.server.read_request') def test_server_shuts_down_during_opening_handshake(self, _read_request): _read_request.side_effect = asyncio.CancelledError self.server.closing = True with self.assertRaises(InvalidHandshake) as raised: self.start_client() # Opening handshake fails with 503 Service Unavailable self.assertEqual(str(raised.exception), "Status code not 101: 503") @with_server() def test_server_shuts_down_during_connection_handling(self): with self.temp_client(): self.server.close() with self.assertRaises(ConnectionClosed): self.loop.run_until_complete(self.client.recv()) # Websocket connection terminates with 1001 Going Away. self.assertEqual(self.client.close_code, 1001) @with_server() @unittest.mock.patch('websockets.server.WebSocketServerProtocol.close') def test_server_shuts_down_during_connection_close(self, _close): _close.side_effect = asyncio.CancelledError self.server.closing = True with self.temp_client(): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") # Websocket connection terminates abnormally. self.assertEqual(self.client.close_code, 1006) @with_server(create_protocol=ForbiddenServerProtocol) def test_invalid_status_error_during_client_connect(self): with self.assertRaises(InvalidStatusCode) as raised: self.start_client() exception = raised.exception self.assertEqual(str(exception), "Status code not 101: 403") self.assertEqual(exception.status_code, 403) @with_server() @unittest.mock.patch( 'websockets.server.WebSocketServerProtocol.write_http_response') @unittest.mock.patch( 'websockets.server.WebSocketServerProtocol.read_http_request') def test_connection_error_during_opening_handshake(self, _read_http_request, _write_http_response): _read_http_request.side_effect = ConnectionError # This exception is currently platform-dependent. It was observed to # be ConnectionResetError on Linux in the non-SSL case, and # InvalidMessage otherwise (including both Linux and macOS). This # doesn't matter though since this test is primarily for testing a # code path on the server side. with self.assertRaises(Exception): self.start_client() # No response must not be written if the network connection is broken. _write_http_response.assert_not_called() @with_server() @unittest.mock.patch('websockets.server.WebSocketServerProtocol.close') def test_connection_error_during_closing_handshake(self, close): close.side_effect = ConnectionError with self.temp_client(): self.loop.run_until_complete(self.client.send("Hello!")) reply = self.loop.run_until_complete(self.client.recv()) self.assertEqual(reply, "Hello!") # Connection ends with an abnormal closure. self.assertEqual(self.client.close_code, 1006)
def process_request(self, path, request_headers): # Test returning headers as a Headers instance (1/3) return UNAUTHORIZED, Headers([('X-Access', 'denied')]), b''
async def handshake( self, wsuri: websockets.WebSocketURI, origin: Optional[websockets.typing.Origin] = None, available_extensions: Optional[ Sequence[websockets.extensions.base.ClientExtensionFactory] ] = None, available_subprotocols: Optional[ Sequence[websockets.typing.Subprotocol] ] = None, extra_headers: Optional[websockets.http.HeadersLike] = None, ): """ Perform the client side of the opening handshake. :param origin: sets the Origin HTTP header :param available_extensions: list of supported extensions in the order in which they should be used :param available_subprotocols: list of supported subprotocols in order of decreasing preference :param extra_headers: sets additional HTTP request headers; it must be a :class:`~websockets.http.Headers` instance, a :class:`~collections.abc.Mapping`, or an iterable of ``(name, value)`` pairs :raises ~websockets.exceptions.InvalidHandshake: if the handshake fails """ request_headers = Headers() if wsuri.port == (443 if wsuri.secure else 80): # pragma: no cover request_headers["Host"] = wsuri.host else: request_headers["Host"] = f"{wsuri.host}:{wsuri.port}" # if wsuri.user_info: # request_headers["Authorization"] = build_authorization_basic( # *wsuri.user_info # ) if origin is not None: request_headers["Origin"] = origin key = build_request(request_headers) if available_extensions is not None: extensions_header = websockets.headers.build_extension( [ (extension_factory.name, extension_factory.get_request_params()) for extension_factory in available_extensions ] ) request_headers["Sec-WebSocket-Extensions"] = extensions_header if available_subprotocols is not None: protocol_header = websockets.headers.build_subprotocol( available_subprotocols ) request_headers["Sec-WebSocket-Protocol"] = protocol_header if extra_headers is not None: if isinstance(extra_headers, Headers): extra_headers = extra_headers.raw_items() elif isinstance(extra_headers, collections.abc.Mapping): extra_headers = extra_headers.items() for name, value in extra_headers: request_headers[name] = value request_headers.setdefault("User-Agent", websockets.http.USER_AGENT) req = Request("GET", f"https://{wsuri.host}/", data="", headers=request_headers) prepared_request = req.prepare() encrypted_request = aws_auth(prepared_request) self.write_http_request(wsuri.resource_name, Headers(encrypted_request.headers)) status_code, response_headers = await self.read_http_response() if status_code in (301, 302, 303, 307, 308): if "Location" not in response_headers: raise websockets.exceptions.InvalidHeader("Location") raise websockets.exceptions.RedirectHandshake(response_headers["Location"]) elif status_code != 101: raise websockets.exceptions.InvalidStatusCode(status_code) websockets.handshake.check_response(response_headers, key) self.extensions = self.process_extensions( response_headers, available_extensions ) self.subprotocol = self.process_subprotocol( response_headers, available_subprotocols ) self.connection_open()
async def handleHTTP( #type: ignore self, path: str, request_headers: Headers) -> Tuple[HTTPStatus, Union[ Headers, Mapping[str, str], Iterable[Tuple[str, str]]], bytes]: status: HTTPStatus if path == "/token": try: token = request_headers['Token'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No token header present' return (status, headers, b'') try: await self.database.tokenSignIn(token=token) except database.InvalidCredentials: status = HTTPStatus.UNAUTHORIZED headers = Headers() headers['error'] = 'Invalid token!' return (status, headers, b'') elif path == "/signIn": try: email = request_headers['Email'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No email header present' return (status, headers, b'') try: password = request_headers['Password'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No password header present' return (status, headers, b'') try: await self.database.signIn(email=email, password=password) except database.InvalidCredentials: status = HTTPStatus.UNAUTHORIZED headers = Headers() headers['error'] = 'Invalid token!' return (status, headers, b'') elif path == "/signUp": try: password = request_headers['Password'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No password header present' return (status, headers, b'') try: email = request_headers['Email'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No email header present' return (status, headers, b'') try: username = request_headers['Username'] except KeyError: status = HTTPStatus.BAD_REQUEST headers = Headers() headers['error'] = 'No username header present' return (status, headers, b'') # try: # response = request_headers['Response'] # except KeyError: # status = HTTPStatus.BAD_REQUEST # headers = Headers() # headers['error'] = 'No response header present' # return (status, headers, b'') # if not (await self.recaptcha.verify(response)): # status = HTTPStatus.UNAUTHORIZED # headers = Headers() # headers['error'] = 'Invalid response token.' # return (status, headers, b'') try: await self.database.signUp(username=username, password=password, email=email) except database.UserExists: status = HTTPStatus.FORBIDDEN headers = Headers() headers['error'] = 'Email is taken.' return (status, headers, b'') else: status = HTTPStatus.NOT_FOUND return (status, Headers(), b'')
def test_str(self): for exception, exception_str in [ # fmt: off ( InvalidHandshake("invalid request"), "invalid request", ), ( AbortHandshake(200, Headers(), b"OK\n"), "HTTP 200, 0 headers, 3 bytes", ), ( SecurityError("redirect from WSS to WS"), "redirect from WSS to WS", ), ( RedirectHandshake("wss://example.com"), "redirect to wss://example.com", ), ( InvalidMessage("malformed HTTP message"), "malformed HTTP message", ), ( InvalidHeader("Name"), "missing Name header", ), ( InvalidHeader("Name", None), "missing Name header", ), ( InvalidHeader("Name", ""), "empty Name header", ), ( InvalidHeader("Name", "Value"), "invalid Name header: Value", ), ( InvalidHeaderFormat("Sec-WebSocket-Protocol", "expected token", "a=|", 3), "invalid Sec-WebSocket-Protocol header: " "expected token at 3 in a=|", ), ( InvalidHeaderValue("Sec-WebSocket-Version", "42"), "invalid Sec-WebSocket-Version header: 42", ), ( InvalidUpgrade("Upgrade"), "missing Upgrade header", ), ( InvalidUpgrade("Connection", "websocket"), "invalid Connection header: websocket", ), ( InvalidOrigin("http://bad.origin"), "invalid Origin header: http://bad.origin", ), ( InvalidStatusCode(403), "server rejected WebSocket connection: HTTP 403", ), ( NegotiationError("unsupported subprotocol: spam"), "unsupported subprotocol: spam", ), ( InvalidParameterName("|"), "invalid parameter name: |", ), ( InvalidParameterValue("a", "|"), "invalid value for parameter a: |", ), ( DuplicateParameter("a"), "duplicate parameter: a", ), ( InvalidState("WebSocket connection isn't established yet"), "WebSocket connection isn't established yet", ), ( ConnectionClosed(1000, ""), "WebSocket connection is closed: code = 1000 " "(OK), no reason", ), ( ConnectionClosedOK(1001, "bye"), "WebSocket connection is closed: code = 1001 " "(going away), reason = bye", ), (ConnectionClosed(1006, None), "WebSocket connection is closed: code = 1006 " "(connection closed abnormally [internal]), no reason"), (ConnectionClosedError(1016, None), "WebSocket connection is closed: code = 1016 " "(unknown), no reason"), (ConnectionClosed(3000, None), "WebSocket connection is closed: code = 3000 " "(registered), no reason"), (ConnectionClosed(4000, None), "WebSocket connection is closed: code = 4000 " "(private use), no reason"), ( InvalidURI("|"), "| isn't a valid URI", ), ( PayloadTooBig("payload length exceeds limit: 2 > 1 bytes"), "payload length exceeds limit: 2 > 1 bytes", ), ( WebSocketProtocolError("invalid opcode: 7"), "invalid opcode: 7", ), # fmt: on ]: with self.subTest(exception=exception): self.assertEqual(str(exception), exception_str)
def test_str(self): for exception, exception_str in [ # fmt: off ( InvalidHandshake("Invalid request"), "Invalid request", ), ( AbortHandshake(200, Headers(), b'OK\n'), "HTTP 200, 0 headers, 3 bytes", ), ( InvalidMessage("Malformed HTTP message"), "Malformed HTTP message", ), ( InvalidHeader('Name'), "Missing Name header", ), ( InvalidHeader('Name', None), "Missing Name header", ), ( InvalidHeader('Name', ''), "Empty Name header", ), ( InvalidHeader('Name', 'Value'), "Invalid Name header: Value", ), ( InvalidHeaderFormat('Sec-WebSocket-Protocol', "expected token", 'a=|', 3), "Invalid Sec-WebSocket-Protocol header: " "expected token at 3 in a=|", ), ( InvalidHeaderValue('Sec-WebSocket-Version', '42'), "Invalid Sec-WebSocket-Version header: 42", ), ( InvalidUpgrade('Upgrade'), "Missing Upgrade header", ), ( InvalidUpgrade('Connection', 'websocket'), "Invalid Connection header: websocket", ), ( InvalidOrigin('http://bad.origin'), 'Invalid Origin header: http://bad.origin', ), ( InvalidStatusCode(403), "Status code not 101: 403", ), ( NegotiationError("Unsupported subprotocol: spam"), "Unsupported subprotocol: spam", ), ( InvalidParameterName('|'), "Invalid parameter name: |", ), ( InvalidParameterValue('a', '|'), "Invalid value for parameter a: |", ), ( DuplicateParameter('a'), "Duplicate parameter: a", ), ( InvalidState("WebSocket connection isn't established yet"), "WebSocket connection isn't established yet", ), ( ConnectionClosed(1000, ''), "WebSocket connection is closed: code = 1000 " "(OK), no reason", ), ( ConnectionClosed(1001, 'bye'), "WebSocket connection is closed: code = 1001 " "(going away), reason = bye", ), (ConnectionClosed(1006, None), "WebSocket connection is closed: code = 1006 " "(connection closed abnormally [internal]), no reason"), (ConnectionClosed(1016, None), "WebSocket connection is closed: code = 1016 " "(unknown), no reason"), (ConnectionClosed(3000, None), "WebSocket connection is closed: code = 3000 " "(registered), no reason"), (ConnectionClosed(4000, None), "WebSocket connection is closed: code = 4000 " "(private use), no reason"), ( InvalidURI("|"), "| isn't a valid URI", ), ( PayloadTooBig("Payload length exceeds limit: 2 > 1 bytes"), "Payload length exceeds limit: 2 > 1 bytes", ), ( WebSocketProtocolError("Invalid opcode: 7"), "Invalid opcode: 7", ), # fmt: on ]: with self.subTest(exception=exception): self.assertEqual(str(exception), exception_str)