Python trio.serve_tcp() Examples

def test_reject_handshake_invalid_info_status(nursery):
    '''
    An informational status code that is not 101 should cause the client to
    reject the handshake. Since it is an informational response, there will not
    be a response body, so this test exercises a different code path.
    '''
    async def handler(stream):
        await stream.send_all(b'HTTP/1.1 100 CONTINUE\r\n\r\n')
        await stream.receive_some(max_bytes=1024)
    serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
    listeners = await nursery.start(serve_fn)
    port = listeners[0].socket.getsockname()[1]

    with pytest.raises(ConnectionRejected) as exc_info:
        async with open_websocket(HOST, port, RESOURCE, use_ssl=False,
                ) as client_ws:
            pass
    exc = exc_info.value
    assert exc.status_code == 100
    assert repr(exc) == 'ConnectionRejected<status_code=100>'
    assert exc.body is None 

def serve_tcp_testbed(unused_tcp_port):
    async def _serve_tcp_testbed(*conns):
        host = "127.0.0.1"
        send_channel, receive_channel = trio.open_memory_channel(0)

        async def _serve_client(stream):
            server_fn = await receive_channel.receive()
            transport = await Transport.init_for_server(stream)
            await server_fn(transport)

        async def _store_handlers(*, task_status=trio.TASK_STATUS_IGNORED):
            async with trio.open_service_nursery() as handler_nursery:
                task_status.started(handler_nursery)
                await trio.sleep_forever()

        async with trio.open_service_nursery() as nursery:
            handler_nursery = await nursery.start(_store_handlers)
            await nursery.start(
                partial(
                    trio.serve_tcp, _serve_client, unused_tcp_port, handler_nursery=handler_nursery
                )
            )
            assert not handler_nursery.child_tasks

            for client_fn, server_fn in conns:
                stream = await trio.open_tcp_stream(host, unused_tcp_port)
                await send_channel.send(server_fn)

                transport = await Transport.init_for_client(stream, host)
                await client_fn(transport)

                await trio.testing.wait_all_tasks_blocked()
                # No pending connections should remain
                assert not handler_nursery.child_tasks

            await send_channel.aclose()
            nursery.cancel_scope.cancel()

    return _serve_tcp_testbed 

def serve(
        self, task_status: TaskStatus[int] = trio.TASK_STATUS_IGNORED
    ) -> None:
        async with trio.open_nursery() as nursery:
            # NOTE: `mypy` does not like the typing here but this
            # should type check...
            listeners: Sequence[trio.SocketListener] = await nursery.start(
                trio.serve_tcp, self.handler, self.port  # type: ignore
            )
            self.port = int(listeners[0].socket.getsockname()[1])
            task_status.started(self.port) 

def test_trio_http_json_server_POST():
    method = "POST"
    path = "/increment"
    api_handlers = {path: {method: post_increment_handler}}
    context = TestContext()
    server = JSONHTTPServer(api_handlers, context)
    with trio.move_on_after(TEST_TIMEOUT):
        async with trio.open_nursery() as nursery:
            some_free_port = 0
            listeners = await nursery.start(
                trio.serve_tcp, server.handler, some_free_port
            )
            client_stream = await open_stream_to_socket_listener(listeners[0])

            body = '{"x": 1}'
            request = (
                f"{method} {path} HTTP/1.0\r\nContent-Length: {len(body)}\r\n\r\n{body}"
            ).encode()
            await client_stream.send_all(request)
            response = bytes()
            async for chunk in client_stream:
                response += chunk
            response_body = response.decode("utf-8").split("\r\n\r\n")[-1]
            result = json.loads(response_body)
            assert result["result"] == 2
            nursery.cancel_scope.cancel() 

