Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_data(split_message: bool) -> None:
client = Connection(CLIENT)
server = Connection(SERVER)
data = "ƒñö®∂😎"
server.receive_data(
client.send(TextMessage(data=data, message_finished=not split_message))
)
event = next(server.events())
assert isinstance(event, TextMessage)
assert event.message_finished is not split_message
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])
server = TCPServer(sanity_framework, Config(), server_stream)
nursery.start_soon(server.run)
client = wsproto.WSConnection(wsproto.ConnectionType.CLIENT)
await client_stream.send_all(client.send(wsproto.events.Request(host="hypercorn", target="/")))
client.receive_data(await client_stream.receive_some(1024))
assert list(client.events()) == [
wsproto.events.AcceptConnection(
extra_headers=[
(b"date", b"Thu, 01 Jan 1970 01:23:20 GMT"),
(b"server", b"hypercorn-h11"),
]
)
]
await client_stream.send_all(client.send(wsproto.events.BytesMessage(data=SANITY_BODY)))
client.receive_data(await client_stream.receive_some(1024))
assert list(client.events()) == [wsproto.events.TextMessage(data="Hello & Goodbye")]
await client_stream.send_all(client.send(wsproto.events.CloseConnection(code=1000)))
client.receive_data(await client_stream.receive_some(1024))
assert list(client.events()) == [wsproto.events.CloseConnection(code=1000, reason="")]
def get_case_count(server: str) -> int:
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: Optional[int] = None
while case_count is None:
in_data = sock.recv(65535)
connection.receive_data(in_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
async def send(self, message: str):
"""
Send a message.
"""
assert isinstance(message, str)
data = self.websocket.send(wsproto.events.TextMessage(data=message))
self.http.send_data(stream_id=self.stream_id, data=data, end_stream=False)
self.transmit()
):
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:
event: wsproto.events.Event
if message.get("bytes") is not None:
event = wsproto.events.BytesMessage(data=bytes(message["bytes"]))
elif not isinstance(message["text"], str):
raise TypeError(f"{message['text']} should be a str")
else:
event = wsproto.events.TextMessage(data=message["text"])
await self.asend(Data(self.connection.send(event)))
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":
data = self.connection.send(wsproto.events.CloseConnection(code=int(message["code"])))
await self.asend(Data(data))
self.state = ASGIWebsocketState.CLOSED
else:
raise UnexpectedMessage(self.state, message["type"])
async def _reader_task(self):
handlers = {
AcceptConnection: self._handle_accept_connection_event,
BytesMessage: self._handle_message_event,
CloseConnection: self._handle_close_connection_event,
Ping: self._handle_ping_event,
Pong: self._handle_pong_event,
RejectConnection: self._handle_reject_connection_event,
RejectData: self._handle_reject_data_event,
# Request: lambda event: None, # We won't handle server-side events.
TextMessage: self._handle_message_event,
}
# We need to initiate the opening handshake.
await self._do_handshake()
while self._reader_running:
for event in self._wsproto.events():
event_type = type(event)
try:
handler = handlers[event_type]
logger.debug('%s received event: %s', self, event_type)
await handler(event)
except KeyError:
logger.warning('%s received unknown event type: "%s"', self, event_type)
except ConnectionClosed:
self._reader_running = False
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))
elif isinstance(event, Pong):
print("Received pong: {!r}".format(event.payload))
else:
raise Exception("Do not know how to handle event: " + str(event))
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("Unknown event: {!r}".format(event))
# 4) Send data from wsproto to network
print("Sending {} bytes".format(len(out_data)))
stream.send(out_data)