예제 #1
0
 async def asgi_send(self, message: dict) -> None:
     if message[
             "type"] == "websocket.accept" and self.state == ASGIWebsocketState.HANDSHAKE:
         self.state = ASGIWebsocketState.CONNECTED
         extensions: List[str] = []
         for name, value in self.scope["headers"]:
             if name == b"sec-websocket-extensions":
                 extensions = split_comma_header(value)
         supported_extensions = [wsproto.extensions.PerMessageDeflate()]
         accepts = handshake_extensions(extensions, supported_extensions)
         headers = [(b":status", b"200")]
         headers.extend(
             build_and_validate_headers(message.get("headers", [])))
         raise_if_subprotocol_present(headers)
         if message.get("subprotocol") is not None:
             headers.append((b"sec-websocket-protocol",
                             message["subprotocol"].encode()))
         if accepts:
             headers.append((b"sec-websocket-extensions", accepts))
         await self.asend(Response(headers))
         self.connection = wsproto.connection.Connection(
             wsproto.connection.ConnectionType.SERVER, supported_extensions)
         self.config.access_logger.access(self.scope, {
             "status": 200,
             "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:
         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"])
예제 #2
0
 def __init__(self, headers: List[Tuple[bytes, bytes]], http_version: str) -> None:
     self.http_version = http_version
     self.connection_tokens: Optional[List[str]] = None
     self.extensions: Optional[List[str]] = None
     self.key: Optional[bytes] = None
     self.subprotocols: Optional[List[str]] = None
     self.upgrade: Optional[bytes] = None
     self.version: Optional[bytes] = None
     for name, value in headers:
         name = name.lower()
         if name == b"connection":
             self.connection_tokens = split_comma_header(value)
         elif name == b"sec-websocket-extensions":
             self.extensions = split_comma_header(value)
         elif name == b"sec-websocket-key":
             self.key = value
         elif name == b"sec-websocket-protocol":
             self.subprotocols = split_comma_header(value)
         elif name == b"sec-websocket-version":
             self.version = value
         elif name == b"upgrade":
             self.upgrade = value