def test_connection_send_state() -> None:
client = WSConnection(CLIENT)
assert client.state is ConnectionState.CONNECTING
server = h11.Connection(h11.SERVER)
server.receive_data(client.send(Request(
host="localhost",
target="/",
)))
headers = normed_header_dict(server.next_event().headers)
response = h11.InformationalResponse(
status_code=101,
headers=[
(b"connection", b"Upgrade"),
(b"upgrade", b"WebSocket"),
(
b"Sec-WebSocket-Accept",
generate_accept_token(headers[b"sec-websocket-key"]),
),
],
)
client.receive_data(server.send(response))
assert len(list(client.events())) == 1
assert client.state is ConnectionState.OPEN # type: ignore # https://github.com/python/mypy/issues/9005
with pytest.raises(LocalProtocolError) as excinfo:
client.send(Request(host="localhost", target="/"))
client.receive_data(b"foobar")
assert len(list(client.events())) == 1
def test_connection_request_simple_extension_no_offer() -> None:
extension = FakeExtension(offer_response=False)
request = _make_connection_request(
Request(host="localhost", target="/", extensions=[extension]))
headers = normed_header_dict(request.headers)
assert b"sec-websocket-extensions" not in headers
def test_connection_request_subprotocols():
request = _make_connection_request(
Request(host="localhost", target="/", subprotocols=["one", "two"])
)
headers = normed_header_dict(request.headers)
assert headers[b"sec-websocket-protocol"] == b"one, two"
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 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 _make_handshake(
response_status,
response_headers,
subprotocols=None,
extensions=None,
auto_accept_key=True,
):
client = WSConnection(CLIENT)
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))
return list(client.events())
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 init_for_client(
cls: Type["Transport"], stream: Stream, host: str, keepalive: Optional[int] = None
) -> "Transport":
ws = WSConnection(ConnectionType.CLIENT)
transport = cls(stream, ws, keepalive)
# Because this is a client WebSocket, we need to initiate the connection
# handshake by sending a Request event.
await transport._net_send(
Request(
host=host,
target=TRANSPORT_TARGET,
extra_headers=[(b"User-Agent", USER_AGENT.encode())],
)
)
# Get handshake answer
event = await transport._next_ws_event()
if isinstance(event, AcceptConnection):
transport.logger.debug("WebSocket negotiation complete", ws_event=event)
else:
transport.logger.warning("Unexpected event during WebSocket handshake", ws_event=event)
reason = f"Unexpected event during WebSocket handshake: {event}"
raise TransportError(reason)
return transport
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 __init__(self, ctx, handshake_flow):
super().__init__(ctx)
self.handshake_flow = handshake_flow
self.flow: WebSocketFlow = None
self.client_frame_buffer = []
self.server_frame_buffer = []
self.connections: dict[object, WSConnection] = {}
client_extensions = []
server_extensions = []
if 'Sec-WebSocket-Extensions' in handshake_flow.response.headers:
if PerMessageDeflate.name in handshake_flow.response.headers['Sec-WebSocket-Extensions']:
client_extensions = [PerMessageDeflate()]
server_extensions = [PerMessageDeflate()]
self.connections[self.client_conn] = WSConnection(ConnectionType.SERVER)
self.connections[self.server_conn] = WSConnection(ConnectionType.CLIENT)
if client_extensions:
client_extensions[0].finalize(handshake_flow.response.headers['Sec-WebSocket-Extensions'])
if server_extensions:
server_extensions[0].finalize(handshake_flow.response.headers['Sec-WebSocket-Extensions'])
request = Request(extensions=client_extensions, host=handshake_flow.request.host, target=handshake_flow.request.path)
data = self.connections[self.server_conn].send(request)
self.connections[self.client_conn].receive_data(data)
event = next(self.connections[self.client_conn].events())
assert isinstance(event, events.Request)
data = self.connections[self.client_conn].send(AcceptConnection(extensions=server_extensions))
self.connections[self.server_conn].receive_data(data)
assert isinstance(next(self.connections[self.server_conn].events()), events.AcceptConnection)
def test_connection_request_simple_extension():
extension = FakeExtension(offer_response=True)
request = _make_connection_request(
Request(host="localhost", target="/", extensions=[extension])
)
headers = normed_header_dict(request.headers)
assert headers[b"sec-websocket-extensions"] == extension.name.encode("ascii")
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.')
async def _do_handshake(self):
try:
handshake = Request(host=self.host,
target=self.path,
subprotocols=self._subprotocols or [],
extra_headers=self._headers or [])
await self._send(handshake)
except ConnectionClosed:
self._reader_running = False
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")))
def test_connection_request_additional_headers() -> None:
request = _make_connection_request(
Request(
host="localhost",
target="/",
extra_headers=[(b"X-Foo", b"Bar"), (b"X-Bar", b"Foo")],
))
headers = normed_header_dict(request.headers)
assert headers[b"x-foo"] == b"Bar"
assert headers[b"x-bar"] == b"Foo"
def handshake(self):
out_data = self.ws.send(Request(host=self.host, target=self.path))
self.sock.send(out_data)
in_data = self.sock.recv(self.receive_bytes)
self.ws.receive_data(in_data)
for event in self.ws.events():
if isinstance(event, AcceptConnection):
break
elif isinstance(event, RejectConnection):
raise ConnectionError(event.status_code)
def test_connection_request_parametrised_extension():
extension = FakeExtension(offer_response="parameter1=value1; parameter2=value2")
request = _make_connection_request(
Request(host="localhost", target="/", extensions=[extension])
)
headers = normed_header_dict(request.headers)
assert headers[b"sec-websocket-extensions"] == b"%s; %s" % (
extension.name.encode("ascii"),
extension.offer_response.encode("ascii"),
)
def test_successful_handshake():
client = H11Handshake(CLIENT)
server = H11Handshake(SERVER)
server.receive_data(client.send(Request(host="localhost", target="/")))
assert isinstance(next(server.events()), Request)
client.receive_data(server.send(AcceptConnection()))
assert isinstance(next(client.events()), AcceptConnection)
assert client.state is ConnectionState.OPEN
assert server.state is ConnectionState.OPEN
def test_rejected_handshake():
client = H11Handshake(CLIENT)
server = H11Handshake(SERVER)
server.receive_data(client.send(Request(host="localhost", target="/")))
assert isinstance(next(server.events()), Request)
client.receive_data(server.send(RejectConnection()))
assert isinstance(next(client.events()), RejectConnection)
assert client.state is ConnectionState.CLOSED
assert server.state is ConnectionState.CLOSED
def test_connection_request():
request = _make_connection_request(Request(host="localhost", target="/"))
assert request.http_version == b"1.1"
assert request.method == b"GET"
assert request.target == b"/"
headers = normed_header_dict(request.headers)
assert headers[b"connection"] == b"Upgrade"
assert headers[b"host"] == b"localhost"
assert headers[b"sec-websocket-version"] == b"13"
assert headers[b"upgrade"] == b"WebSocket"
assert b"sec-websocket-key" in headers
def __init__(
self,
path: str,
event_loop: asyncio.AbstractEventLoop,
*,
framework: Type[ASGIFramework] = EchoFramework,
) -> None:
self.transport = MockTransport()
self.server = WebsocketServer( # type: ignore
framework, event_loop, Config(), self.transport)
self.connection = WSConnection(ConnectionType.CLIENT)
self.server.data_received(
self.connection.send(Request(target=path, host="hypercorn")))
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 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)]
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())
async def test_bad_framework_http(path: str) -> 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=path, host="hypercorn", extra_headers=headers)
await server.handle_websocket(request)
assert server.sent_events == [
RejectConnection(status_code=500,
headers=[(b"server", b"hypercorn")],
has_body=False)
]
async def init_for_client(cls, stream: Stream, host: str) -> "Transport":
ws = WSConnection(ConnectionType.CLIENT)
transport = cls(stream, ws)
# Because this is a client WebSocket, we need to initiate the connection
# handshake by sending a Request event.
await transport._net_send(Request(host=host, target=TRANSPORT_TARGET))
# Get handshake answer
event = await transport._next_ws_event()
if isinstance(event, AcceptConnection):
transport.logger.debug("WebSocket negotiation complete", ws_event=event)
else:
transport.logger.warning("Unexpected event during WebSocket handshake", ws_event=event)
reason = f"Unexpected event during WebSocket handshake: {event}"
raise TransportError(reason)
return transport
async def test_asgi_scope() -> None:
server = MockWebsocket()
connection_request = Request(target="/path?a=b", host="hypercorn")
await server.handle_websocket(connection_request)
scope = server.scope
assert scope == {
"type": "websocket",
"asgi": {
"version": "2.0"
},
"http_version": "1.1",
"scheme": "ws",
"path": "/path",
"query_string": b"a=b",
"root_path": "",
"headers": [(b"host", b"hypercorn")],
"client": ("127.0.0.1", 5000),
"server": ("remote", 5000),
"subprotocols": [],
"extensions": {
"websocket.http.response": {}
},
}
def test_send_invalid_event() -> None:
client = Connection(CLIENT)
with pytest.raises(LocalProtocolError):
client.send(Request(target="/", host="wsproto"))
sock.sendall(out_data)
except CONNECTION_EXCEPTIONS:
break
sock.close()
return case_count
def run_case(server, case, agent):
uri = urlparse(server + '/runCase?case=%d&agent=%s' % (case, 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),
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):