def test_trio_http_json_server_GET():
    method = "GET"
    some_param = 33
    path = "/info"
    api_handlers = {path: {method: get_info_handler}}
    context = TestContext()
    server = JSONHTTPServer(api_handlers, context)
    with trio.move_on_after(TEST_TIMEOUT):
        async with trio.open_nursery() as nursery:
            some_free_port = 0
            listeners = await nursery.start(
                trio.serve_tcp, server.handler, some_free_port
            )
            client_stream = await open_stream_to_socket_listener(listeners[0])

            request = (
                f"{method} {path}?some-param={some_param} HTTP/1.0\r\n\r\n"
            ).encode()
            await client_stream.send_all(request)
            response = bytes()
            async for chunk in client_stream:
                response += chunk
            response_body = response.decode("utf-8").split("\r\n\r\n")[-1]
            result = json.loads(response_body)
            assert result == SOME_INFO + str(some_param)
            nursery.cancel_scope.cancel() 

def test_server_does_not_close_handshake(nursery):
    async def handler(stream):
        request = await wrap_server_stream(nursery, stream)
        server_ws = await request.accept()
        async with server_ws:
            await stream.aclose()
            with pytest.raises(ConnectionClosed):
                await server_ws.send_message('Hello from client!')
    serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
    listeners = await nursery.start(serve_fn)
    port = listeners[0].socket.getsockname()[1]
    async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as client:
        with pytest.raises(ConnectionClosed):
            await client.get_message() 

def test_wrap_server_stream(nursery):
    async def handler(stream):
        request = await wrap_server_stream(nursery, stream)
        server_ws = await request.accept()
        async with server_ws:
            assert not server_ws.closed
            msg = await server_ws.get_message()
            assert msg == 'Hello from client!'
        assert server_ws.closed
    serve_fn = partial(trio.serve_tcp, handler, 0, host=HOST)
    listeners = await nursery.start(serve_fn)
    port = listeners[0].socket.getsockname()[1]
    async with open_websocket(HOST, port, RESOURCE, use_ssl=False) as client:
        await client.send_message('Hello from client!') 

def serve(port):
    print("listening on http://localhost:{}".format(port))
    try:
        await trio.serve_tcp(http_serve, port)
    except KeyboardInterrupt:
        print("KeyboardInterrupt - shutting down")

################################################################
# Run the server
################################################################ 

def run(self):
        try:
            self.logger.debug(f'Starting server on {self.bind}:{self.port}')
            ctx = self._ssl_context()
            if ctx:
                await trio.serve_ssl_over_tcp(self.handle, self.port, ctx, host=self.bind)
            else:
                await trio.serve_tcp(self.handle, self.port, host=self.bind)
        except KeyboardInterrupt:
            self.logger.debug('Stopping server')
            sys.exit(0) 

def _run(self):
        self.logger.info(f'Ready to serve requests on {self.bind}:{self.port}')
        ctx = self._ssl_context()
        if ctx:
            await trio.serve_ssl_over_tcp(self.handle, self.port, ctx, host=self.bind)
        else:
            await trio.serve_tcp(self.handle, self.port, host=self.bind) 

def run(self):
        try:
            self.logger.debug(f'Starting server on {self.bind}:{self.port}')
            ctx = self._ssl_context()
            if ctx:
                await trio.serve_ssl_over_tcp(self.handle, self.port, ctx, host=self.bind)
            else:
                await trio.serve_tcp(self.handle, self.port, host=self.bind)
        except KeyboardInterrupt:
            self.logger.debug('Stopping server')
            sys.exit(0) 

def _run(self):
        self.logger.info(f'Ready to serve requests on {self.bind}:{self.port}')
        ctx = self._ssl_context()
        if ctx:
            await trio.serve_ssl_over_tcp(self.handle, self.port, ctx, host=self.bind)
        else:
            await trio.serve_tcp(self.handle, self.port, host=self.bind) 

