import time from typing import Callable from unittest import mock, skipUnless import DNS from django.conf import settings from django.core.exceptions import ValidationError from django.http import HttpResponse from django.utils.timezone import now as timezone_now from zerver.forms import email_is_not_mit_mailing_list from zerver.lib.rate_limiter import ( RateLimitedIPAddr, RateLimitedUser, RateLimiterLockingException, add_ratelimit_rule, remove_ratelimit_rule, ) from zerver.lib.test_classes import ZulipTestCase from zerver.lib.zephyr import compute_mit_user_fullname from zerver.models import PushDeviceToken, UserProfile if settings.ZILENCER_ENABLED: from zilencer.models import RateLimitedRemoteZulipServer, RemoteZulipServer class MITNameTest(ZulipTestCase): def test_valid_hesiod(self) -> None: with mock.patch( "DNS.dnslookup", return_value=[ ["starnine:*:84233:101:Athena Consulting Exchange User,,,:/mit/starnine:/bin/bash"] ], ): self.assertEqual( compute_mit_user_fullname(self.mit_email("starnine")), "Athena Consulting Exchange User", ) with mock.patch( "DNS.dnslookup", return_value=[["sipbexch:*:87824:101:Exch Sipb,,,:/mit/sipbexch:/bin/athena/bash"]], ): self.assertEqual(compute_mit_user_fullname("sipbexch@mit.edu"), "Exch Sipb") def test_invalid_hesiod(self) -> None: with mock.patch( "DNS.dnslookup", side_effect=DNS.Base.ServerError("DNS query status: NXDOMAIN", 3) ): self.assertEqual(compute_mit_user_fullname("1234567890@mit.edu"), "1234567890@mit.edu") with mock.patch( "DNS.dnslookup", side_effect=DNS.Base.ServerError("DNS query status: NXDOMAIN", 3) ): self.assertEqual(compute_mit_user_fullname("ec-discuss@mit.edu"), "ec-discuss@mit.edu") def test_mailinglist(self) -> None: with mock.patch( "DNS.dnslookup", side_effect=DNS.Base.ServerError("DNS query status: NXDOMAIN", 3) ): self.assertRaises(ValidationError, email_is_not_mit_mailing_list, "1234567890@mit.edu") with mock.patch( "DNS.dnslookup", side_effect=DNS.Base.ServerError("DNS query status: NXDOMAIN", 3) ): self.assertRaises(ValidationError, email_is_not_mit_mailing_list, "ec-discuss@mit.edu") def test_notmailinglist(self) -> None: with mock.patch("DNS.dnslookup", return_value=[["POP IMAP.EXCHANGE.MIT.EDU starnine"]]): email_is_not_mit_mailing_list("sipbexch@mit.edu") class RateLimitTests(ZulipTestCase): def setUp(self) -> None: super().setUp() # Some tests here can be somewhat timing-sensitive in a way # that can't be eliminated, e.g. due to testing things that rely # on Redis' internal timing mechanism which we can't mock. # The first API request when running a suite of tests is slow # and can take multiple seconds. This is not a problem when running # multiple tests, but if an individual, time-sensitive test from this class # is run, the first API request it makes taking a lot of time can throw things off # and cause the test to fail. Thus we do a dummy API request here to warm up # the system and allow the tests to assume their requests won't take multiple seconds. user = self.example_user("hamlet") self.api_get(user, "/api/v1/messages") settings.RATE_LIMITING = True add_ratelimit_rule(1, 5) def tearDown(self) -> None: settings.RATE_LIMITING = False remove_ratelimit_rule(1, 5) super().tearDown() def send_api_message(self, user: UserProfile, content: str) -> HttpResponse: return self.api_post( user, "/api/v1/messages", { "type": "stream", "to": "Verona", "client": "test suite", "content": content, "topic": "whatever", }, ) def send_unauthed_api_request(self) -> HttpResponse: result = self.client_get("/json/messages") # We're not making a correct request here, but rate-limiting is supposed # to happen before the request fails due to not being correctly made. Thus # we expect either an 400 error if the request is allowed by the rate limiter, # or 429 if we're above the limit. We don't expect to see other status codes here, # so we assert for safety. self.assertIn(result.status_code, [400, 429]) return result def test_headers(self) -> None: user = self.example_user("hamlet") RateLimitedUser(user).clear_history() result = self.send_api_message(user, "some stuff") self.assertTrue("X-RateLimit-Remaining" in result) self.assertTrue("X-RateLimit-Limit" in result) self.assertTrue("X-RateLimit-Reset" in result) def test_ratelimit_decrease(self) -> None: user = self.example_user("hamlet") RateLimitedUser(user).clear_history() result = self.send_api_message(user, "some stuff") limit = int(result["X-RateLimit-Remaining"]) result = self.send_api_message(user, "some stuff 2") newlimit = int(result["X-RateLimit-Remaining"]) self.assertEqual(limit, newlimit + 1) def do_test_hit_ratelimits(self, request_func: Callable[[], HttpResponse]) -> HttpResponse: start_time = time.time() for i in range(6): with mock.patch("time.time", return_value=(start_time + i * 0.1)): result = request_func() self.assertEqual(result.status_code, 429) json = result.json() self.assertEqual(json.get("result"), "error") self.assertIn("API usage exceeded rate limit", json.get("msg")) self.assertEqual(json.get("retry-after"), 0.5) self.assertTrue("Retry-After" in result) self.assertEqual(result["Retry-After"], "0.5") # We simulate waiting a second here, rather than force-clearing our history, # to make sure the rate-limiting code automatically forgives a user # after some time has passed. with mock.patch("time.time", return_value=(start_time + 1.01)): result = request_func() self.assertNotEqual(result, 429) def test_hit_ratelimits_as_user(self) -> None: user = self.example_user("cordelia") RateLimitedUser(user).clear_history() self.do_test_hit_ratelimits(lambda: self.send_api_message(user, "some stuff")) def test_hit_ratelimits_as_ip(self) -> None: add_ratelimit_rule(1, 5, domain="api_by_ip") try: RateLimitedIPAddr("127.0.0.1").clear_history() self.do_test_hit_ratelimits(self.send_unauthed_api_request) finally: # We need this in a finally block to ensure the test cleans up after itself # even in case of failure, to avoid polluting the rules state. remove_ratelimit_rule(1, 5, domain="api_by_ip") @skipUnless(settings.ZILENCER_ENABLED, "requires zilencer") def test_hit_ratelimits_as_remote_server(self) -> None: add_ratelimit_rule(1, 5, domain="api_by_remote_server") server_uuid = "1234-abcd" server = RemoteZulipServer( uuid=server_uuid, api_key="magic_secret_api_key", hostname="demo.example.com", last_updated=timezone_now(), ) server.save() endpoint = "/api/v1/remotes/push/register" payload = {"user_id": 10, "token": "111222", "token_kind": PushDeviceToken.GCM} try: # Remote servers can only make requests to the root subdomain. original_default_subdomain = self.DEFAULT_SUBDOMAIN self.DEFAULT_SUBDOMAIN = "" RateLimitedRemoteZulipServer(server).clear_history() with self.assertLogs("zerver.lib.rate_limiter", level="WARNING") as m: self.do_test_hit_ratelimits(lambda: self.uuid_post(server_uuid, endpoint, payload)) self.assertEqual( m.output, [ "WARNING:zerver.lib.rate_limiter:Remote server exceeded rate limits on domain api_by_remote_server" ], ) finally: self.DEFAULT_SUBDOMAIN = original_default_subdomain remove_ratelimit_rule(1, 5, domain="api_by_remote_server") def test_hit_ratelimiterlockingexception(self) -> None: user = self.example_user("cordelia") RateLimitedUser(user).clear_history() with mock.patch( "zerver.lib.rate_limiter.RedisRateLimiterBackend.incr_ratelimit", side_effect=RateLimiterLockingException, ): with self.assertLogs("zerver.lib.rate_limiter", level="WARNING") as m: result = self.send_api_message(user, "some stuff") self.assertEqual(result.status_code, 429) self.assertEqual( m.output, [ "WARNING:zerver.lib.rate_limiter:Deadlock trying to incr_ratelimit for {}".format( f"RateLimitedUser:{user.id}:api_by_user" ) ], )