synapse/tests/synapse_rust/test_http_client.py

# This file is licensed under the Affero General Public License (AGPL) version 3.
#
# Copyright (C) 2025 New Vector, Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# See the GNU Affero General Public License for more details:
# <https://www.gnu.org/licenses/agpl-3.0.html>.

import json
import logging
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any, TypeVar

from twisted.internet.testing import MemoryReactor

from synapse.logging.context import (
    LoggingContext,
    PreserveLoggingContext,
    _Sentinel,
    current_context,
    run_in_background,
)
from synapse.server import HomeServer
from synapse.synapse_rust.http_client import HttpClient
from synapse.util.clock import Clock
from synapse.util.json import json_decoder

from tests.unittest import HomeserverTestCase

logger = logging.getLogger(__name__)

T = TypeVar("T")


class StubRequestHandler(BaseHTTPRequestHandler):
    server: "StubServer"

    def do_GET(self) -> None:
        self.server.calls += 1

        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(json.dumps({"ok": True}).encode("utf-8"))

    def log_message(self, format: str, *args: Any) -> None:
        # Don't log anything; by default, the server logs to stderr
        pass


class StubServer(HTTPServer):
    """A stub HTTP server that we can send requests to for testing.

    This opens a real HTTP server on a random port, on a separate thread.
    """

    calls: int = 0
    """How many times has the endpoint been requested."""

    _thread: threading.Thread

    def __init__(self) -> None:
        super().__init__(("127.0.0.1", 0), StubRequestHandler)

        self._thread = threading.Thread(
            target=self.serve_forever,
            name="StubServer",
            kwargs={"poll_interval": 0.01},
            daemon=True,
        )
        self._thread.start()

    def shutdown(self) -> None:
        super().shutdown()
        self._thread.join()

    @property
    def endpoint(self) -> str:
        return f"http://127.0.0.1:{self.server_port}/"


class HttpClientTestCase(HomeserverTestCase):
    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        self._http_client = hs.get_proxied_http_client()
        self._rust_http_client = HttpClient(
            runtime=hs.get_rust_runtime(),
            user_agent=self._http_client.user_agent.decode("utf8"),
        )

        self.server = StubServer()

    def tearDown(self) -> None:
        # MemoryReactor doesn't trigger the shutdown phases, and we want the
        # Tokio thread pool to be stopped
        # XXX: This logic should probably get moved somewhere else
        shutdown_triggers = self.reactor.triggers.get("shutdown", {})
        for phase in ["before", "during", "after"]:
            triggers = shutdown_triggers.get(phase, [])
            for callbable, args, kwargs in triggers:
                callbable(*args, **kwargs)

    def _check_current_logcontext(self, expected_logcontext_string: str) -> None:
        context = current_context()
        assert isinstance(context, LoggingContext) or isinstance(context, _Sentinel), (
            f"Expected LoggingContext({expected_logcontext_string}) but saw {context}"
        )
        self.assertEqual(
            str(context),
            expected_logcontext_string,
            f"Expected LoggingContext({expected_logcontext_string}) but saw {context}",
        )

    def test_request_response(self) -> None:
        """
        Test to make sure we can make a basic request and get the expected
        response.
        """

        async def do_request() -> None:
            resp_body = await self._rust_http_client.get(
                url=self.server.endpoint,
                response_limit=1 * 1024 * 1024,
            )
            raw_response = json_decoder.decode(resp_body.decode("utf-8"))
            self.assertEqual(raw_response, {"ok": True})

        self.get_success(do_request())
        self.assertEqual(self.server.calls, 1)

    def test_request_response_limit_exceeded(self) -> None:
        """
        Test to make sure we handle the response limit being exceeded
        """

        async def do_request() -> None:
            await self._rust_http_client.get(
                url=self.server.endpoint,
                # Small limit so we hit the limit
                response_limit=1,
            )

        self.get_failure(
            do_request(),
            RuntimeError,
        )
        self.assertEqual(self.server.calls, 1)

    async def test_logging_context(self) -> None:
        """
        Test to make sure the `LoggingContext` (logcontext) is handled correctly
        when making requests.
        """
        # Sanity check that we start in the sentinel context
        self._check_current_logcontext("sentinel")

        callback_finished = False

        async def do_request() -> None:
            nonlocal callback_finished
            try:
                # Should have the same logcontext as the caller
                self._check_current_logcontext("foo")

                with LoggingContext(name="competing", server_name="test_server"):
                    # Make the actual request
                    await self._rust_http_client.get(
                        url=self.server.endpoint,
                        response_limit=1 * 1024 * 1024,
                    )
                    self._check_current_logcontext("competing")

                # Back to the caller's context outside of the `LoggingContext` block
                self._check_current_logcontext("foo")
            finally:
                # When exceptions happen, we still want to mark the callback as finished
                # so that the test can complete and we see the underlying error.
                callback_finished = True

        with LoggingContext(name="foo", server_name="test_server"):
            # Fire off the function, but don't wait on it.
            run_in_background(do_request)

            # Now wait for the function under test to have run
            with PreserveLoggingContext():
                while not callback_finished:
                    # Allow the async Rust to run
                    #
                    # Suspend execution of this thread to allow other the Tokio thread
                    # pool to do work.
                    time.sleep(0)
                    # Advance the Twisted reactor and run any scheduled callbacks
                    #
                    # In terms of other threads, they may have scheduled something on the
                    # reactor to run (like `reactor.callFromThread(...)`)
                    self.reactor.advance(0)

            # check that the logcontext is left in a sane state.
            self._check_current_logcontext("foo")

        self.assertTrue(
            callback_finished,
            "Callback never finished which means the test probably didn't wait long enough",
        )

        # Back to the sentinel context
        self._check_current_logcontext("sentinel")