def test_http_get(nursery, asyncio_loop):
    # Create a server:
    async def handler(stream):
        request = await stream.receive_some(4096)
        assert request.startswith(b'GET /foo?user=john HTTP/1.1\r\n')
        assert request.endswith(b'\r\n\r\n')
        await stream.send_all(
                b'HTTP/1.1 200 OK\r\n'
                b'Content-type: text/html\r\n'
                b'\r\n'
                b'<html><head><title>Unit Test</title></head>\r\n'
                b'<body><h1>This is a unit test</h1></body></html>\r\n'
                b'\r\n'
            )
        await stream.aclose()
    serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1')
    http_server = await nursery.start(serve_tcp)
    addr, port = http_server[0].socket.getsockname()

    # Run the test:
    job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb'
    policy = make_policy()
    request_send, request_recv = trio.open_memory_channel(0)
    response_send, response_recv = trio.open_memory_channel(0)
    semaphore = trio.Semaphore(1)
    rate_limiter_reset = Mock()
    rate_limiter_reset.send = AsyncMock()
    stats = make_stats()
    dl = Downloader(job_id, policy, response_send, request_recv, semaphore,
        rate_limiter_reset, stats)
    nursery.start_soon(dl.run)
    request = make_request('http://{}:{}/foo'.format(addr, port),
        form_data={'user': 'john'})
    await request_send.send(request)
    response = await response_recv.receive()
    nursery.cancel_scope.cancel()
    assert response.status_code == 200
    assert response.content_type == 'text/html'
    assert response.body.startswith(b'<html>')
    assert stats['item_count'] == 1
    assert stats['http_success_count'] == 1
    assert stats['http_status_counts'][200] == 1 

def test_http_post(nursery):
    # Create a server:
    async def handler(stream):
        request = b''
        while True:
            request = await stream.receive_some(4096)
            if request.endswith(b'\r\n\r\n'):
                break
        assert request.startswith(b'POST /foo HTTP/1.1\r\n')
        await stream.send_all(
                b'HTTP/1.1 200 OK\r\n'
                b'Content-type: text/html\r\n'
                b'\r\n'
                b'<html><head><title>Unit Test</title></head>\r\n'
                b'<body><h1>This is a unit test</h1></body></html>\r\n'
                b'\r\n'
            )
        await stream.aclose()
    serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1')
    http_server = await nursery.start(serve_tcp)
    addr, port = http_server[0].socket.getsockname()

    # Run the test:
    job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb'
    policy = make_policy()
    request_send, request_recv = trio.open_memory_channel(0)
    response_send, response_recv = trio.open_memory_channel(0)
    semaphore = trio.Semaphore(1)
    rate_limiter_reset = Mock()
    rate_limiter_reset.send = AsyncMock()
    stats = make_stats()
    dl = Downloader(job_id, policy, response_send, request_recv, semaphore,
        rate_limiter_reset, stats)
    nursery.start_soon(dl.run)
    request = make_request('http://{}:{}/foo'.format(addr, port), method='POST',
        form_data={'user': 'john'})
    await request_send.send(request)
    response = await response_recv.receive()
    nursery.cancel_scope.cancel()
    assert response.status_code == 200
    assert response.content_type == 'text/html'
    assert response.body.startswith(b'<html>')
    assert stats['item_count'] == 1
    assert stats['http_success_count'] == 1
    assert stats['http_status_counts'][200] == 1 

def test_http_proxy_get(asyncio_loop, nursery):
    # Create a server:
    async def handler(stream):
        request = await stream.receive_some(4096)
        assert request.startswith(b'GET http://test.example/foo HTTP/1.1')
        assert request.endswith(b'\r\n\r\n')
        await stream.send_all(
                b'HTTP/1.1 200 OK\n'
                b'Content-type: text/html\n'
                b'\n'
                b'<html><head><title>Unit Test</title></head>\n'
                b'<body><h1>This is a unit test</h1></body></html>\n'
            )
        await stream.aclose()
    serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1')
    http_proxy = await nursery.start(serve_tcp)
    addr, port = http_proxy[0].socket.getsockname()

    # Run the test:
    job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb'
    policy = make_policy(proxy=[{
        'proxy_url': 'http://{}:{}'.format(addr, port),
    }])
    request_send, request_recv = trio.open_memory_channel(0)
    response_send, response_recv = trio.open_memory_channel(0)
    semaphore = trio.Semaphore(1)
    rate_limiter_reset = Mock()
    rate_limiter_reset.send = AsyncMock()
    stats = make_stats()
    dl = Downloader(job_id, policy, response_send, request_recv, semaphore,
        rate_limiter_reset, stats)
    nursery.start_soon(dl.run)
    request = make_request('http://test.example/foo', policy=policy)
    await request_send.send(request)
    response = await response_recv.receive()
    nursery.cancel_scope.cancel()
    assert response.status_code == 200
    assert response.content_type == 'text/html'
    assert response.body.startswith(b'<html>')
    assert stats['item_count'] == 1
    assert stats['http_success_count'] == 1
    assert stats['http_status_counts'][200] == 1 

