async def app_send(self, message: Optional[dict]) -> None: if self.closed: # Allow app to finish after close return if message is None: # App has errored if self.state == ASGIWebsocketState.HANDSHAKE: await self._send_error_response(500) await self.config.log.access( self.scope, {"status": 500, "headers": []}, time() - self.start_time ) elif self.state == ASGIWebsocketState.CONNECTED: await self._send_wsproto_event(CloseConnection(code=CloseReason.ABNORMAL_CLOSURE)) await self.send(StreamClosed(stream_id=self.stream_id)) else: if message["type"] == "websocket.accept" and self.state == ASGIWebsocketState.HANDSHAKE: self.state = ASGIWebsocketState.CONNECTED status_code, headers, self.connection = self.handshake.accept( message.get("subprotocol") ) await self.send( Response(stream_id=self.stream_id, status_code=status_code, headers=headers) ) await self.config.log.access( self.scope, {"status": status_code, "headers": []}, time() - self.start_time ) elif ( message["type"] == "websocket.http.response.start" and self.state == ASGIWebsocketState.HANDSHAKE ): self.response = message elif message["type"] == "websocket.http.response.body" and self.state in { ASGIWebsocketState.HANDSHAKE, ASGIWebsocketState.RESPONSE, }: await self._send_rejection(message) elif message["type"] == "websocket.send" and self.state == ASGIWebsocketState.CONNECTED: event: WSProtoEvent if message.get("bytes") is not None: event = BytesMessage(data=bytes(message["bytes"])) elif not isinstance(message["text"], str): raise TypeError(f"{message['text']} should be a str") else: event = TextMessage(data=message["text"]) await self._send_wsproto_event(event) elif ( message["type"] == "websocket.close" and self.state == ASGIWebsocketState.HANDSHAKE ): self.state = ASGIWebsocketState.HTTPCLOSED await self._send_error_response(403) elif message["type"] == "websocket.close": self.state = ASGIWebsocketState.CLOSED await self._send_wsproto_event( CloseConnection(code=int(message.get("code", CloseReason.NORMAL_CLOSURE))) ) await self.send(EndData(stream_id=self.stream_id)) else: raise UnexpectedMessage(self.state, message["type"])
def handle_events(self) -> None: for event in self.connection.events(): if isinstance(event, Request): self.task = self.loop.create_task(self.handle_websocket(event)) self.task.add_done_callback(self.maybe_close) elif isinstance(event, Message): try: self.buffer.extend(event) except FrameTooLarge: self.write( self.connection.send( CloseConnection(code=CloseReason.MESSAGE_TOO_BIG))) self.app_queue.put_nowait({"type": "websocket.disconnect"}) self.close() break if event.message_finished: self.app_queue.put_nowait(self.buffer.to_message()) self.buffer.clear() elif isinstance(event, Ping): self.write(self.connection.send(event.response())) elif isinstance(event, CloseConnection): if self.connection.state == ConnectionState.REMOTE_CLOSING: self.write(self.connection.send(event.response())) self.app_queue.put_nowait({"type": "websocket.disconnect"}) self.close() break
def handle_todo_post_save(sender, instance, created, **kwargs): if not hasattr(sender, 'APP_PORT'): return conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM) conn.connect(('localhost', int(sender.APP_PORT))) ws = WSConnection(ConnectionType.CLIENT) print("hello") net_send( ws.send(Request( host=f"localhost:{sender.APP_PORT}", target=f"ws/todos")), conn ) net_recv(ws, conn) handle_events(ws) net_send(ws.send(Message(data=str(instance.pk))), conn) net_recv(ws, conn) handle_events(ws) net_send(ws.send(CloseConnection(code=1000)), conn) net_recv(ws, conn) conn.shutdown(socket.SHUT_WR) net_recv(ws, conn) del sender.APP_PORT
def _handle_close(self, event: CloseConnection) -> None: self._running = False if EventType.DISCONNECT in self._callbacks: self._callbacks[EventType.DISCONNECT](event) # if the server sends first a close connection we need to reply with another one if self._ws.state is ConnectionState.REMOTE_CLOSING: self._sock.sendall(self._ws.send(event.response()))
async def read_messages(self) -> None: while True: data = await self.stream.receive_some(MAX_RECV) if data == b"": data = None # wsproto expects None rather than b"" for EOF self.connection.receive_data(data) for event in self.connection.events(): if isinstance(event, Message): try: self.buffer.extend(event) except FrameTooLarge: await self.asend( CloseConnection(code=CloseReason.MESSAGE_TOO_BIG)) await self.asgi_put({"type": "websocket.disconnect"}) raise MustCloseError() if event.message_finished: await self.asgi_put(self.buffer.to_message()) self.buffer.clear() elif isinstance(event, Ping): await self.asend(event.response()) elif isinstance(event, CloseConnection): if self.connection.state == ConnectionState.REMOTE_CLOSING: await self.asend(event.response()) await self.asgi_put({"type": "websocket.disconnect"}) raise MustCloseError()
def get_case_count(server): uri = urlparse(server + '/getCaseCount') connection = WSConnection(CLIENT) sock = socket.socket() sock.connect((uri.hostname, uri.port or 80)) sock.sendall(connection.send(Request(host=uri.netloc, target=uri.path))) case_count = None while case_count is None: data = sock.recv(65535) connection.receive_data(data) data = "" out_data = b"" for event in connection.events(): if isinstance(event, TextMessage): data += event.data if event.message_finished: case_count = json.loads(data) out_data += connection.send(CloseConnection(code=CloseReason.NORMAL_CLOSURE)) try: sock.sendall(out_data) except CONNECTION_EXCEPTIONS: break sock.close() return case_count
def close(self, reason=None, message=None): out_data = self.ws.send( CloseConnection(reason or CloseReason.NORMAL_CLOSURE, message)) try: self.sock.send(out_data) except BrokenPipeError: pass
def update_reports(server, agent): uri = urlparse(server + '/updateReports?agent=%s' % agent) connection = WSConnection(CLIENT) sock = socket.socket() sock.connect((uri.hostname, uri.port or 80)) sock.sendall( connection.send( Request(host=uri.netloc, target='%s?%s' % (uri.path, uri.query)))) closed = False while not closed: data = sock.recv(65535) connection.receive_data(data) for event in connection.events(): if isinstance(event, AcceptConnection): sock.sendall( connection.send( CloseConnection(code=CloseReason.NORMAL_CLOSURE))) try: sock.close() except CONNECTION_EXCEPTIONS: pass finally: closed = True
def test_closure(client_sends, code, reason): client = Connection(CLIENT) server = Connection(SERVER) if client_sends: local = client remote = server else: local = server remote = client remote.receive_data(local.send(CloseConnection(code=code, reason=reason))) event = next(remote.events()) assert isinstance(event, CloseConnection) assert event.code is code assert event.reason == reason assert remote.state is ConnectionState.REMOTE_CLOSING assert local.state is ConnectionState.LOCAL_CLOSING local.receive_data(remote.send(event.response())) event = next(local.events()) assert isinstance(event, CloseConnection) assert event.code is code assert event.reason == reason assert remote.state is ConnectionState.CLOSED assert local.state is ConnectionState.CLOSED
async def _handle_message_event(self, event: Message): self._message_size += len(event.data) self._message_parts.append(event.data) if self._message_size > self.max_message_size: error = 'Exceeded maximum message size: {} bytes'.format( self.max_message_size) self._message_size = 0 self._message_parts = [] self._close_code = 1009 self._close_reason = error await self._send(CloseConnection(code=1009, reason=error)) self._reader_running = False elif event.message_finished: message = (b'' if isinstance(event, BytesMessage) else '').join( self._message_parts) self._message_size = 0 self._message_parts = [] await self._event_queue.put(message)
def test_closure(client_sends: bool, code: CloseReason, reason: str) -> None: client = Connection(CLIENT) server = Connection(SERVER) if client_sends: local = client remote = server else: local = server remote = client remote.receive_data(local.send(CloseConnection(code=code, reason=reason))) event = next(remote.events()) assert isinstance(event, CloseConnection) assert event.code is code assert event.reason == reason assert remote.state is ConnectionState.REMOTE_CLOSING assert local.state is ConnectionState.LOCAL_CLOSING local.receive_data(remote.send(event.response())) event = next(local.events()) assert isinstance(event, CloseConnection) assert event.code is code assert event.reason == reason assert remote.state is ConnectionState.CLOSED # type: ignore[comparison-overlap] assert local.state is ConnectionState.CLOSED with pytest.raises(LocalProtocolError): local.receive_data(b"foobar")
async def asgi_send(self, message: dict) -> None: """Called by the ASGI instance to send a message.""" if message["type"] == "websocket.accept" and self.state == ASGIWebsocketState.HANDSHAKE: await self.asend(AcceptConnection(extensions=[PerMessageDeflate()])) self.state = ASGIWebsocketState.CONNECTED self.config.access_logger.access( self.scope, {"status": 101, "headers": []}, time() - self.start_time ) elif ( message["type"] == "websocket.http.response.start" and self.state == ASGIWebsocketState.HANDSHAKE ): self.response = message self.config.access_logger.access(self.scope, self.response, time() - self.start_time) elif message["type"] == "websocket.http.response.body" and self.state in { ASGIWebsocketState.HANDSHAKE, ASGIWebsocketState.RESPONSE, }: await self._asgi_send_rejection(message) elif message["type"] == "websocket.send" and self.state == ASGIWebsocketState.CONNECTED: data: Union[bytes, str] if message.get("bytes") is not None: await self.asend(BytesMessage(data=bytes(message["bytes"]))) elif not isinstance(message["text"], str): raise TypeError(f"{message['text']} should be a str") else: await self.asend(TextMessage(data=message["text"])) elif message["type"] == "websocket.close" and self.state == ASGIWebsocketState.HANDSHAKE: await self.send_http_error(403) self.state = ASGIWebsocketState.HTTPCLOSED elif message["type"] == "websocket.close": await self.asend(CloseConnection(code=int(message["code"]))) self.state = ASGIWebsocketState.CLOSED else: raise UnexpectedMessage(self.state, message["type"])
def close_request(self, code: int = 1000, reason: str = None) -> None: if not isinstance(code, int): raise TypeError('code must be an integer') if not isinstance(reason, str): raise TypeError('reason must be a string') self._client.sendall(self._ws.send(CloseConnection(code, reason)))
def websocket(request): # The underlying socket must be provided by the server. Gunicorn and # Werkzeug's dev server are known to support this. stream = request.environ.get("werkzeug.socket") if stream is None: stream = request.environ.get("gunicorn.socket") if stream is None: raise InternalServerError() # Initialize the wsproto connection. Need to recreate the request # data that was read by the WSGI server already. ws = WSConnection(ConnectionType.SERVER) in_data = b"GET %s HTTP/1.1\r\n" % request.path.encode("utf8") for header, value in request.headers.items(): in_data += f"{header}: {value}\r\n".encode() in_data += b"\r\n" ws.receive_data(in_data) running = True while True: out_data = b"" for event in ws.events(): if isinstance(event, WSRequest): out_data += ws.send(AcceptConnection()) elif isinstance(event, CloseConnection): out_data += ws.send(event.response()) running = False elif isinstance(event, Ping): out_data += ws.send(event.response()) elif isinstance(event, TextMessage): # echo the incoming message back to the client if event.data == "quit": out_data += ws.send( CloseConnection(CloseReason.NORMAL_CLOSURE, "bye") ) running = False else: out_data += ws.send(Message(data=event.data)) if out_data: stream.send(out_data) if not running: break in_data = stream.recv(4096) ws.receive_data(in_data) # The connection will be closed at this point, but WSGI still # requires a response. return Response("", status=204)
async def _abort_websocket(self): close_code = CloseReason.ABNORMAL_CLOSURE if self.state is ConnectionState.OPEN: self._wsproto.send(CloseConnection(code=close_code.value)) if self._close_code is None: await self._close_websocket(close_code) await self._close_stream() await self._close_handshake.set()
def _handle_close_connection(self, event, source_conn, other_conn, is_server): self.flow.close_sender = "server" if is_server else "client" self.flow.close_code = event.code self.flow.close_reason = event.reason data = self.connections[other_conn].send(CloseConnection(code=event.code, reason=event.reason)) other_conn.send(data) data = self.connections[source_conn].send(event.response()) source_conn.send(data) return False
async def _next_event(self): """ Gets the next event. """ while True: for event in self._connection.events(): if isinstance(event, Message): # check if we need to buffer if event.message_finished: return self._wrap_data(self._gather_buffers(event)) self._buffer(event) break # exit for loop return event if self._sock is None: return CloseConnection(code=500, reason="Socket closed") data = await self._sock.receive(4096) if not data: return CloseConnection(code=500, reason="Socket closed") self._connection.receive_data(data)
async def test_asgi_send() -> None: server = MockWebsocket() server.app = bad_framework await server.asgi_send({"type": "websocket.accept"}) await server.asgi_send({"type": "websocket.send", "bytes": b"abc"}) await server.asgi_send({"type": "websocket.close", "code": 1000}) assert isinstance(server.sent_events[0], AcceptConnection) assert server.sent_events[1:] == [ BytesMessage(data=b"abc"), CloseConnection(code=1000) ]
async def aclose(self) -> None: try: try: await self.stream.send_all( self.ws.send( CloseConnection(code=CloseReason.NORMAL_CLOSURE))) except LocalProtocolError: # TODO: exception occurs when ws.state is already closed... pass await self.stream.aclose() except (BrokenResourceError, TransportError): pass
def wsproto_demo(host, port): ''' Demonstrate wsproto: 0) Open TCP connection 1) Negotiate WebSocket opening handshake 2) Send a message and display response 3) Send ping and display pong 4) Negotiate WebSocket closing handshake :param stream: a socket stream ''' # 0) Open TCP connection print('Connecting to {}:{}'.format(host, port)) conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM) conn.connect((host, port)) # 1) Negotiate WebSocket opening handshake print('Opening WebSocket') ws = WSConnection(ConnectionType.CLIENT) # Because this is a client WebSocket, we need to initiate the connection # handshake by sending a Request event. net_send(ws.send(Request(host=host, target='server')), conn) net_recv(ws, conn) handle_events(ws) # 2) Send a message and display response message = "wsproto is great" print('Sending message: {}'.format(message)) net_send(ws.send(Message(data=message)), conn) net_recv(ws, conn) handle_events(ws) # 3) Send ping and display pong payload = b"table tennis" print('Sending ping: {}'.format(payload)) net_send(ws.send(Ping(payload=payload)), conn) net_recv(ws, conn) handle_events(ws) # 4) Negotiate WebSocket closing handshake print('Closing WebSocket') net_send(ws.send(CloseConnection(code=1000, reason='sample reason')), conn) # After sending the closing frame, we won't get any more events. The server # should send a reply and then close the connection, so we need to receive # twice: net_recv(ws, conn) conn.shutdown(socket.SHUT_WR) net_recv(ws, conn)
async def test_bad_framework() -> None: server = MockWebsocket() server.app = bad_framework headers = [ (b"sec-websocket-key", b"ZdCqRHQRNflNt6o7yU48Pg=="), (b"sec-websocket-version", b"13"), (b"connection", b"upgrade"), (b"upgrade", b"connection"), ] request = Request(target="/accept", host="hypercorn", extra_headers=headers) await server.handle_websocket(request) assert isinstance(server.sent_events[0], AcceptConnection) assert server.sent_events[1:] == [CloseConnection(code=1006)]
async def close(self, code: int = 1006, reason: str = "Connection closed"): """ Closes the websocket. """ if self._closed: return self._closed = True if self._scope is not None: await self._scope.cancel() # cancel any outstanding listeners data = self._connection.send(CloseConnection(code=code, reason=reason)) await self._sock.send_all(data) # No, we don't wait for the correct reply await self._sock.close()
async def _handle_events(self) -> None: for event in self.connection.events(): if isinstance(event, Message): try: self.buffer.extend(event) except FrameTooLarge: await self._send_wsproto_event( CloseConnection(code=CloseReason.MESSAGE_TOO_BIG)) break if event.message_finished: await self.app_put(self.buffer.to_message()) self.buffer.clear() elif isinstance(event, Ping): await self._send_wsproto_event(event.response()) elif isinstance(event, CloseConnection): if self.connection.state == ConnectionState.REMOTE_CLOSING: await self._send_wsproto_event(event.response()) await self.send(StreamClosed(stream_id=self.stream_id))
async def close(self, code: int = 1006, reason: str = "Connection closed"): """ Closes the websocket. """ sock, self._sock = self._sock, None if not sock: return try: data = self._connection.send( CloseConnection(code=code, reason=reason)) except LocalProtocolError: # not yet fully open pass else: await sock.send(data) # No, we don't wait for the correct reply await sock.aclose()
async def test_server_tcp_closed_on_close_connection_event(nursery): """ensure server closes TCP immediately after receiving CloseConnection""" server_stream_closed = trio.Event() async def _close_stream_stub(): assert not server_stream_closed.is_set() server_stream_closed.set() async def handle_connection(request): ws = await request.accept() ws._close_stream = _close_stream_stub await trio.sleep_forever() server = await nursery.start( partial(serve_websocket, handle_connection, HOST, 0, ssl_context=None)) client = await connect_websocket(nursery, HOST, server.port, RESOURCE, use_ssl=False) # send a CloseConnection event to server but leave client connected await client._send(CloseConnection(code=1000)) await server_stream_closed.wait()
async def handle_asgi_app(self, event: Request) -> None: self.start_time = time() await self.asgi_put({"type": "websocket.connect"}) try: asgi_instance = self.app(self.scope) await asgi_instance(self.asgi_receive, self.asgi_send) except asyncio.CancelledError: pass except Exception: if self.config.error_logger is not None: self.config.error_logger.exception("Error in ASGI Framework") if self.state == ASGIWebsocketState.CONNECTED: await self.asend(CloseConnection(code=CloseReason.ABNORMAL_CLOSURE)) self.state = ASGIWebsocketState.CLOSED # If the application hasn't accepted the connection (or sent a # response) send a 500 for it. Otherwise if the connection # hasn't been closed then close it. if self.state == ASGIWebsocketState.HANDSHAKE: await self.send_http_error(500) self.state = ASGIWebsocketState.HTTPCLOSED
async def close(self, code: int = 1000, reason: str = ''): """Terminates the open WebSocket connection. This sends a closing frame and suspends until the connection is closed. After calling this method, any further I/O on this WebSocket, e.g. :meth:`~WebSocketConnection.get_message` will raise :exc:`~anysocks.exceptions.ConnectionClosed`. Parameters ---------- code : int A 4-digit WebSocket close code. reason : str An optional string indicating why the connection was closed. """ if self._closed: return logger.info('Closing connection to %s...', self.host) self._closed = True await self._close_websocket(code, reason) try: if self.state is ConnectionState.OPEN: await self._send(CloseConnection(code=code, reason=reason)) elif self.state in (ConnectionState.CONNECTING, ConnectionState.REJECTING): await self._close_handshake.set() await self._close_handshake.wait() except ConnectionClosed: # If _send() raised ConnectionClosed, we can opt out. pass finally: await self._close_stream()
def close(self) -> None: self.server.data_received( self.connection.send(CloseConnection(code=1000)))
async def close(self, code=1000, reason='Closed.'): await self.outgoing.put(CloseConnection(code=code, reason=reason))
def test_close_whilst_closing() -> None: client = Connection(CLIENT) client.send(CloseConnection(code=CloseReason.NORMAL_CLOSURE)) with pytest.raises(LocalProtocolError): client.send(CloseConnection(code=CloseReason.NORMAL_CLOSURE))