From 9d8e380d2e8267129de921b9b926257c36417cd2 Mon Sep 17 00:00:00 2001 From: Sean Quah <8349537+squahtx@users.noreply.github.com> Date: Wed, 11 May 2022 12:25:13 +0100 Subject: Respect the `@cancellable` flag for `RestServlet`s and `BaseFederationServlet`s (#12699) Both `RestServlet`s and `BaseFederationServlet`s register their handlers with `HttpServer.register_paths` / `JsonResource.register_paths`. Update `JsonResource` to respect the `@cancellable` flag on handlers registered in this way. Although `ReplicationEndpoint` also registers itself using `register_paths`, it does not pass the handler method that would have the `@cancellable` flag directly, and so needs separate handling. Signed-off-by: Sean Quah --- tests/federation/transport/server/__init__.py | 13 +++ tests/federation/transport/server/test__base.py | 112 ++++++++++++++++++++++++ tests/http/test_servlet.py | 60 ++++++++++++- tests/unittest.py | 2 +- 4 files changed, 185 insertions(+), 2 deletions(-) create mode 100644 tests/federation/transport/server/__init__.py create mode 100644 tests/federation/transport/server/test__base.py (limited to 'tests') diff --git a/tests/federation/transport/server/__init__.py b/tests/federation/transport/server/__init__.py new file mode 100644 index 0000000000..3a5f22c022 --- /dev/null +++ b/tests/federation/transport/server/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The Matrix.org Foundation C.I.C. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/federation/transport/server/test__base.py b/tests/federation/transport/server/test__base.py new file mode 100644 index 0000000000..98a951f03e --- /dev/null +++ b/tests/federation/transport/server/test__base.py @@ -0,0 +1,112 @@ +# Copyright 2022 The Matrix.org Foundation C.I.C. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from http import HTTPStatus +from typing import Dict, List, Tuple + +from synapse.api.errors import Codes +from synapse.federation.transport.server import BaseFederationServlet +from synapse.federation.transport.server._base import Authenticator +from synapse.http.server import JsonResource, cancellable +from synapse.server import HomeServer +from synapse.types import JsonDict +from synapse.util.ratelimitutils import FederationRateLimiter + +from tests import unittest +from tests.http.server._base import EndpointCancellationTestHelperMixin + + +class CancellableFederationServlet(BaseFederationServlet): + PATH = "/sleep" + + def __init__( + self, + hs: HomeServer, + authenticator: Authenticator, + ratelimiter: FederationRateLimiter, + server_name: str, + ): + super().__init__(hs, authenticator, ratelimiter, server_name) + self.clock = hs.get_clock() + + @cancellable + async def on_GET( + self, origin: str, content: None, query: Dict[bytes, List[bytes]] + ) -> Tuple[int, JsonDict]: + await self.clock.sleep(1.0) + return HTTPStatus.OK, {"result": True} + + async def on_POST( + self, origin: str, content: JsonDict, query: Dict[bytes, List[bytes]] + ) -> Tuple[int, JsonDict]: + await self.clock.sleep(1.0) + return HTTPStatus.OK, {"result": True} + + +class BaseFederationServletCancellationTests( + unittest.FederatingHomeserverTestCase, EndpointCancellationTestHelperMixin +): + """Tests for `BaseFederationServlet` cancellation.""" + + path = f"{CancellableFederationServlet.PREFIX}{CancellableFederationServlet.PATH}" + + def create_test_resource(self): + """Overrides `HomeserverTestCase.create_test_resource`.""" + resource = JsonResource(self.hs) + + CancellableFederationServlet( + hs=self.hs, + authenticator=Authenticator(self.hs), + ratelimiter=self.hs.get_federation_ratelimiter(), + server_name=self.hs.hostname, + ).register(resource) + + return resource + + def test_cancellable_disconnect(self) -> None: + """Test that handlers with the `@cancellable` flag can be cancelled.""" + channel = self.make_signed_federation_request( + "GET", self.path, await_result=False + ) + + # Advance past all the rate limiting logic. If we disconnect too early, the + # request won't be processed. + self.pump() + + self._test_disconnect( + self.reactor, + channel, + expect_cancellation=True, + expected_body={"error": "Request cancelled", "errcode": Codes.UNKNOWN}, + ) + + def test_uncancellable_disconnect(self) -> None: + """Test that handlers without the `@cancellable` flag cannot be cancelled.""" + channel = self.make_signed_federation_request( + "POST", + self.path, + content={}, + await_result=False, + ) + + # Advance past all the rate limiting logic. If we disconnect too early, the + # request won't be processed. + self.pump() + + self._test_disconnect( + self.reactor, + channel, + expect_cancellation=False, + expected_body={"result": True}, + ) diff --git a/tests/http/test_servlet.py b/tests/http/test_servlet.py index a80bfb9f4e..ad521525cf 100644 --- a/tests/http/test_servlet.py +++ b/tests/http/test_servlet.py @@ -12,16 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. import json +from http import HTTPStatus from io import BytesIO +from typing import Tuple from unittest.mock import Mock -from synapse.api.errors import SynapseError +from synapse.api.errors import Codes, SynapseError +from synapse.http.server import cancellable from synapse.http.servlet import ( + RestServlet, parse_json_object_from_request, parse_json_value_from_request, ) +from synapse.http.site import SynapseRequest +from synapse.rest.client._base import client_patterns +from synapse.server import HomeServer +from synapse.types import JsonDict from tests import unittest +from tests.http.server._base import EndpointCancellationTestHelperMixin def make_request(content): @@ -76,3 +85,52 @@ class TestServletUtils(unittest.TestCase): # Test not an object with self.assertRaises(SynapseError): parse_json_object_from_request(make_request(b'["foo"]')) + + +class CancellableRestServlet(RestServlet): + """A `RestServlet` with a mix of cancellable and uncancellable handlers.""" + + PATTERNS = client_patterns("/sleep$") + + def __init__(self, hs: HomeServer): + super().__init__() + self.clock = hs.get_clock() + + @cancellable + async def on_GET(self, request: SynapseRequest) -> Tuple[int, JsonDict]: + await self.clock.sleep(1.0) + return HTTPStatus.OK, {"result": True} + + async def on_POST(self, request: SynapseRequest) -> Tuple[int, JsonDict]: + await self.clock.sleep(1.0) + return HTTPStatus.OK, {"result": True} + + +class TestRestServletCancellation( + unittest.HomeserverTestCase, EndpointCancellationTestHelperMixin +): + """Tests for `RestServlet` cancellation.""" + + servlets = [ + lambda hs, http_server: CancellableRestServlet(hs).register(http_server) + ] + + def test_cancellable_disconnect(self) -> None: + """Test that handlers with the `@cancellable` flag can be cancelled.""" + channel = self.make_request("GET", "/sleep", await_result=False) + self._test_disconnect( + self.reactor, + channel, + expect_cancellation=True, + expected_body={"error": "Request cancelled", "errcode": Codes.UNKNOWN}, + ) + + def test_uncancellable_disconnect(self) -> None: + """Test that handlers without the `@cancellable` flag cannot be cancelled.""" + channel = self.make_request("POST", "/sleep", await_result=False) + self._test_disconnect( + self.reactor, + channel, + expect_cancellation=False, + expected_body={"result": True}, + ) diff --git a/tests/unittest.py b/tests/unittest.py index 9afa68c164..e7f255b4fa 100644 --- a/tests/unittest.py +++ b/tests/unittest.py @@ -831,7 +831,7 @@ class FederatingHomeserverTestCase(HomeserverTestCase): self.site, method=method, path=path, - content=content or "", + content=content if content is not None else "", shorthand=False, await_result=await_result, custom_headers=custom_headers, -- cgit 1.4.1