async def test_asgi_send_http() -> None:
server = MockWebsocket()
await server.asgi_send({
"type": "websocket.http.response.start",
"headers": [(b"X-Header", b"Value")],
"status": 200,
})
# Server must not send a response till the receipt of the first
# body chunk.
assert server.sent_events == []
await server.asgi_send({
"type": "websocket.http.response.body",
"body": b"a",
"more_body": True
})
await server.asgi_send({
"type": "websocket.http.response.body",
"more_body": False
})
server.sent_events[0].headers = list(
server.sent_events[0].headers) # To allow comparison
assert server.sent_events == [
RejectConnection(
status_code=200,
headers=[(b"x-header", b"Value"), (b"server", b"hypercorn")],
has_body=True,
),
RejectData(data=b"a", body_finished=False),
RejectData(data=b"", body_finished=True),
]
async def _asgi_send_rejection(self, message: dict) -> None:
body_suppressed = suppress_body("GET", self.response["status"])
if self.state == ASGIWebsocketState.HANDSHAKE:
headers = chain(
[
(bytes(key).strip(), bytes(value).strip())
for key, value in self.response["headers"]
],
self.response_headers(),
)
await self.asend(
RejectConnection(
status_code=int(self.response["status"]),
headers=headers,
has_body=not body_suppressed,
)
)
self.state = ASGIWebsocketState.RESPONSE
if not body_suppressed:
await self.asend(
RejectData(
data=bytes(message.get("body", b"")),
body_finished=not message.get("more_body", False),
)
)
if not message.get("more_body", False):
await self.asgi_put({"type": "websocket.disconnect"})
self.state = ASGIWebsocketState.HTTPCLOSED
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
def test_handshake_rejection():
events = _make_handshake_rejection(400)
assert events == [
RejectConnection(
headers=[(b"connection", b"close")], has_body=True, status_code=400
),
RejectData(body_finished=True, data=b""),
]
def test_handshake_rejection_with_body():
events = _make_handshake_rejection(400, b"Hello")
assert events == [
RejectConnection(
headers=[(b"content-length", b"5")], has_body=True, status_code=400
),
RejectData(body_finished=False, data=b"Hello"),
RejectData(body_finished=True, data=b""),
]
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 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 _make_handshake_rejection(
status_code: int, body: Optional[bytes] = None
) -> List[Event]:
client = h11.Connection(h11.CLIENT)
server = WSConnection(SERVER)
nonce = generate_nonce()
server.receive_data(
client.send(
h11.Request(
method="GET",
target="/",
headers=[
(b"Host", b"localhost"),
(b"Connection", b"Keep-Alive, Upgrade"),
(b"Upgrade", b"WebSocket"),
(b"Sec-WebSocket-Version", b"13"),
(b"Sec-WebSocket-Key", nonce),
],
)
)
)
if body is not None:
client.receive_data(
server.send(
RejectConnection(
headers=[(b"content-length", b"%d" % len(body))],
status_code=status_code,
has_body=True,
)
)
)
client.receive_data(server.send(RejectData(data=body)))
else:
client.receive_data(server.send(RejectConnection(status_code=status_code)))
events = []
while True:
event = client.next_event()
events.append(event)
if isinstance(event, h11.EndOfMessage):
return events
def test_connection_request_bad_version_header(version: bytes) -> None:
with pytest.raises(RemoteProtocolError) as excinfo:
event = _make_connection_request([
(b"Host", b"localhost"),
(b"Connection", b"Keep-Alive, Upgrade"),
(b"Upgrade", b"WebSocket"),
(b"Sec-WebSocket-Version", version),
(b"Sec-WebSocket-Key", generate_nonce()),
])
assert str(excinfo.value) == "Missing header, 'Sec-WebSocket-Version'"
assert excinfo.value.event_hint == RejectConnection(headers=[
(b"Sec-WebSocket-Version", b"13")
],
status_code=426)
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 send_http_error(self, status: int) -> None:
await self.asend(RejectConnection(status_code=status, headers=self.response_headers()))
self.config.access_logger.access(
self.scope, {"status": status, "headers": []}, time() - self.start_time
)