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))
