Open Source Monorail

appengine/monorail/framework/test/ratelimiter_test.py

Index: appengine/monorail/framework/test/ratelimiter_test.py
diff --git a/appengine/monorail/framework/test/ratelimiter_test.py b/appengine/monorail/framework/test/ratelimiter_test.py
new file mode 100644
index 0000000000000000000000000000000000000000..82537ccd4391c17e91671178aaca5c1949d49503
--- /dev/null
+++ b/appengine/monorail/framework/test/ratelimiter_test.py
@@ -0,0 +1,330 @@
+# Copyright 2016 The Chromium Authors. All rights reserved.
+# Use of this source code is govered by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+"""Unit tests for RateLimiter.
+"""
+import unittest
+
+from google.appengine.api import memcache
+from google.appengine.ext import testbed
+
+import mox
+import os
+import settings
+
+from framework import ratelimiter
+from services import service_manager
+from testing import fake
+from testing import testing_helpers
+
+
+class RateLimiterTest(unittest.TestCase):
+  def setUp(self):
+    settings.ratelimiting_enabled = True
+    self.testbed = testbed.Testbed()
+    self.testbed.activate()
+    self.testbed.init_memcache_stub()
+    self.testbed.init_user_stub()
+
+    self.mox = mox.Mox()
+    self.services = service_manager.Services(
+      config=fake.ConfigService(),
+      issue=fake.IssueService(),
+      user=fake.UserService(),
+      project=fake.ProjectService(),
+    )
+    self.project = self.services.project.TestAddProject('proj', project_id=987)
+
+    self.ratelimiter = ratelimiter.RateLimiter()
+    ratelimiter.COUNTRY_LIMITS = {}
+    os.environ['USER_EMAIL'] = ''
+    settings.ratelimiting_enabled = True
+    settings.ratelimiting_cost_enabled = True
+    ratelimiter.DEFAULT_LIMIT = 10
+
+  def tearDown(self):
+    self.testbed.deactivate()
+    self.mox.UnsetStubs()
+    self.mox.ResetAll()
+    # settings.ratelimiting_enabled = True
+
+  def testCheckStart_pass(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers['X-AppEngine-Country'] = 'US'
+    request.remote_addr = '192.168.1.0'
+    self.ratelimiter.CheckStart(request)
+    # Should not throw an exception.
+
+  def testCheckStart_fail(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers['X-AppEngine-Country'] = 'US'
+    request.remote_addr = '192.168.1.0'
+    now = 0.0
+    cachekeysets, _, _, _ = ratelimiter._CacheKeys(request, now)
+    values = [{key: ratelimiter.DEFAULT_LIMIT for key in cachekeys} for
+              cachekeys in cachekeysets]
+    for value in values:
+      memcache.add_multi(value)
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      self.ratelimiter.CheckStart(request, now)
+
+  def testCheckStart_expiredEntries(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers['X-AppEngine-Country'] = 'US'
+    request.remote_addr = '192.168.1.0'
+    now = 0.0
+    cachekeysets, _, _, _ = ratelimiter._CacheKeys(request, now)
+    values = [{key: ratelimiter.DEFAULT_LIMIT for key in cachekeys} for
+              cachekeys in cachekeysets]
+    for value in values:
+      memcache.add_multi(value)
+
+    now = now + 2 * ratelimiter.EXPIRE_AFTER_SECS
+    self.ratelimiter.CheckStart(request, now)
+    # Should not throw an exception.
+
+  def testCheckStart_repeatedCalls(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers['X-AppEngine-Country'] = 'US'
+    request.remote_addr = '192.168.1.0'
+    now = 0.0
+
+    # Call CheckStart once every minute.  Should be ok.
+    for _ in range(ratelimiter.N_MINUTES):
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 120.0
+
+    # Call CheckStart more than DEFAULT_LIMIT times in the same minute.
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      for _ in range(ratelimiter.DEFAULT_LIMIT + 2):
+        now = now + 0.001
+        self.ratelimiter.CheckStart(request, now)
+
+  def testCheckStart_differentIPs(self):
+    now = 0.0
+
+    ratelimiter.COUNTRY_LIMITS = {}
+    # Exceed DEFAULT_LIMIT calls, but vary remote_addr so different
+    # remote addresses aren't ratelimited together.
+    for m in range(ratelimiter.DEFAULT_LIMIT * 2):
+      request, _ = testing_helpers.GetRequestObjects(
+        project=self.project)
+      request.headers['X-AppEngine-Country'] = 'US'
+      request.remote_addr = '192.168.1.%d' % (m % 16)
+      ratelimiter._CacheKeys(request, now)
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 0.001
+
+    # Exceed the limit, but only for one IP address. The
+    # others should be fine.
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      for m in range(ratelimiter.DEFAULT_LIMIT):
+        request, _ = testing_helpers.GetRequestObjects(
+          project=self.project)
+        request.headers['X-AppEngine-Country'] = 'US'
+        request.remote_addr = '192.168.1.0'
+        ratelimiter._CacheKeys(request, now)
+        self.ratelimiter.CheckStart(request, now)
+        now = now + 0.001
+
+    # Now proceed to make requests for all of the other IP
+    # addresses besides .0.
+    for m in range(ratelimiter.DEFAULT_LIMIT * 2):
+      request, _ = testing_helpers.GetRequestObjects(
+        project=self.project)
+      request.headers['X-AppEngine-Country'] = 'US'
+      # Skip .0 since it's already exceeded the limit.
+      request.remote_addr = '192.168.1.%d' % (m + 1)
+      ratelimiter._CacheKeys(request, now)
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 0.001
+
+  def testCheckStart_sameIPDifferentUserIDs(self):
+    # Behind a NAT, e.g.
+    now = 0.0
+
+    # Exceed DEFAULT_LIMIT calls, but vary user_id so different
+    # users behind the same IP aren't ratelimited together.
+    for m in range(ratelimiter.DEFAULT_LIMIT * 2):
+      request, _ = testing_helpers.GetRequestObjects(
+        project=self.project)
+      request.remote_addr = '192.168.1.0'
+      os.environ['USER_EMAIL'] = '%s@example.com' % m
+      request.headers['X-AppEngine-Country'] = 'US'
+      ratelimiter._CacheKeys(request, now)
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 0.001
+
+    # Exceed the limit, but only for one userID+IP address. The
+    # others should be fine.
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      for m in range(ratelimiter.DEFAULT_LIMIT + 2):
+        request, _ = testing_helpers.GetRequestObjects(
+          project=self.project)
+        request.headers['X-AppEngine-Country'] = 'US'
+        request.remote_addr = '192.168.1.0'
+        os.environ['USER_EMAIL'] = '42@example.com'
+        ratelimiter._CacheKeys(request, now)
+        self.ratelimiter.CheckStart(request, now)
+        now = now + 0.001
+
+    # Now proceed to make requests for other user IDs
+    # besides 42.
+    for m in range(ratelimiter.DEFAULT_LIMIT * 2):
+      request, _ = testing_helpers.GetRequestObjects(
+        project=self.project)
+      request.headers['X-AppEngine-Country'] = 'US'
+      # Skip .0 since it's already exceeded the limit.
+      request.remote_addr = '192.168.1.0'
+      os.environ['USER_EMAIL'] = '%s@example.com' % (43 + m)
+      ratelimiter._CacheKeys(request, now)
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 0.001
+
+  def testCheckStart_ratelimitingDisabled(self):
+    settings.ratelimiting_enabled = False
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers['X-AppEngine-Country'] = 'US'
+    request.remote_addr = '192.168.1.0'
+    now = 0.0
+
+    # Call CheckStart a lot.  Should be ok.
+    for _ in range(ratelimiter.DEFAULT_LIMIT):
+      self.ratelimiter.CheckStart(request, now)
+      now = now + 0.001
+
+  def testCheckStart_perCountryLoggedOutLimit(self):
+    ratelimiter.COUNTRY_LIMITS['US'] = 10
+
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers[ratelimiter.COUNTRY_HEADER] = 'US'
+    request.remote_addr = '192.168.1.1'
+    now = 0.0
+
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      for m in range(ratelimiter.DEFAULT_LIMIT + 2):
+        self.ratelimiter.CheckStart(request, now)
+        # Vary remote address to make sure the limit covers
+        # the whole country, regardless of IP.
+        request.remote_addr = '192.168.1.%d' % m
+        now = now + 0.001
+
+    # CheckStart for a country that isn't covered by a country-specific limit.
+    request.headers['X-AppEngine-Country'] = 'UK'
+    for m in range(11):
+      self.ratelimiter.CheckStart(request, now)
+      # Vary remote address to make sure the limit covers
+      # the whole country, regardless of IP.
+      request.remote_addr = '192.168.1.%d' % m
+      now = now + 0.001
+
+    # And regular rate limits work per-IP.
+    request.remote_addr = '192.168.1.1'
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      for m in range(ratelimiter.DEFAULT_LIMIT):
+        self.ratelimiter.CheckStart(request, now)
+        # Vary remote address to make sure the limit covers
+        # the whole country, regardless of IP.
+        now = now + 0.001
+
+  def testCheckEnd_overCostThresh(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers[ratelimiter.COUNTRY_HEADER] = 'US'
+    request.remote_addr = '192.168.1.1'
+    start_time = 0.0
+
+    # Send some requests, all under the limit.
+    for _ in range(ratelimiter.DEFAULT_LIMIT-1):
+      start_time = start_time + 0.001
+      self.ratelimiter.CheckStart(request, start_time)
+      now = start_time + 0.010
+      self.ratelimiter.CheckEnd(request, now, start_time)
+
+    # Now issue some more request, this time taking long
+    # enough to get the cost threshold penalty.
+    # Fast forward enough to impact a later bucket than the
+    # previous requests.
+    start_time = now + 120.0
+    self.ratelimiter.CheckStart(request, start_time)
+
+    # Take longer than the threshold to process the request.
+    now = start_time + (settings.ratelimiting_cost_thresh_ms + 1) / 1000
+
+    # The request finished, taking longer than the cost
+    # threshold.
+    self.ratelimiter.CheckEnd(request, now, start_time)
+
+    with self.assertRaises(ratelimiter.RateLimitExceeded):
+      # One more request after the expensive query should
+      # throw an excpetion.
+      self.ratelimiter.CheckStart(request, start_time)
+
+  def testCheckEnd_overCostThreshButDisabled(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers[ratelimiter.COUNTRY_HEADER] = 'US'
+    request.remote_addr = '192.168.1.1'
+    start_time = 0.0
+    settings.ratelimiting_cost_enabled = False
+
+    # Send some requests, all under the limit.
+    for _ in range(ratelimiter.DEFAULT_LIMIT-1):
+      start_time = start_time + 0.001
+      self.ratelimiter.CheckStart(request, start_time)
+      now = start_time + 0.010
+      self.ratelimiter.CheckEnd(request, now, start_time)
+
+    # Now issue some more request, this time taking long
+    # enough to get the cost threshold penalty.
+    # Fast forward enough to impact a later bucket than the
+    # previous requests.
+    start_time = now + 120.0
+    self.ratelimiter.CheckStart(request, start_time)
+
+    # Take longer than the threshold to process the request.
+    now = start_time + (settings.ratelimiting_cost_thresh_ms + 10)/1000
+
+    # The request finished, taking longer than the cost
+    # threshold.
+    self.ratelimiter.CheckEnd(request, now, start_time)
+
+    # One more request after the expensive query should
+    # throw an excpetion, but cost thresholds are disabled.
+    self.ratelimiter.CheckStart(request, start_time)
+
+  def testChekcEnd_underCostThresh(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers[ratelimiter.COUNTRY_HEADER] = 'asdasd'
+    request.remote_addr = '192.168.1.1'
+    start_time = 0.0
+
+    # Send some requests, all under the limit.
+    for _ in range(ratelimiter.DEFAULT_LIMIT):
+      self.ratelimiter.CheckStart(request, start_time)
+      now = start_time + 0.010
+      self.ratelimiter.CheckEnd(request, now, start_time)
+      start_time = now + 0.010
+
+  def testChekcEnd_underCostThresh(self):
+    request, _ = testing_helpers.GetRequestObjects(
+      project=self.project)
+    request.headers[ratelimiter.COUNTRY_HEADER] = 'asdasd'
+    request.remote_addr = '192.168.1.1'
+    start_time = 0.0
+
+    # Send some requests, all under the limit.
+    for _ in range(ratelimiter.DEFAULT_LIMIT):
+      self.ratelimiter.CheckStart(request, start_time)
+      now = start_time + 0.01
+      self.ratelimiter.CheckEnd(request, now, start_time)
+      start_time = now + 0.01