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_upgrade_request() -> None: server = WSConnection(SERVER) server.initiate_upgrade_connection( [ (b"Host", b"localhost"), (b"Connection", b"Keep-Alive, Upgrade"), (b"Upgrade", b"WebSocket"), (b"Sec-WebSocket-Version", b"13"), (b"Sec-WebSocket-Key", generate_nonce()), (b"X-Foo", b"bar"), ], "/", ) event = next(server.events()) event = cast(Request, event) assert event.extensions == [] assert event.host == "localhost" assert event.subprotocols == [] assert event.target == "/" headers = normed_header_dict(event.extra_headers) assert b"host" not in headers assert b"sec-websocket-extensions" not in headers assert b"sec-websocket-protocol" not in headers assert headers[b"connection"] == b"Keep-Alive, Upgrade" assert headers[b"sec-websocket-version"] == b"13" assert headers[b"upgrade"] == b"WebSocket" assert headers[b"x-foo"] == b"bar"
def _make_handshake( response_status: int, response_headers: Headers, subprotocols: Optional[List[str]] = None, extensions: Optional[List[Extension]] = None, auto_accept_key: bool = True, ) -> List[Event]: client = WSConnection(CLIENT) assert client.state is ConnectionState.CONNECTING server = h11.Connection(h11.SERVER) server.receive_data( client.send( Request( host="localhost", target="/", subprotocols=subprotocols or [], extensions=extensions or [], ))) request = server.next_event() if auto_accept_key: full_request_headers = normed_header_dict(request.headers) response_headers.append(( b"Sec-WebSocket-Accept", generate_accept_token(full_request_headers[b"sec-websocket-key"]), )) response = h11.InformationalResponse(status_code=response_status, headers=response_headers) client.receive_data(server.send(response)) assert client.state is not ConnectionState.CONNECTING return list(client.events())
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
class MockWebsocketConnection: def __init__(self, path: str, *, framework: ASGIFramework = echo_framework) -> None: self.client_stream, server_stream = trio.testing.memory_stream_pair() server_stream.socket = MockSocket() self.server = WebsocketServer(framework, Config(), server_stream) self.connection = WSConnection(ConnectionType.CLIENT) self.server.connection.receive_data( self.connection.send(Request(target=path, host="hypercorn"))) async def send(self, data: AnyStr) -> None: await self.client_stream.send_all( self.connection.send(Message(data=data))) await trio.sleep(0) # Allow the server to respond async def receive(self) -> List[Message]: data = await self.client_stream.receive_some(2**16) self.connection.receive_data(data) return [event for event in self.connection.events()] async def close(self) -> None: await self.client_stream.send_all( self.connection.send(CloseConnection(code=1000)))
def _make_connection_request(request_headers, method="GET"): # type: (List[Tuple[str, str]]) -> Request client = h11.Connection(h11.CLIENT) server = WSConnection(SERVER) server.receive_data( client.send(h11.Request(method=method, target="/", headers=request_headers)) ) return next(server.events())
def __init__(self, stream: Stream, ws: WSConnection, keepalive: Optional[int] = None): self.stream = stream self.ws = ws self.keepalive = keepalive self.conn_id = uuid4().hex self.logger = logger.bind(conn_id=self.conn_id) self._ws_events = ws.events() self._handshake: Optional[ServerHandshake] = None
def handle_events(ws: WSConnection) -> None: for event in ws.events(): if isinstance(event, AcceptConnection): print("WebSocket negotiation complete") elif isinstance(event, TextMessage): print("Received message: {}".format(event.data)) else: raise Exception("Do not know how to handle event: " + str(event))
def _make_connection_request(request_headers: Headers, method: str = "GET") -> Request: client = h11.Connection(h11.CLIENT) server = WSConnection(SERVER) server.receive_data( client.send( h11.Request(method=method, target="/", headers=request_headers))) return next(server.events()) # type: ignore
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)
def handle_connection(stream: socket.socket) -> None: """ Handle a connection. The server operates a request/response cycle, so it performs a synchronous loop: 1) Read data from network into wsproto 2) Get new events and handle them 3) Send data from wsproto to network :param stream: a socket stream """ ws = WSConnection(ConnectionType.SERVER) running = True while running: # 1) Read data from network in_data = stream.recv(RECEIVE_BYTES) print("Received {} bytes".format(len(in_data))) ws.receive_data(in_data) # 2) Get new events and handle them out_data = b"" for event in ws.events(): if isinstance(event, Request): # Negotiate new WebSocket connection print("Accepting WebSocket upgrade") out_data += ws.send(AcceptConnection()) elif isinstance(event, CloseConnection): # Print log message and break out print("Connection closed: code={} reason={}".format( event.code, event.reason)) out_data += ws.send(event.response()) running = False elif isinstance(event, TextMessage): # Reverse text and send it back to wsproto print("Received request and sending response") out_data += ws.send(Message(data=event.data[::-1])) elif isinstance(event, Ping): # wsproto handles ping events for you by placing a pong frame in # the outgoing buffer. You should not call pong() unless you want to # send an unsolicited pong frame. print("Received ping and sending pong") out_data += ws.send(event.response()) else: print(f"Unknown event: {event!r}") # 4) Send data from wsproto to network print("Sending {} bytes".format(len(out_data))) stream.send(out_data)
def _make_handshake_rejection(status_code, body=None): client = WSConnection(CLIENT) server = h11.Connection(h11.SERVER) server.receive_data(client.send(Request(host="localhost", target="/"))) headers = [] if body is not None: headers.append(("Content-Length", str(len(body)))) client.receive_data( server.send(h11.Response(status_code=status_code, headers=headers)) ) if body is not None: client.receive_data(server.send(h11.Data(data=body))) client.receive_data(server.send(h11.EndOfMessage())) return list(client.events())
class Websocket(WebsocketPrototype): def __init__(self): super().__init__() self.socket = curio.socket.socket(curio.socket.AF_INET, curio.socket.SOCK_STREAM) self.protocol = WSConnection(ConnectionType.CLIENT) async def connect(self, path, host, port): await self.socket.connect((host, port)) request = Request(host=f'{host}:{port}', target=path) await self.socket.sendall(self.protocol.send(request)) upgrade_response = await self.socket.recv(8096) self.protocol.receive_data(upgrade_response) event = next(self.protocol.events()) if not isinstance(event, AcceptConnection): raise Exception('Websocket handshake failed.')
def new_conn(sock: socket.socket) -> None: global count print("test_server.py received connection {}".format(count)) count += 1 ws = WSConnection(SERVER) closed = False while not closed: try: data: Optional[bytes] = sock.recv(65535) except socket.error: data = None ws.receive_data(data or None) outgoing_data = b"" for event in ws.events(): if isinstance(event, Request): outgoing_data += ws.send( AcceptConnection(extensions=[PerMessageDeflate()])) elif isinstance(event, Message): outgoing_data += ws.send( Message(data=event.data, message_finished=event.message_finished)) elif isinstance(event, Ping): outgoing_data += ws.send(event.response()) elif isinstance(event, CloseConnection): closed = True if ws.state is not ConnectionState.CLOSED: outgoing_data += ws.send(event.response()) if not data: closed = True try: sock.sendall(outgoing_data) except socket.error: closed = True sock.close()
async def __call__(self, scope: Scope, receive: Receive, send: Send): if scope["type"] != "http": return await self.app(scope, receive, send) # A WebSocket-Request must be a GET request. if scope["method"].lower() != "get": return await self.app(scope, receive, send) # We did not detect an upgrade header. upgrade = tuple(get_header_from_scope(b"upgrade", scope)) if not upgrade or upgrade[0].lower() != b"websocket": return await self.app(scope, receive, send) request = scope_to_request(scope) websocket = WSConnection(SERVER) connection = h11.Connection(h11.CLIENT) # Receive the actual connection. websocket.receive_data(connection.send(request)) event: Request = next(websocket.events()) assert isinstance(event, Request) headers = event.extra_headers host = event.host headers.append([b"host", host.encode("utf-8")]) subprotocols = event.subprotocols extensions = scope.get("extensions", {}).copy() extensions["websocket.http.response"] = {} ws_scope = { "type": "websocket", ############################################################################### "path": scope["path"], "headers": headers, "subprotocols": subprotocols, ############################################################################### "asgi": scope.get("asgi", { "version": "3.0", "spec_version": "2.1" }), "http_version": scope.get("http_version", "1.1"), "scheme": "ws" + (scope.get("scheme", "http").endswith("s") * "s"), "query_string": scope.get("query_string", b""), "root_path": scope.get("root_path", b""), "client": scope.get("client", None), "server": scope.get("server", None), "extensions": extensions } async def finish_events(port: QueuePort): buffer = [] for event in websocket.events(): if isinstance(event, CloseConnection): await port.close({ "type": "websocket.disconnect", "code": event.code }) return elif isinstance(event, Message): buffer += event.data if event.message_finished: base = '' if isinstance(event, TextMessage) else b'' key = 'text' if isinstance(event, TextMessage) else 'bytes' data = base.join(buffer) buffer = [] await port.put({ "type": "websocket.receive", key: data }) async def receiver(receive: Receive, port: QueuePort): while websocket.state not in (ConnectionState.CLOSED, ConnectionState.REMOTE_CLOSING): event = await receive() websocket.receive_data(event.get("body", b"")) await finish_events(port) if not event.get("more_body", False): websocket.receive_data(None) await finish_events(port) return def instrument_errors(func): async def _wrapper(*args, **kwargs): try: return await func(*args, **kwargs) except BaseException as e: import traceback traceback.print_exception(type(e), e, e.__traceback__) return _wrapper async def sender(send: Send, port: QueuePort): state = 'connecting' while websocket.state not in (ConnectionState.CLOSED, ConnectionState.LOCAL_CLOSING): event = await port.pull() if event is QueuePort.PORT_CLOSED_SENTINEL: return assert isinstance(event, dict) if event["type"] == "websocket.send": if "bytes" in event: await send({ 'type': 'http.response.body', 'body': websocket.send(BytesMessage(data=event["bytes"])), 'more_body': True }) elif "text" in event: await send({ 'type': 'http.response.body', 'body': websocket.send(TextMessage(data=event["text"])), 'more_body': True }) elif event["type"] == "websocket.close": if state == "connected": code = event.get("code", 1000) await send({ 'type': 'http.response.body', 'body': websocket.send(CloseConnection(code=code)), 'more_body': False }) await port.close({ 'type': 'websocket.close', 'code': code }) else: state = "denied" await send({ 'type': 'http.response.start', 'status': 403, 'headers': [] }) await send({ 'type': 'http.response.body', 'body': b'', 'more_body': False }) await port.close({'type': 'websocket.close'}) elif event["type"] in ("websocket.http.response.start", "websocket.http.response.body"): if state == "connected": raise ValueError( "You already accepted a websocket connection.") if state == "connecting" and event[ "type"] == "websocket.http.response.body": raise ValueError("You did not start a response.") elif state == "denied" and event[ "type"] == "websocket.http.response.start": raise ValueError("You already started a response.") state = "denied" event = event.copy() event["type"] = event["type"][len("websocket."):] await send(event) elif event["type"] == "websocket.accept": raw = websocket.send( AcceptConnection( extra_headers=event.get("headers", []), subprotocol=event.get("subprotocol", None))) connection.receive_data(raw) response = connection.next_event() state = "connected" await send({ 'type': 'http.response.start', 'status': response.status_code, 'headers': response.headers }) port = QueuePort() await port.put({'type': 'websocket.connect'}) loop = get_running_loop() receive_task = loop.create_task( instrument_errors(receiver)(receive, port)) send_task = loop.create_task(instrument_errors(sender)(send, port)) try: await self.app(ws_scope, port.receive, port.send) finally: await send_task await receive_task
class BaseServer(ABC): bytes_to_receive: int = 65535 buffer_size: int = io.DEFAULT_BUFFER_SIZE # noinspection PyTypeChecker def __init__(self, host: str, port: int): self._check_init_arguments(host, port) self._host = host self._port = port self._server: StreamServer = None self._ws = WSConnection(ConnectionType.SERVER) self._client: socket = None # client socket provided by the StreamServer self._running = True @staticmethod def _check_ws_headers(headers: Headers) -> None: if headers is None: return error_message = 'headers must of a list of tuples of the form [(bytes, bytes), ..]' if not isinstance(headers, list): raise TypeError(error_message) try: for key, value in headers: if not isinstance(key, bytes) or not isinstance(value, bytes): raise TypeError(error_message) except ValueError: # in case it is not a list of tuples raise TypeError(error_message) @abstractmethod def handle_request(self, request: Request) -> None: pass def accept_request(self, extra_headers: Headers = None, sub_protocol: str = None) -> None: self._check_ws_headers(extra_headers) if sub_protocol is not None and not isinstance(sub_protocol, str): raise TypeError('sub_protocol must be a string') extra_headers = extra_headers if extra_headers else [] self._client.sendall( self._ws.send( AcceptConnection(extra_headers=extra_headers, subprotocol=sub_protocol))) def reject_request(self, status_code: int = 400, reason: str = None) -> None: if not isinstance(status_code, int): raise TypeError('status_code must be an integer') if reason is not None and not isinstance(reason, str): raise TypeError('reason must be a string') if not reason: self._client.sendall( self._ws.send(RejectConnection(status_code=status_code))) else: data = bytearray( self._ws.send( RejectConnection(has_body=True, headers=[(b'Content-type', b'text/txt') ]))) data.extend(self._ws.send(RejectData(reason.encode()))) self._client.sendall(bytes(data)) 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 _handle_close_event(self, event: CloseConnection) -> None: if self._ws.state is ConnectionState.REMOTE_CLOSING: self._client.sendall(self._ws.send(event.response())) def _handle_ping(self, event: Ping) -> None: self._client.sendall(self._ws.send(event.response())) @abstractmethod def receive_text(self, data: str) -> None: pass @abstractmethod def receive_json(self, data: Any) -> None: pass @abstractmethod def receive_bytes(self, data: bytes) -> None: pass @abstractmethod def handle_pong(self, data: bytes) -> None: pass def _send_data(self, data: AnyStr) -> None: if isinstance(data, str): io_object = io.StringIO(data) else: io_object = io.BytesIO(data) with io_object as f: chunk = f.read(self.buffer_size) while chunk: if len(chunk) < self.buffer_size: self._client.sendall( self._ws.send(Message(data, message_finished=True))) break else: self._client.sendall( self._ws.send(Message(data, message_finished=False))) chunk = f.read(self.buffer_size) def ping(self, data: bytes = b'hello') -> None: if not isinstance(data, bytes): raise TypeError('data must be bytes') self._client.sendall(self._ws.send(Ping(data))) def send(self, data: AnyStr) -> None: if not isinstance(data, (bytes, str)): raise TypeError('data must be either a string or binary data') self._send_data(data) def send_json(self, data: Any) -> None: self.send(json.dumps(data)) @staticmethod def _check_init_arguments(host: str, port: int) -> None: if not isinstance(host, str): raise TypeError('host must be a string') error_message = 'custom_port must a positive integer' if not isinstance(port, int): raise TypeError(error_message) if port < 0: raise TypeError(error_message) def _handler(self, client: socket, address: Tuple[str, int]) -> None: self._client = client text_message = [] binary_message = bytearray() while self._running: data = client.recv(self.bytes_to_receive) self._ws.receive_data(data) for event in self._ws.events(): if isinstance(event, Request): self.handle_request(event) elif isinstance(event, CloseConnection): self._handle_close_event(event) self._running = False elif isinstance(event, Ping): self._handle_ping(event) elif isinstance(event, Pong): self.handle_pong(event.payload) elif isinstance(event, TextMessage): text_message.append(event.data) if event.message_finished: str_data = ''.join(text_message) try: self.receive_json(json.loads(str_data)) except json.JSONDecodeError: self.receive_text(str_data) text_message.clear() elif isinstance(event, BytesMessage): binary_message.extend(event.data) if event.message_finished: self.receive_bytes(bytes(binary_message)) binary_message.clear() else: print('unknown event:', event) def run(self, backlog: int = 256, spawn: str = 'default', **kwargs) -> None: self._server = StreamServer((self._host, self._port), self._handler, backlog=backlog, spawn=spawn, **kwargs) self._server.serve_forever() def close(self) -> None: if self._server is not None: self._server.close()
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) net_send(ws.send(Request(host=host, target='server')), conn) net_recv(ws, conn) # events is a generator that yields websocket event objects. Usually you # would say `for event in ws.events()`, but the synchronous nature of this # client requires us to use next(event) instead so that we can interleave # the network I/O. It will raise StopIteration when it runs out of events # (i.e. needs more network data), but since this script is synchronous, we # will explicitly resume the generator whenever we have new network data. events = ws.events() # Because this is a client WebSocket, wsproto has automatically queued up # a handshake, and we need to send it and wait for a response. event = next(events) if isinstance(event, AcceptConnection): print('WebSocket negotiation complete') else: raise Exception('Expected AcceptConnection event!') # 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) event = next(events) if isinstance(event, TextMessage): print('Received message: {}'.format(event.data)) else: raise Exception('Expected TextMessage event!') # 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) event = next(events) if isinstance(event, Pong): print('Received pong: {}'.format(event.payload)) else: raise Exception('Expected Pong event!') # 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)
class Base: def __init__(self, sock=None, connection_type=None, receive_bytes=4096, thread_class=threading.Thread, event_class=threading.Event): self.sock = sock self.receive_bytes = receive_bytes self.input_buffer = [] self.event = event_class() self.connected = False self.is_server = (connection_type == ConnectionType.SERVER) self.ws = WSConnection(connection_type) self.handshake() self.thread = thread_class(target=self._thread) self.thread.start() self.event.wait() self.event.clear() def handshake(self): """To be implemented by subclasses.""" pass def send(self, data): if not self.connected: raise ConnectionClosed() if isinstance(data, bytes): out_data = self.ws.send(Message(data=data)) else: out_data = self.ws.send(TextMessage(data=str(data))) self.sock.send(out_data) def receive(self, timeout=None): while self.connected and not self.input_buffer: if not self.event.wait(timeout=timeout): return None self.event.clear() if not self.connected: raise ConnectionClosed() return self.input_buffer.pop(0) 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 _thread(self): self.connected = self._handle_events() self.event.set() while self.connected: try: in_data = self.sock.recv(self.receive_bytes) except (OSError, ConnectionResetError): self.connected = False self.event.set() break self.ws.receive_data(in_data) self.connected = self._handle_events() def _handle_events(self): keep_going = True out_data = b'' try: for event in self.ws.events(): if isinstance(event, Request): out_data += self.ws.send(AcceptConnection()) elif isinstance(event, CloseConnection): if self.is_server: out_data += self.ws.send(event.response()) self.event.set() keep_going = False elif isinstance(event, Ping): out_data += self.ws.send(event.response()) elif isinstance(event, TextMessage): self.input_buffer.append(event.data) self.event.set() elif isinstance(event, BytesMessage): self.input_buffer.append(event.data) self.event.set() if out_data: self.sock.send(out_data) return keep_going except: return False
class Client: _callbacks: Dict[EventType, Callable] = {} receive_bytes: int = 65535 buffer_size: int = io.DEFAULT_BUFFER_SIZE # noinspection PyTypeChecker def __init__(self, connect_uri: str, headers: Headers = None, extensions: List[str] = None, sub_protocols: List[str] = None): self._check_ws_headers(headers) self._check_list_argument('extensions', extensions) self._check_list_argument('sub_protocols', sub_protocols) self._sock: socket = None self._ws: WSConnection = None # wsproto does not seem to like empty path, so we provide an arbitrary one self._default_path = '/path' self._running = True self._handshake_finished = AsyncResult() host, port, path = self._get_connect_information(connect_uri) self._establish_tcp_connection(host, port) self._establish_websocket_handshake(host, path, headers, extensions, sub_protocols) self._green = spawn(self._run) @staticmethod def _check_ws_headers(headers: Headers) -> None: if headers is None: return error_message = 'headers must of a list of tuples of the form [(bytes, bytes), ..]' if not isinstance(headers, list): raise TypeError(error_message) try: for key, value in headers: if not isinstance(key, bytes) or not isinstance(value, bytes): raise TypeError(error_message) except ValueError: # in case it is not a list of tuples raise TypeError(error_message) @staticmethod def _check_list_argument(name: str, ws_argument: List[str]) -> None: if ws_argument is None: return error_message = f'{name} must be a list of strings' if not isinstance(ws_argument, list): raise TypeError(error_message) for item in ws_argument: if not isinstance(item, str): raise TypeError(error_message) def _get_connect_information(self, connect_uri: str) -> Tuple[str, int, str]: if not isinstance(connect_uri, str): raise TypeError('Your uri must be a string') regex = re.match(r'ws://(\w+)(:\d+)?(/\w+)?', connect_uri) if not regex: raise ValueError( 'Your uri must follow the syntax ws://<host>[:port][/path]') host = regex.group(1) port = int(regex.group(2)[1:]) if regex.group(2) is not None else 80 path = regex.group(3)[1:] if regex.group( 3) is not None else self._default_path return host, port, path @staticmethod def _check_callable(method: str, callback: Callable) -> None: if not isinstance(callback, callable): raise TypeError(f'{method} callback must be a callable') @classmethod def _on_callback(cls, event_type: EventType, func: Callback) -> Callback: cls._callbacks[event_type] = func return func @classmethod def on_connect(cls, func: EventCallback) -> EventCallback: return cls._on_callback(EventType.CONNECT, func) @classmethod def on_disconnect(cls, func: EventCallback) -> EventCallback: return cls._on_callback(EventType.DISCONNECT, func) @classmethod def on_ping(cls, func: BytesCallback) -> BytesCallback: return cls._on_callback(EventType.PING, func) @classmethod def on_pong(cls, func: BytesCallback) -> BytesCallback: return cls._on_callback(EventType.PONG, func) @classmethod def on_text_message(cls, func: StrCallback) -> StrCallback: return cls._on_callback(EventType.TEXT_MESSAGE, func) @classmethod def on_json_message(cls, func: JsonCallback) -> JsonCallback: return cls._on_callback(EventType.JSON_MESSAGE, func) @classmethod def on_binary_message(cls, func: BytesCallback) -> BytesCallback: return cls._on_callback(EventType.BINARY_MESSAGE, func) def _establish_tcp_connection(self, host: str, port: int) -> None: self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self._sock.connect((host, port)) def _establish_websocket_handshake(self, host: str, path: str, headers: Headers, extensions: List[str], sub_protocols: List[str]) -> None: self._ws = WSConnection(ConnectionType.CLIENT) headers = headers if headers is not None else [] extensions = extensions if extensions is not None else [] sub_protocols = sub_protocols if sub_protocols is not None else [] request = Request(host=host, target=path, extra_headers=headers, extensions=extensions, subprotocols=sub_protocols) self._sock.sendall(self._ws.send(request)) def _handle_accept(self, event: AcceptConnection) -> None: self._handshake_finished.set() if EventType.CONNECT in self._callbacks: self._callbacks[EventType.CONNECT](self, event) def _handle_reject(self, event: RejectConnection) -> None: self._handshake_finished.set_exception( ConnectionRejectedError(event.status_code, event.headers, b'')) self._running = False def _handle_reject_data(self, event: RejectData, data: bytearray, status_code: int, headers: Headers) -> None: data.extend(event.data) if event.body_finished: self._handshake_finished.set_exception( ConnectionRejectedError(status_code, headers, data)) self._running = False 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())) def _handle_ping(self, event: Ping) -> None: if EventType.PING in self._callbacks: self._callbacks[EventType.PING](event.payload) self._sock.sendall(self._ws.send(event.response())) def _handle_pong(self, event: Pong) -> None: if EventType.PONG in self._callbacks: self._callbacks[EventType.PONG](event.payload) def _handle_text_or_json_message(self, event: TextMessage, text_message: List[str]) -> None: text_message.append(event.data) if event.message_finished: if EventType.JSON_MESSAGE in self._callbacks: str_message = ''.join(text_message) try: self._callbacks[EventType.JSON_MESSAGE]( self, json.loads(str_message)) text_message.clear() return # no need to process text handler if json handler already does the job except json.JSONDecodeError: pass if EventType.TEXT_MESSAGE in self._callbacks: self._callbacks[EventType.TEXT_MESSAGE](self, ''.join(text_message)) text_message.clear() def _handle_binary_message(self, event: BytesMessage, binary_message: bytearray) -> None: binary_message.extend(event.data) if event.message_finished: if EventType.BINARY_MESSAGE in self._callbacks: self._callbacks[EventType.BINARY_MESSAGE](self, binary_message) binary_message.clear() def _run(self) -> None: reject_data = bytearray() reject_status_code = 400 reject_headers = [] text_message = [] binary_message = bytearray() while self._running: data = self._sock.recv(self.receive_bytes) if not data: data = None self._ws.receive_data(data) for event in self._ws.events(): if isinstance(event, AcceptConnection): self._handle_accept(event) elif isinstance(event, RejectConnection): if not event.has_body: self._handle_reject(event) else: reject_status_code = event.status_code reject_headers = event.headers elif isinstance(event, RejectData): self._handle_reject_data(event, reject_data, reject_status_code, reject_headers) elif isinstance(event, CloseConnection): self._handle_close(event) elif isinstance(event, Ping): self._handle_ping(event) elif isinstance(event, Pong): self._handle_pong(event) elif isinstance(event, TextMessage): self._handle_text_or_json_message(event, text_message) elif isinstance(event, BytesMessage): self._handle_binary_message(event, binary_message) else: print('unknown event', event) self._sock.close() def ping(self, data: bytes = b'hello') -> None: self._handshake_finished.get() if not isinstance(data, bytes): raise TypeError('data must be bytes') self._sock.sendall(self._ws.send(Ping(data))) def _send_data(self, data: AnyStr) -> None: if isinstance(data, str): io_object = io.StringIO(data) else: io_object = io.BytesIO(data) with io_object as f: chunk = f.read(self.buffer_size) while chunk: if len(chunk) < self.buffer_size: self._sock.sendall( self._ws.send(Message(data, message_finished=True))) break else: self._sock.sendall( self._ws.send(Message(data, message_finished=False))) chunk = f.read(self.buffer_size) def send(self, data: AnyStr) -> None: self._handshake_finished.get() if not isinstance(data, (bytes, str)): raise TypeError('data must be bytes or string') self._send_data(data) def send_json(self, data: Any) -> None: self.send(json.dumps(data)) def _close_ws_connection(self): close_data = self._ws.send( CloseConnection(code=1000, reason='nothing more to do')) self._sock.sendall(close_data) def close(self) -> None: self._handshake_finished.get() if self._ws.state is ConnectionState.OPEN: self._close_ws_connection() # don't forget to join the run greenlet, if not, you will have some surprises with your event handlers! self._green.join() def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close()
class WebsocketServer(HTTPServer, WebsocketMixin): def __init__( self, app: ASGIFramework, config: Config, stream: trio.abc.Stream, *, upgrade_request: Optional[h11.Request] = None, ) -> None: super().__init__(stream, "wsproto") self.app = app self.config = config self.connection = WSConnection(ConnectionType.SERVER) self.response: Optional[dict] = None self.scope: Optional[dict] = None self.send_lock = trio.Lock() self.state = ASGIWebsocketState.HANDSHAKE self.buffer = WebsocketBuffer(self.config.websocket_max_message_size) self.app_send_channel, self.app_receive_channel = trio.open_memory_channel( 10) if upgrade_request is not None: self.connection.initiate_upgrade_connection( upgrade_request.headers, upgrade_request.target) async def handle_connection(self) -> None: try: request = await self.read_request() async with trio.open_nursery() as nursery: nursery.start_soon(self.read_messages) await self.handle_websocket(request) if self.state == ASGIWebsocketState.HTTPCLOSED: raise MustCloseError() except (trio.BrokenResourceError, trio.ClosedResourceError): await self.asgi_put({"type": "websocket.disconnect"}) except MustCloseError: pass finally: await self.aclose() async def read_request(self) -> Request: for event in self.connection.events(): if isinstance(event, Request): return event 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() async def asend(self, event: Event) -> None: async with self.send_lock: await self.stream.send_all(self.connection.send(event)) async def asgi_put(self, message: dict) -> None: await self.app_send_channel.send(message) async def asgi_receive(self) -> dict: return await self.app_receive_channel.receive() @property def scheme(self) -> str: return "wss" if self._is_ssl else "ws"
class WebsocketServer(HTTPServer, WebsocketMixin): def __init__( self, app: Type[ASGIFramework], loop: asyncio.AbstractEventLoop, config: Config, transport: asyncio.BaseTransport, *, upgrade_request: Optional[h11.Request] = None, ) -> None: super().__init__(loop, config, transport, "wsproto") self.stop_keep_alive_timeout() self.app = app self.connection = WSConnection(ConnectionType.SERVER) self.app_queue: asyncio.Queue = asyncio.Queue() self.response: Optional[dict] = None self.scope: Optional[dict] = None self.state = ASGIWebsocketState.HANDSHAKE self.task: Optional[asyncio.Future] = None self.buffer = WebsocketBuffer(self.config.websocket_max_message_size) if upgrade_request is not None: self.connection.initiate_upgrade_connection( upgrade_request.headers, upgrade_request.target) self.handle_events() def connection_lost(self, error: Optional[Exception]) -> None: if error is not None: self.app_queue.put_nowait({"type": "websocket.disconnect"}) def eof_received(self) -> bool: self.data_received(None) return True def data_received(self, data: Optional[bytes]) -> None: self.connection.receive_data(data) self.handle_events() 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 maybe_close(self, future: asyncio.Future) -> None: # Close the connection iff a HTTP response was sent if self.state == ASGIWebsocketState.HTTPCLOSED: self.close() async def asend(self, event: Event) -> None: await self.drain() self.write(self.connection.send(event)) async def asgi_put(self, message: dict) -> None: await self.app_queue.put(message) async def asgi_receive(self) -> dict: """Called by the ASGI instance to receive a message.""" return await self.app_queue.get() @property def scheme(self) -> str: return "wss" if self.ssl_info is not None else "ws"
class Websocket: _scope = None _sock = None _connection = None def __init__(self): self._byte_buffer = BytesIO() self._string_buffer = StringIO() self._closed = False async def __ainit__(self, addr, path: str, headers: Optional[List] = None, subprotocols=None, **connect_kw): sock = await anyio.connect_tcp(*addr, **connect_kw) await self.start_client( sock, addr, path=path, headers=headers, subprotocols=subprotocols) async def start_client(self, sock: anyio.abc.SocketStream, addr, path: str, headers: Optional[List] = None, subprotocols: Optional[List[str]] = None): """Start a client WS connection on this socket. Returns: the AcceptConnection message. """ self._sock = sock self._connection = WSConnection(ConnectionType.CLIENT) if headers is None: headers = [] if subprotocols is None: subprotocols = [] data = self._connection.send( Request( host=addr[0], target=path, extra_headers=headers, subprotocols=subprotocols)) await self._sock.send_all(data) assert self._scope is None self._scope = True try: event = await self._next_event() if not isinstance(event, AcceptConnection): raise ConnectionError("Failed to establish a connection", event) return event finally: self._scope = None async def start_server(self, sock: anyio.abc.SocketStream, filter=None): # pylint: disable=W0622 """Start a server WS connection on this socket. Filter: an async callable that gets passed the initial Request. It may return an AcceptConnection message, a bool, or a string (the subprotocol to use). Returns: the Request message. """ assert self._scope is None self._scope = True self._sock = sock self._connection = WSConnection(ConnectionType.SERVER) try: event = await self._next_event() if not isinstance(event, Request): raise ConnectionError("Failed to establish a connection", event) msg = None if filter is not None: msg = await filter(event) if not msg: msg = RejectConnection() elif msg is True: msg = None elif isinstance(msg, str): msg = AcceptConnection(subprotocol=msg) if not msg: msg = AcceptConnection(subprotocol=event.subprotocols[0]) data = self._connection.send(msg) await self._sock.send_all(data) if not isinstance(msg, AcceptConnection): raise ConnectionError("Not accepted", msg) finally: self._scope = None 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 else: return event data = await self._sock.receive_some(4096) if not data: return CloseConnection(code=500, reason="Socket closed") self._connection.receive_data(data) 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 send(self, data: Union[bytes, str], final: bool = True): """ Sends some data down the connection. """ MsgType = TextMessage if isinstance(data, str) else BytesMessage data = MsgType(data=data, message_finished=final) data = self._connection.send(event=data) await self._sock.send_all(data) def _buffer(self, event: Message): """ Buffers an event, if applicable. """ if isinstance(event, BytesMessage): self._byte_buffer.write(event.data) elif isinstance(event, TextMessage): self._string_buffer.write(event.data) def _gather_buffers(self, event: Message): """ Gathers all the data from a buffer. """ if isinstance(event, BytesMessage): buf = self._byte_buffer else: buf = self._string_buffer # yay for code shortening buf.write(event.data) buf.seek(0) data = buf.read() buf.seek(0) buf.truncate() return data @staticmethod def _wrap_data(data: Union[str, bytes]): """ Wraps data into the right event. """ MsgType = TextMessage if isinstance(data, str) else BytesMessage return MsgType(data=data, frame_finished=True, message_finished=True) async def __aiter__(self): async with anyio.open_cancel_scope() as scope: if self._scope is not None: raise RuntimeError( "Only one task may iterate on this web socket") self._scope = scope try: while True: msg = await self._next_event() if isinstance(msg, CloseConnection): return yield msg finally: self._scope = None
sock.sendall( connection.send(Request( host=uri.netloc, target='%s?%s' % (uri.path, uri.query), extensions=[PerMessageDeflate()], )) ) closed = False while not closed: try: data = sock.recv(65535) except CONNECTION_EXCEPTIONS: data = None connection.receive_data(data or None) out_data = b"" for event in connection.events(): if isinstance(event, Message): out_data += connection.send(Message(data=event.data, message_finished=event.message_finished)) elif isinstance(event, Ping): out_data += connection.send(event.response()) elif isinstance(event, CloseConnection): closed = True out_data += connection.send(event.response()) # else: # print("??", event) if out_data is None: break try: sock.sendall(out_data) except CONNECTION_EXCEPTIONS: closed = True
def handle_connection(stream): ''' Handle a connection. The server operates a request/response cycle, so it performs a synchronous loop: 1) Read data from network into wsproto 2) Get next wsproto event 3) Handle event 4) Send data from wsproto to network :param stream: a socket stream ''' ws = WSConnection(ConnectionType.SERVER) # events is a generator that yields websocket event objects. Usually you # would say `for event in ws.events()`, but the synchronous nature of this # server requires us to use next(event) instead so that we can interleave # the network I/O. events = ws.events() running = True while running: # 1) Read data from network in_data = stream.recv(RECEIVE_BYTES) print('Received {} bytes'.format(len(in_data))) ws.receive_data(in_data) # 2) Get next wsproto event try: event = next(events) except StopIteration: print('Client connection dropped unexpectedly') return # 3) Handle event if isinstance(event, Request): # Negotiate new WebSocket connection print('Accepting WebSocket upgrade') out_data = ws.send(AcceptConnection()) elif isinstance(event, CloseConnection): # Print log message and break out print('Connection closed: code={}/{} reason={}'.format( event.code.value, event.code.name, event.reason)) out_data = ws.send(event.response()) running = False elif isinstance(event, TextMessage): # Reverse text and send it back to wsproto print('Received request and sending response') out_data = ws.send(Message(data=event.data[::-1])) elif isinstance(event, Ping): # wsproto handles ping events for you by placing a pong frame in # the outgoing buffer. You should not call pong() unless you want to # send an unsolicited pong frame. print('Received ping and sending pong') out_data = ws.send(event.response()) else: print('Unknown event: {!r}'.format(event)) # 4) Send data from wsproto to network print('Sending {} bytes'.format(len(out_data))) stream.send(out_data)