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 test_buffer() -> None:
buffer_ = WebsocketBuffer(10)
buffer_.extend(TextMessage(data="abc", frame_finished=False, message_finished=True))
assert buffer_.to_message() == {"type": "websocket.receive", "bytes": None, "text": "abc"}
buffer_.clear()
buffer_.extend(BytesMessage(data=b"abc", frame_finished=False, message_finished=True))
assert buffer_.to_message() == {"type": "websocket.receive", "bytes": b"abc", "text": None}
async def send_message(self, message: Union[bytes, str]):
"""Sends a WebSocket message to the peer.
Parameters
----------
message : Union[bytes, str]
The message to send.
Raises
------
:exc:`anysocks.exceptions.ConnectionClosed`
If the connection is already closed.
:exc:`ValueError`
If the type of ``message`` isn't ``str`` or ``bytes``.
"""
if self._closed:
raise ConnectionClosed(self._close_code, self._close_reason)
if isinstance(message, str):
event = TextMessage(data=message)
elif isinstance(message, bytes):
event = BytesMessage(data=message)
else:
raise ValueError('Message must be bytes or string')
await self._send(event)
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"])
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)
]
def test_send_message(client_sends: bool, final: bool) -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
if client_sends:
local = client
remote = server
else:
local = server
remote = client
data = b"x" * 23
remote.receive_data(
local.send(BytesMessage(data=data, message_finished=final)))
event = next(remote.events())
assert isinstance(event, BytesMessage)
assert event.data == data
assert event.message_finished is final
async def send(self, msg: bytes) -> None:
"""
Raises:
TransportError
"""
await self._net_send(BytesMessage(data=msg))
def test_buffer_frame_too_large() -> None:
buffer_ = WebsocketBuffer(2)
with pytest.raises(FrameTooLarge):
buffer_.extend(
TextMessage(data="abc",
frame_finished=False,
message_finished=True))
@pytest.mark.parametrize(
"data",
[
(
TextMessage(
data="abc", frame_finished=False, message_finished=True),
BytesMessage(
data=b"abc", frame_finished=False, message_finished=True),
),
(
BytesMessage(
data=b"abc", frame_finished=False, message_finished=True),
TextMessage(
data="abc", frame_finished=False, message_finished=True),
),
],
)
def test_buffer_mixed_types(data: list) -> None:
buffer_ = WebsocketBuffer(10)
buffer_.extend(data[0])
with pytest.raises(TypeError):
buffer_.extend(data[1])
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
})