def test_socks_proxy_get(asyncio_loop, nursery):
    # Create a server:
    async def handler(stream):
        # Negotiate SOCKS authentication (no auth)
        request = await stream.receive_some(4096)
        assert request.startswith(b'\x05')
        await stream.send_all(b'\x05\x00')
        # Negotiate SOCKS command.
        request = await stream.receive_some(4096)
        assert request == b'\x05\x01\x00\x01\x7f\x00\x00\x01\x00\x50'
        await stream.send_all(b'\x05\x00\x00\x01\x7f\x00\x00\x01\x00\x50')
        # Get HTTP request and send response.
        request = await stream.receive_some(4096)
        assert request.startswith(b'GET /foo HTTP/1.1')
        assert request.endswith(b'\r\n\r\n')
        await stream.send_all(
                b'HTTP/1.1 200 OK\r\n'
                b'Content-type: text/html\r\n'
                b'\r\n'
                b'<html><head><title>Unit Test</title></head>\r\n'
                b'<body><h1>This is a unit test</h1></body></html>\r\n'
                b'\r\n'
            )
        await stream.aclose()
    serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1')
    socks_proxy = await nursery.start(serve_tcp)
    addr, port = socks_proxy[0].socket.getsockname()

    # Run the test:
    job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb'
    policy = make_policy(proxy=[{
        'proxy_url': 'socks5://{}:{}'.format(addr, port),
    }])
    request_send, request_recv = trio.open_memory_channel(0)
    response_send, response_recv = trio.open_memory_channel(0)
    semaphore = trio.Semaphore(1)
    rate_limiter_reset = Mock()
    rate_limiter_reset.send = AsyncMock()
    stats = make_stats()
    dl = Downloader(job_id, policy, response_send, request_recv, semaphore,
        rate_limiter_reset, stats)
    nursery.start_soon(dl.run)
    request = make_request('http://127.0.0.1/foo', policy=policy)
    await request_send.send(request)
    response = await response_recv.receive()
    nursery.cancel_scope.cancel()
    assert response.status_code == 200
    assert response.content_type == 'text/html'
    assert response.body.startswith(b'<html>')
    assert stats['item_count'] == 1
    assert stats['http_success_count'] == 1
    assert stats['http_status_counts'][200] == 1 

def _run_tcp_server(socket_file: Path, cmd_handler):
    async def _client_handler(stream):

        # General exception handling
        try:

            # Stream handling
            try:

                unpacker = Unpacker()
                async for raw in stream:
                    unpacker.feed(raw)
                    for cmd in unpacker:
                        cmd = cmd_req_serializer.load(cmd)
                        rep = await cmd_handler(cmd)
                        raw_rep = cmd_rep_serializer.dumps(rep)
                        logger.info("Command processed", cmd=cmd["cmd"], rep_status=rep["status"])
                        await stream.send_all(raw_rep)

            except SerdeError as exc:
                await stream.send_all(packb({"status": "invalid_format", "reason": str(exc)}))

            finally:
                await stream.aclose()

        except trio.BrokenResourceError:
            pass  # Peer has closed the connection while we were sending a response

        except Exception:
            logger.exception("Unexpected crash")

    try:
        async with trio.open_service_nursery() as nursery:
            listeners = await nursery.start(
                partial(trio.serve_tcp, _client_handler, 0, host="127.0.0.1")
            )
            port = listeners[0].socket.getsockname()[1]

            # Make sure the path exists and write the socket file
            socket_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
            socket_file.write_text(str(port))

            logger.info("IPC server ready", port=port)
            try:
                yield
            finally:
                nursery.cancel_scope.cancel()

    except OSError as exc:
        raise IPCServerError(f"Cannot start IPC server: {exc}") from exc