async def wrap_server_stream(nursery,
stream,
message_queue_size=MESSAGE_QUEUE_SIZE,
max_message_size=MAX_MESSAGE_SIZE):
'''
Wrap an arbitrary stream in a server-side WebSocket.
This is a low-level function only needed in rare cases. In most cases, you
should use :func:`serve_websocket`.
:param nursery: A nursery to run background tasks in.
:param stream: A stream to be wrapped.
:param int message_queue_size: The maximum number of messages that will be
buffered in the library's internal message queue.
:param int max_message_size: The maximum message size as measured by
``len()``. If a message is received that is larger than this size,
then the connection is closed with code 1009 (Message Too Big).
:type stream: trio.abc.Stream
:rtype: WebSocketConnection
'''
wsproto = wsconnection.WSConnection(wsconnection.SERVER)
connection = WebSocketConnection(stream,
wsproto,
message_queue_size=message_queue_size,
max_message_size=max_message_size)
nursery.start_soon(connection._reader_task)
request = await connection._get_request()
return request
async def wrap_client_stream(nursery, stream, host, resource, *,
subprotocols=None, message_queue_size=MESSAGE_QUEUE_SIZE,
max_message_size=MAX_MESSAGE_SIZE):
'''
Wrap an arbitrary stream in a WebSocket connection.
This is a low-level function only needed in rare cases. In most cases, you
should use :func:`open_websocket` or :func:`open_websocket_url`.
:param nursery: A Trio nursery to run background tasks in.
:param stream: A Trio stream to be wrapped.
:type stream: trio.abc.Stream
:param str host: A host string that will be sent in the ``Host:`` header.
:param str resource: A resource string, i.e. the path component to be
accessed on the server.
:param subprotocols: An iterable of strings representing preferred
subprotocols.
:param int message_queue_size: The maximum number of messages that will be
buffered in the library's internal message queue.
:param int max_message_size: The maximum message size as measured by
``len()``. If a message is received that is larger than this size,
then the connection is closed with code 1009 (Message Too Big).
:rtype: WebSocketConnection
'''
wsproto = wsconnection.WSConnection(wsconnection.CLIENT, host=host,
resource=resource, subprotocols=subprotocols)
connection = WebSocketConnection(stream, wsproto, path=resource,
message_queue_size=message_queue_size,
max_message_size=max_message_size)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
return connection
async def _handle_connection(self, stream):
'''
Handle an incoming connection by spawning a connection background task
and a handler task inside a new nursery.
:param stream:
:type stream: trio.abc.Stream
'''
async with trio.open_nursery() as nursery:
wsproto = wsconnection.WSConnection(wsconnection.SERVER)
connection = WebSocketConnection(
stream,
wsproto,
message_queue_size=self._message_queue_size,
max_message_size=self._max_message_size)
nursery.start_soon(connection._reader_task)
with trio.move_on_after(self._connect_timeout) as connect_scope:
request = await connection._get_request()
if connect_scope.cancelled_caught:
nursery.cancel_scope.cancel()
await stream.aclose()
return
try:
await self._handler(request)
finally:
with trio.move_on_after(self._disconnect_timeout):
# aclose() will shut down the reader task even if its
# cancelled:
await connection.aclose()
async def connect(self, nursery):
'''
Connect to WebSocket server.
:param nursery: a Trio nursery to run background connection tasks in
:raises: OSError if connection attempt fails
'''
logger.info('Connecting to http%s://%s:%d/%s',
'' if self._ssl is None else 's', self._host, self._port,
self._resource)
if self._ssl is None:
stream = await trio.open_tcp_stream(self._host, self._port)
else:
stream = await trio.open_ssl_over_tcp_stream(self._host,
self._port,
ssl_context=self._ssl,
https_compatible=True)
if self._port in (80, 443):
host_header = self._host
else:
host_header = '{}:{}'.format(self._host, self._port)
wsproto = wsconnection.WSConnection(wsconnection.CLIENT,
host=host_header,
resource=self._resource)
connection = WebSocketConnection(stream, wsproto)
nursery.start_soon(connection._reader_task)
return connection
async def _handle_connection(self, stream):
''' Handle an incoming connection. '''
async with trio.open_nursery() as nursery:
wsproto = wsconnection.WSConnection(wsconnection.SERVER)
connection = WebSocketConnection(stream, wsproto)
nursery.start_soon(connection._reader_task)
nursery.start_soon(self._handler, connection)
async def connect_websocket(nursery, host, port, resource, *, use_ssl,
subprotocols=None, message_queue_size=MESSAGE_QUEUE_SIZE,
max_message_size=MAX_MESSAGE_SIZE):
'''
Return an open WebSocket client connection to a host.
This function is used to specify a custom nursery to run connection
background tasks in. The caller is responsible for closing the connection.
If you don't need a custom nursery, you should probably use
:func:`open_websocket` instead.
:param nursery: A Trio nursery to run background tasks in.
:param str host: The host to connect to.
:param int port: The port to connect to.
:param str resource: The resource, i.e. URL path.
:type use_ssl: bool or ssl.SSLContext
:param subprotocols: An iterable of strings representing preferred
subprotocols.
:param int message_queue_size: The maximum number of messages that will be
buffered in the library's internal message queue.
:param int max_message_size: The maximum message size as measured by
``len()``. If a message is received that is larger than this size,
then the connection is closed with code 1009 (Message Too Big).
:rtype: WebSocketConnection
'''
if use_ssl == True:
ssl_context = ssl.create_default_context()
elif use_ssl == False:
ssl_context = None
elif isinstance(use_ssl, ssl.SSLContext):
ssl_context = use_ssl
else:
raise TypeError('`use_ssl` argument must be bool or ssl.SSLContext')
logger.debug('Connecting to ws%s://%s:%d%s',
'' if ssl_context is None else 's', host, port, resource)
if ssl_context is None:
stream = await trio.open_tcp_stream(host, port)
else:
stream = await trio.open_ssl_over_tcp_stream(host, port,
ssl_context=ssl_context, https_compatible=True)
if port in (80, 443):
host_header = host
else:
host_header = '{}:{}'.format(host, port)
wsproto = wsconnection.WSConnection(wsconnection.CLIENT,
host=host_header, resource=resource, subprotocols=subprotocols)
connection = WebSocketConnection(stream, wsproto, path=resource,
message_queue_size=message_queue_size,
max_message_size=max_message_size)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
return connection
async def _handle_connection(self, stream):
'''
Handle an incoming connection by spawning a connection background task
and a handler task inside a new nursery.
:param stream:
:type stream: trio.abc.Stream
'''
async with trio.open_nursery() as nursery:
wsproto = wsconnection.WSConnection(wsconnection.SERVER)
connection = WebSocketConnection(stream, wsproto)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
nursery.start_soon(self._handler, connection)
async def wrap_server_stream(nursery, stream):
'''
Wrap an arbitrary stream in a server-side ``WebSocketConnection``.
This is a low-level function only needed in rare cases. Most users should
call ``serve_websocket()`.
:param nursery: A Trio nursery to run background tasks in.
:param stream: A Trio stream to be wrapped.
:param task_status: part of Trio nursery start protocol
:rtype: WebSocketConnection
'''
wsproto = wsconnection.WSConnection(wsconnection.SERVER)
connection = WebSocketConnection(stream, wsproto)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
return connection
async def connect_websocket(nursery, host, port, resource, use_ssl):
'''
Return a WebSocket client connection to a host.
Most users should use ``open_websocket(…)`` instead of this function. This
function is not an async context manager, and it requires a nursery argument
for the connection's background task[s]. The caller is responsible for
closing the connection.
:param nursery: a Trio nursery to run background tasks in
:param str host: the host to connect to
:param int port: the port to connect to
:param str resource: the resource a.k.a. path
:param use_ssl: a bool or SSLContext
:rtype: WebSocketConnection
'''
if use_ssl == True:
ssl_context = ssl.create_default_context()
elif use_ssl == False:
ssl_context = None
elif isinstance(use_ssl, ssl.SSLContext):
ssl_context = use_ssl
else:
raise TypeError('`use_ssl` argument must be bool or ssl.SSLContext')
logger.debug('Connecting to ws%s://%s:%d%s',
'' if ssl_context is None else 's', host, port, resource)
if ssl_context is None:
stream = await trio.open_tcp_stream(host, port)
else:
stream = await trio.open_ssl_over_tcp_stream(host, port,
ssl_context=ssl_context, https_compatible=True)
if port in (80, 443):
host_header = host
else:
host_header = '{}:{}'.format(host, port)
wsproto = wsconnection.WSConnection(wsconnection.CLIENT,
host=host_header, resource=resource)
connection = WebSocketConnection(stream, wsproto, path=resource)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
return connection
async def wrap_client_stream(nursery, stream, host, resource):
'''
Wrap an arbitrary stream in a client-side ``WebSocketConnection``.
This is a low-level function only needed in rare cases. Most users should
call ``open_websocket()`` or ``open_websocket_url()``.
:param nursery: A Trio nursery to run background tasks in.
:param stream: A Trio stream to be wrapped.
:param str host: A host string that will be sent in the ``Host:`` header.
:param str resource: A resource string, i.e. the path component to be
accessed on the server.
:rtype: WebSocketConnection
'''
wsproto = wsconnection.WSConnection(wsconnection.CLIENT, host=host,
resource=resource)
connection = WebSocketConnection(stream, wsproto, path=resource)
nursery.start_soon(connection._reader_task)
await connection._open_handshake.wait()
return connection
