1
0

test_ratelimiting.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390
  1. from synapse.api.ratelimiting import LimitExceededError, Ratelimiter
  2. from synapse.appservice import ApplicationService
  3. from synapse.types import create_requester
  4. from tests import unittest
  5. class TestRatelimiter(unittest.HomeserverTestCase):
  6. def test_allowed_via_can_do_action(self):
  7. limiter = Ratelimiter(
  8. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  9. )
  10. allowed, time_allowed = self.get_success_or_raise(
  11. limiter.can_do_action(None, key="test_id", _time_now_s=0)
  12. )
  13. self.assertTrue(allowed)
  14. self.assertEqual(10.0, time_allowed)
  15. allowed, time_allowed = self.get_success_or_raise(
  16. limiter.can_do_action(None, key="test_id", _time_now_s=5)
  17. )
  18. self.assertFalse(allowed)
  19. self.assertEqual(10.0, time_allowed)
  20. allowed, time_allowed = self.get_success_or_raise(
  21. limiter.can_do_action(None, key="test_id", _time_now_s=10)
  22. )
  23. self.assertTrue(allowed)
  24. self.assertEqual(20.0, time_allowed)
  25. def test_allowed_appservice_ratelimited_via_can_requester_do_action(self):
  26. appservice = ApplicationService(
  27. None,
  28. id="foo",
  29. rate_limited=True,
  30. sender="@as:example.com",
  31. )
  32. as_requester = create_requester("@user:example.com", app_service=appservice)
  33. limiter = Ratelimiter(
  34. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  35. )
  36. allowed, time_allowed = self.get_success_or_raise(
  37. limiter.can_do_action(as_requester, _time_now_s=0)
  38. )
  39. self.assertTrue(allowed)
  40. self.assertEqual(10.0, time_allowed)
  41. allowed, time_allowed = self.get_success_or_raise(
  42. limiter.can_do_action(as_requester, _time_now_s=5)
  43. )
  44. self.assertFalse(allowed)
  45. self.assertEqual(10.0, time_allowed)
  46. allowed, time_allowed = self.get_success_or_raise(
  47. limiter.can_do_action(as_requester, _time_now_s=10)
  48. )
  49. self.assertTrue(allowed)
  50. self.assertEqual(20.0, time_allowed)
  51. def test_allowed_appservice_via_can_requester_do_action(self):
  52. appservice = ApplicationService(
  53. None,
  54. id="foo",
  55. rate_limited=False,
  56. sender="@as:example.com",
  57. )
  58. as_requester = create_requester("@user:example.com", app_service=appservice)
  59. limiter = Ratelimiter(
  60. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  61. )
  62. allowed, time_allowed = self.get_success_or_raise(
  63. limiter.can_do_action(as_requester, _time_now_s=0)
  64. )
  65. self.assertTrue(allowed)
  66. self.assertEqual(-1, time_allowed)
  67. allowed, time_allowed = self.get_success_or_raise(
  68. limiter.can_do_action(as_requester, _time_now_s=5)
  69. )
  70. self.assertTrue(allowed)
  71. self.assertEqual(-1, time_allowed)
  72. allowed, time_allowed = self.get_success_or_raise(
  73. limiter.can_do_action(as_requester, _time_now_s=10)
  74. )
  75. self.assertTrue(allowed)
  76. self.assertEqual(-1, time_allowed)
  77. def test_allowed_via_ratelimit(self):
  78. limiter = Ratelimiter(
  79. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  80. )
  81. # Shouldn't raise
  82. self.get_success_or_raise(limiter.ratelimit(None, key="test_id", _time_now_s=0))
  83. # Should raise
  84. with self.assertRaises(LimitExceededError) as context:
  85. self.get_success_or_raise(
  86. limiter.ratelimit(None, key="test_id", _time_now_s=5)
  87. )
  88. self.assertEqual(context.exception.retry_after_ms, 5000)
  89. # Shouldn't raise
  90. self.get_success_or_raise(
  91. limiter.ratelimit(None, key="test_id", _time_now_s=10)
  92. )
  93. def test_allowed_via_can_do_action_and_overriding_parameters(self):
  94. """Test that we can override options of can_do_action that would otherwise fail
  95. an action
  96. """
  97. # Create a Ratelimiter with a very low allowed rate_hz and burst_count
  98. limiter = Ratelimiter(
  99. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  100. )
  101. # First attempt should be allowed
  102. allowed, time_allowed = self.get_success_or_raise(
  103. limiter.can_do_action(
  104. None,
  105. ("test_id",),
  106. _time_now_s=0,
  107. )
  108. )
  109. self.assertTrue(allowed)
  110. self.assertEqual(10.0, time_allowed)
  111. # Second attempt, 1s later, will fail
  112. allowed, time_allowed = self.get_success_or_raise(
  113. limiter.can_do_action(
  114. None,
  115. ("test_id",),
  116. _time_now_s=1,
  117. )
  118. )
  119. self.assertFalse(allowed)
  120. self.assertEqual(10.0, time_allowed)
  121. # But, if we allow 10 actions/sec for this request, we should be allowed
  122. # to continue.
  123. allowed, time_allowed = self.get_success_or_raise(
  124. limiter.can_do_action(None, ("test_id",), _time_now_s=1, rate_hz=10.0)
  125. )
  126. self.assertTrue(allowed)
  127. self.assertEqual(1.1, time_allowed)
  128. # Similarly if we allow a burst of 10 actions
  129. allowed, time_allowed = self.get_success_or_raise(
  130. limiter.can_do_action(None, ("test_id",), _time_now_s=1, burst_count=10)
  131. )
  132. self.assertTrue(allowed)
  133. self.assertEqual(1.0, time_allowed)
  134. def test_allowed_via_ratelimit_and_overriding_parameters(self):
  135. """Test that we can override options of the ratelimit method that would otherwise
  136. fail an action
  137. """
  138. # Create a Ratelimiter with a very low allowed rate_hz and burst_count
  139. limiter = Ratelimiter(
  140. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  141. )
  142. # First attempt should be allowed
  143. self.get_success_or_raise(
  144. limiter.ratelimit(None, key=("test_id",), _time_now_s=0)
  145. )
  146. # Second attempt, 1s later, will fail
  147. with self.assertRaises(LimitExceededError) as context:
  148. self.get_success_or_raise(
  149. limiter.ratelimit(None, key=("test_id",), _time_now_s=1)
  150. )
  151. self.assertEqual(context.exception.retry_after_ms, 9000)
  152. # But, if we allow 10 actions/sec for this request, we should be allowed
  153. # to continue.
  154. self.get_success_or_raise(
  155. limiter.ratelimit(None, key=("test_id",), _time_now_s=1, rate_hz=10.0)
  156. )
  157. # Similarly if we allow a burst of 10 actions
  158. self.get_success_or_raise(
  159. limiter.ratelimit(None, key=("test_id",), _time_now_s=1, burst_count=10)
  160. )
  161. def test_pruning(self):
  162. limiter = Ratelimiter(
  163. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=1
  164. )
  165. self.get_success_or_raise(
  166. limiter.can_do_action(None, key="test_id_1", _time_now_s=0)
  167. )
  168. self.assertIn("test_id_1", limiter.actions)
  169. self.get_success_or_raise(
  170. limiter.can_do_action(None, key="test_id_2", _time_now_s=10)
  171. )
  172. self.assertNotIn("test_id_1", limiter.actions)
  173. def test_db_user_override(self):
  174. """Test that users that have ratelimiting disabled in the DB aren't
  175. ratelimited.
  176. """
  177. store = self.hs.get_datastores().main
  178. user_id = "@user:test"
  179. requester = create_requester(user_id)
  180. self.get_success(
  181. store.db_pool.simple_insert(
  182. table="ratelimit_override",
  183. values={
  184. "user_id": user_id,
  185. "messages_per_second": None,
  186. "burst_count": None,
  187. },
  188. desc="test_db_user_override",
  189. )
  190. )
  191. limiter = Ratelimiter(store=store, clock=None, rate_hz=0.1, burst_count=1)
  192. # Shouldn't raise
  193. for _ in range(20):
  194. self.get_success_or_raise(limiter.ratelimit(requester, _time_now_s=0))
  195. def test_multiple_actions(self):
  196. limiter = Ratelimiter(
  197. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
  198. )
  199. # Test that 4 actions aren't allowed with a maximum burst of 3.
  200. allowed, time_allowed = self.get_success_or_raise(
  201. limiter.can_do_action(None, key="test_id", n_actions=4, _time_now_s=0)
  202. )
  203. self.assertFalse(allowed)
  204. # Test that 3 actions are allowed with a maximum burst of 3.
  205. allowed, time_allowed = self.get_success_or_raise(
  206. limiter.can_do_action(None, key="test_id", n_actions=3, _time_now_s=0)
  207. )
  208. self.assertTrue(allowed)
  209. self.assertEqual(10.0, time_allowed)
  210. # Test that, after doing these 3 actions, we can't do any more actions without
  211. # waiting.
  212. allowed, time_allowed = self.get_success_or_raise(
  213. limiter.can_do_action(None, key="test_id", n_actions=1, _time_now_s=0)
  214. )
  215. self.assertFalse(allowed)
  216. self.assertEqual(10.0, time_allowed)
  217. # Test that after waiting we would be able to do only 1 action.
  218. # Note that we don't actually do it (update=False) here.
  219. allowed, time_allowed = self.get_success_or_raise(
  220. limiter.can_do_action(
  221. None,
  222. key="test_id",
  223. update=False,
  224. n_actions=1,
  225. _time_now_s=10,
  226. )
  227. )
  228. self.assertTrue(allowed)
  229. # We would be able to do the 5th action at t=20.
  230. self.assertEqual(20.0, time_allowed)
  231. # Attempt (but fail) to perform TWO actions at t=10.
  232. # Those would be the 4th and 5th actions.
  233. allowed, time_allowed = self.get_success_or_raise(
  234. limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=10)
  235. )
  236. self.assertFalse(allowed)
  237. # The returned time allowed for the next action is now even though we weren't
  238. # allowed to perform the action because whilst we don't allow 2 actions,
  239. # we could still do 1.
  240. self.assertEqual(10.0, time_allowed)
  241. # Test that after waiting until t=20, we can do perform 2 actions.
  242. # These are the 4th and 5th actions.
  243. allowed, time_allowed = self.get_success_or_raise(
  244. limiter.can_do_action(None, key="test_id", n_actions=2, _time_now_s=20)
  245. )
  246. self.assertTrue(allowed)
  247. # We would be able to do the 6th action at t=30.
  248. self.assertEqual(30.0, time_allowed)
  249. def test_rate_limit_burst_only_given_once(self) -> None:
  250. """
  251. Regression test against a bug that meant that you could build up
  252. extra tokens by timing requests.
  253. """
  254. limiter = Ratelimiter(
  255. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
  256. )
  257. def consume_at(time: float) -> bool:
  258. success, _ = self.get_success_or_raise(
  259. limiter.can_do_action(requester=None, key="a", _time_now_s=time)
  260. )
  261. return success
  262. # Use all our 3 burst tokens
  263. self.assertTrue(consume_at(0.0))
  264. self.assertTrue(consume_at(0.1))
  265. self.assertTrue(consume_at(0.2))
  266. # Wait to recover 1 token (10 seconds at 0.1 Hz).
  267. self.assertTrue(consume_at(10.1))
  268. # Check that we get rate limited after using that token.
  269. self.assertFalse(consume_at(11.1))
  270. def test_record_action_which_doesnt_fill_bucket(self) -> None:
  271. limiter = Ratelimiter(
  272. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
  273. )
  274. # Observe two actions, leaving room in the bucket for one more.
  275. limiter.record_action(requester=None, key="a", n_actions=2, _time_now_s=0.0)
  276. # We should be able to take a new action now.
  277. success, _ = self.get_success_or_raise(
  278. limiter.can_do_action(requester=None, key="a", _time_now_s=0.0)
  279. )
  280. self.assertTrue(success)
  281. # ... but not two.
  282. success, _ = self.get_success_or_raise(
  283. limiter.can_do_action(requester=None, key="a", _time_now_s=0.0)
  284. )
  285. self.assertFalse(success)
  286. def test_record_action_which_fills_bucket(self) -> None:
  287. limiter = Ratelimiter(
  288. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
  289. )
  290. # Observe three actions, filling up the bucket.
  291. limiter.record_action(requester=None, key="a", n_actions=3, _time_now_s=0.0)
  292. # We should be unable to take a new action now.
  293. success, _ = self.get_success_or_raise(
  294. limiter.can_do_action(requester=None, key="a", _time_now_s=0.0)
  295. )
  296. self.assertFalse(success)
  297. # If we wait 10 seconds to leak a token, we should be able to take one action...
  298. success, _ = self.get_success_or_raise(
  299. limiter.can_do_action(requester=None, key="a", _time_now_s=10.0)
  300. )
  301. self.assertTrue(success)
  302. # ... but not two.
  303. success, _ = self.get_success_or_raise(
  304. limiter.can_do_action(requester=None, key="a", _time_now_s=10.0)
  305. )
  306. self.assertFalse(success)
  307. def test_record_action_which_overfills_bucket(self) -> None:
  308. limiter = Ratelimiter(
  309. store=self.hs.get_datastores().main, clock=None, rate_hz=0.1, burst_count=3
  310. )
  311. # Observe four actions, exceeding the bucket.
  312. limiter.record_action(requester=None, key="a", n_actions=4, _time_now_s=0.0)
  313. # We should be prevented from taking a new action now.
  314. success, _ = self.get_success_or_raise(
  315. limiter.can_do_action(requester=None, key="a", _time_now_s=0.0)
  316. )
  317. self.assertFalse(success)
  318. # If we wait 10 seconds to leak a token, we should be unable to take an action
  319. # because the bucket is still full.
  320. success, _ = self.get_success_or_raise(
  321. limiter.can_do_action(requester=None, key="a", _time_now_s=10.0)
  322. )
  323. self.assertFalse(success)
  324. # But after another 10 seconds we leak a second token, giving us room for
  325. # action.
  326. success, _ = self.get_success_or_raise(
  327. limiter.can_do_action(requester=None, key="a", _time_now_s=20.0)
  328. )
  329. self.assertTrue(success)