test_srv_resolver.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  1. # Copyright 2014-2016 OpenMarket Ltd
  2. # Copyright 2019 New Vector Ltd
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. from typing import Dict, Generator, List, Tuple, cast
  16. from unittest.mock import Mock
  17. from twisted.internet import defer
  18. from twisted.internet.defer import Deferred
  19. from twisted.internet.error import ConnectError
  20. from twisted.names import dns, error
  21. from synapse.http.federation.srv_resolver import Server, SrvResolver
  22. from synapse.logging.context import LoggingContext, current_context
  23. from tests import unittest
  24. from tests.utils import MockClock
  25. class SrvResolverTestCase(unittest.TestCase):
  26. def test_resolve(self) -> None:
  27. dns_client_mock = Mock()
  28. service_name = b"test_service.example.com"
  29. host_name = b"example.com"
  30. answer_srv = dns.RRHeader(
  31. type=dns.SRV, payload=dns.Record_SRV(target=host_name)
  32. )
  33. result_deferred: "Deferred[Tuple[List[dns.RRHeader], None, None]]" = Deferred()
  34. dns_client_mock.lookupService.return_value = result_deferred
  35. cache: Dict[bytes, List[Server]] = {}
  36. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  37. @defer.inlineCallbacks
  38. def do_lookup() -> Generator["Deferred[object]", object, List[Server]]:
  39. with LoggingContext("one") as ctx:
  40. resolve_d = resolver.resolve_service(service_name)
  41. result: List[Server]
  42. result = yield defer.ensureDeferred(resolve_d) # type: ignore[assignment]
  43. # should have restored our context
  44. self.assertIs(current_context(), ctx)
  45. return result
  46. test_d = do_lookup()
  47. self.assertNoResult(test_d)
  48. dns_client_mock.lookupService.assert_called_once_with(service_name)
  49. result_deferred.callback(([answer_srv], None, None))
  50. servers = self.successResultOf(test_d)
  51. self.assertEqual(len(servers), 1)
  52. self.assertEqual(servers, cache[service_name])
  53. self.assertEqual(servers[0].host, host_name)
  54. @defer.inlineCallbacks
  55. def test_from_cache_expired_and_dns_fail(
  56. self,
  57. ) -> Generator["Deferred[object]", object, None]:
  58. dns_client_mock = Mock()
  59. dns_client_mock.lookupService.return_value = defer.fail(error.DNSServerError())
  60. service_name = b"test_service.example.com"
  61. entry = Mock(spec_set=["expires", "priority", "weight"])
  62. entry.expires = 0
  63. entry.priority = 0
  64. entry.weight = 0
  65. cache = {service_name: [cast(Server, entry)]}
  66. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  67. servers: List[Server]
  68. servers = yield defer.ensureDeferred(
  69. resolver.resolve_service(service_name)
  70. ) # type: ignore[assignment]
  71. dns_client_mock.lookupService.assert_called_once_with(service_name)
  72. self.assertEqual(len(servers), 1)
  73. self.assertEqual(servers, cache[service_name])
  74. @defer.inlineCallbacks
  75. def test_from_cache(self) -> Generator["Deferred[object]", object, None]:
  76. clock = MockClock()
  77. dns_client_mock = Mock(spec_set=["lookupService"])
  78. dns_client_mock.lookupService = Mock(spec_set=[])
  79. service_name = b"test_service.example.com"
  80. entry = Mock(spec_set=["expires", "priority", "weight"])
  81. entry.expires = 999999999
  82. entry.priority = 0
  83. entry.weight = 0
  84. cache = {service_name: [cast(Server, entry)]}
  85. resolver = SrvResolver(
  86. dns_client=dns_client_mock, cache=cache, get_time=clock.time
  87. )
  88. servers: List[Server]
  89. servers = yield defer.ensureDeferred(
  90. resolver.resolve_service(service_name)
  91. ) # type: ignore[assignment]
  92. self.assertFalse(dns_client_mock.lookupService.called)
  93. self.assertEqual(len(servers), 1)
  94. self.assertEqual(servers, cache[service_name])
  95. @defer.inlineCallbacks
  96. def test_empty_cache(self) -> Generator["Deferred[object]", object, None]:
  97. dns_client_mock = Mock()
  98. dns_client_mock.lookupService.return_value = defer.fail(error.DNSServerError())
  99. service_name = b"test_service.example.com"
  100. cache: Dict[bytes, List[Server]] = {}
  101. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  102. with self.assertRaises(error.DNSServerError):
  103. yield defer.ensureDeferred(resolver.resolve_service(service_name))
  104. @defer.inlineCallbacks
  105. def test_name_error(self) -> Generator["Deferred[object]", object, None]:
  106. dns_client_mock = Mock()
  107. dns_client_mock.lookupService.return_value = defer.fail(error.DNSNameError())
  108. service_name = b"test_service.example.com"
  109. cache: Dict[bytes, List[Server]] = {}
  110. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  111. servers: List[Server]
  112. servers = yield defer.ensureDeferred(
  113. resolver.resolve_service(service_name)
  114. ) # type: ignore[assignment]
  115. self.assertEqual(len(servers), 0)
  116. self.assertEqual(len(cache), 0)
  117. def test_disabled_service(self) -> None:
  118. """
  119. test the behaviour when there is a single record which is ".".
  120. """
  121. service_name = b"test_service.example.com"
  122. lookup_deferred: "Deferred[Tuple[List[dns.RRHeader], None, None]]" = Deferred()
  123. dns_client_mock = Mock()
  124. dns_client_mock.lookupService.return_value = lookup_deferred
  125. cache: Dict[bytes, List[Server]] = {}
  126. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  127. # Old versions of Twisted don't have an ensureDeferred in failureResultOf.
  128. resolve_d = defer.ensureDeferred(resolver.resolve_service(service_name))
  129. # returning a single "." should make the lookup fail with a ConenctError
  130. lookup_deferred.callback(
  131. (
  132. [dns.RRHeader(type=dns.SRV, payload=dns.Record_SRV(target=b"."))],
  133. None,
  134. None,
  135. )
  136. )
  137. self.failureResultOf(resolve_d, ConnectError)
  138. def test_non_srv_answer(self) -> None:
  139. """
  140. test the behaviour when the dns server gives us a spurious non-SRV response
  141. """
  142. service_name = b"test_service.example.com"
  143. lookup_deferred: "Deferred[Tuple[List[dns.RRHeader], None, None]]" = Deferred()
  144. dns_client_mock = Mock()
  145. dns_client_mock.lookupService.return_value = lookup_deferred
  146. cache: Dict[bytes, List[Server]] = {}
  147. resolver = SrvResolver(dns_client=dns_client_mock, cache=cache)
  148. # Old versions of Twisted don't have an ensureDeferred in successResultOf.
  149. resolve_d = defer.ensureDeferred(resolver.resolve_service(service_name))
  150. lookup_deferred.callback(
  151. (
  152. [
  153. dns.RRHeader(type=dns.A, payload=dns.Record_A()),
  154. dns.RRHeader(type=dns.SRV, payload=dns.Record_SRV(target=b"host")),
  155. ],
  156. None,
  157. None,
  158. )
  159. )
  160. servers = self.successResultOf(resolve_d)
  161. self.assertEqual(len(servers), 1)
  162. self.assertEqual(servers, cache[service_name])
  163. self.assertEqual(servers[0].host, b"host")