OLD | NEW |
(Empty) | |
| 1 # Copyright 2016 The Chromium Authors. All rights reserved. |
| 2 # Use of this source code is govered by a BSD-style |
| 3 # license that can be found in the LICENSE file or at |
| 4 # https://developers.google.com/open-source/licenses/bsd |
| 5 |
| 6 """Unit tests for RateLimiter. |
| 7 """ |
| 8 import unittest |
| 9 |
| 10 from google.appengine.api import memcache |
| 11 from google.appengine.ext import testbed |
| 12 |
| 13 import mox |
| 14 import os |
| 15 import settings |
| 16 |
| 17 from framework import ratelimiter |
| 18 from services import service_manager |
| 19 from testing import fake |
| 20 from testing import testing_helpers |
| 21 |
| 22 |
| 23 class RateLimiterTest(unittest.TestCase): |
| 24 def setUp(self): |
| 25 settings.ratelimiting_enabled = True |
| 26 self.testbed = testbed.Testbed() |
| 27 self.testbed.activate() |
| 28 self.testbed.init_memcache_stub() |
| 29 self.testbed.init_user_stub() |
| 30 |
| 31 self.mox = mox.Mox() |
| 32 self.services = service_manager.Services( |
| 33 config=fake.ConfigService(), |
| 34 issue=fake.IssueService(), |
| 35 user=fake.UserService(), |
| 36 project=fake.ProjectService(), |
| 37 ) |
| 38 self.project = self.services.project.TestAddProject('proj', project_id=987) |
| 39 |
| 40 self.ratelimiter = ratelimiter.RateLimiter() |
| 41 ratelimiter.COUNTRY_LIMITS = {} |
| 42 os.environ['USER_EMAIL'] = '' |
| 43 settings.ratelimiting_enabled = True |
| 44 settings.ratelimiting_cost_enabled = True |
| 45 ratelimiter.DEFAULT_LIMIT = 10 |
| 46 |
| 47 def tearDown(self): |
| 48 self.testbed.deactivate() |
| 49 self.mox.UnsetStubs() |
| 50 self.mox.ResetAll() |
| 51 # settings.ratelimiting_enabled = True |
| 52 |
| 53 def testCheckStart_pass(self): |
| 54 request, _ = testing_helpers.GetRequestObjects( |
| 55 project=self.project) |
| 56 request.headers['X-AppEngine-Country'] = 'US' |
| 57 request.remote_addr = '192.168.1.0' |
| 58 self.ratelimiter.CheckStart(request) |
| 59 # Should not throw an exception. |
| 60 |
| 61 def testCheckStart_fail(self): |
| 62 request, _ = testing_helpers.GetRequestObjects( |
| 63 project=self.project) |
| 64 request.headers['X-AppEngine-Country'] = 'US' |
| 65 request.remote_addr = '192.168.1.0' |
| 66 now = 0.0 |
| 67 cachekeysets, _, _, _ = ratelimiter._CacheKeys(request, now) |
| 68 values = [{key: ratelimiter.DEFAULT_LIMIT for key in cachekeys} for |
| 69 cachekeys in cachekeysets] |
| 70 for value in values: |
| 71 memcache.add_multi(value) |
| 72 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 73 self.ratelimiter.CheckStart(request, now) |
| 74 |
| 75 def testCheckStart_expiredEntries(self): |
| 76 request, _ = testing_helpers.GetRequestObjects( |
| 77 project=self.project) |
| 78 request.headers['X-AppEngine-Country'] = 'US' |
| 79 request.remote_addr = '192.168.1.0' |
| 80 now = 0.0 |
| 81 cachekeysets, _, _, _ = ratelimiter._CacheKeys(request, now) |
| 82 values = [{key: ratelimiter.DEFAULT_LIMIT for key in cachekeys} for |
| 83 cachekeys in cachekeysets] |
| 84 for value in values: |
| 85 memcache.add_multi(value) |
| 86 |
| 87 now = now + 2 * ratelimiter.EXPIRE_AFTER_SECS |
| 88 self.ratelimiter.CheckStart(request, now) |
| 89 # Should not throw an exception. |
| 90 |
| 91 def testCheckStart_repeatedCalls(self): |
| 92 request, _ = testing_helpers.GetRequestObjects( |
| 93 project=self.project) |
| 94 request.headers['X-AppEngine-Country'] = 'US' |
| 95 request.remote_addr = '192.168.1.0' |
| 96 now = 0.0 |
| 97 |
| 98 # Call CheckStart once every minute. Should be ok. |
| 99 for _ in range(ratelimiter.N_MINUTES): |
| 100 self.ratelimiter.CheckStart(request, now) |
| 101 now = now + 120.0 |
| 102 |
| 103 # Call CheckStart more than DEFAULT_LIMIT times in the same minute. |
| 104 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 105 for _ in range(ratelimiter.DEFAULT_LIMIT + 2): |
| 106 now = now + 0.001 |
| 107 self.ratelimiter.CheckStart(request, now) |
| 108 |
| 109 def testCheckStart_differentIPs(self): |
| 110 now = 0.0 |
| 111 |
| 112 ratelimiter.COUNTRY_LIMITS = {} |
| 113 # Exceed DEFAULT_LIMIT calls, but vary remote_addr so different |
| 114 # remote addresses aren't ratelimited together. |
| 115 for m in range(ratelimiter.DEFAULT_LIMIT * 2): |
| 116 request, _ = testing_helpers.GetRequestObjects( |
| 117 project=self.project) |
| 118 request.headers['X-AppEngine-Country'] = 'US' |
| 119 request.remote_addr = '192.168.1.%d' % (m % 16) |
| 120 ratelimiter._CacheKeys(request, now) |
| 121 self.ratelimiter.CheckStart(request, now) |
| 122 now = now + 0.001 |
| 123 |
| 124 # Exceed the limit, but only for one IP address. The |
| 125 # others should be fine. |
| 126 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 127 for m in range(ratelimiter.DEFAULT_LIMIT): |
| 128 request, _ = testing_helpers.GetRequestObjects( |
| 129 project=self.project) |
| 130 request.headers['X-AppEngine-Country'] = 'US' |
| 131 request.remote_addr = '192.168.1.0' |
| 132 ratelimiter._CacheKeys(request, now) |
| 133 self.ratelimiter.CheckStart(request, now) |
| 134 now = now + 0.001 |
| 135 |
| 136 # Now proceed to make requests for all of the other IP |
| 137 # addresses besides .0. |
| 138 for m in range(ratelimiter.DEFAULT_LIMIT * 2): |
| 139 request, _ = testing_helpers.GetRequestObjects( |
| 140 project=self.project) |
| 141 request.headers['X-AppEngine-Country'] = 'US' |
| 142 # Skip .0 since it's already exceeded the limit. |
| 143 request.remote_addr = '192.168.1.%d' % (m + 1) |
| 144 ratelimiter._CacheKeys(request, now) |
| 145 self.ratelimiter.CheckStart(request, now) |
| 146 now = now + 0.001 |
| 147 |
| 148 def testCheckStart_sameIPDifferentUserIDs(self): |
| 149 # Behind a NAT, e.g. |
| 150 now = 0.0 |
| 151 |
| 152 # Exceed DEFAULT_LIMIT calls, but vary user_id so different |
| 153 # users behind the same IP aren't ratelimited together. |
| 154 for m in range(ratelimiter.DEFAULT_LIMIT * 2): |
| 155 request, _ = testing_helpers.GetRequestObjects( |
| 156 project=self.project) |
| 157 request.remote_addr = '192.168.1.0' |
| 158 os.environ['USER_EMAIL'] = '%s@example.com' % m |
| 159 request.headers['X-AppEngine-Country'] = 'US' |
| 160 ratelimiter._CacheKeys(request, now) |
| 161 self.ratelimiter.CheckStart(request, now) |
| 162 now = now + 0.001 |
| 163 |
| 164 # Exceed the limit, but only for one userID+IP address. The |
| 165 # others should be fine. |
| 166 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 167 for m in range(ratelimiter.DEFAULT_LIMIT + 2): |
| 168 request, _ = testing_helpers.GetRequestObjects( |
| 169 project=self.project) |
| 170 request.headers['X-AppEngine-Country'] = 'US' |
| 171 request.remote_addr = '192.168.1.0' |
| 172 os.environ['USER_EMAIL'] = '42@example.com' |
| 173 ratelimiter._CacheKeys(request, now) |
| 174 self.ratelimiter.CheckStart(request, now) |
| 175 now = now + 0.001 |
| 176 |
| 177 # Now proceed to make requests for other user IDs |
| 178 # besides 42. |
| 179 for m in range(ratelimiter.DEFAULT_LIMIT * 2): |
| 180 request, _ = testing_helpers.GetRequestObjects( |
| 181 project=self.project) |
| 182 request.headers['X-AppEngine-Country'] = 'US' |
| 183 # Skip .0 since it's already exceeded the limit. |
| 184 request.remote_addr = '192.168.1.0' |
| 185 os.environ['USER_EMAIL'] = '%s@example.com' % (43 + m) |
| 186 ratelimiter._CacheKeys(request, now) |
| 187 self.ratelimiter.CheckStart(request, now) |
| 188 now = now + 0.001 |
| 189 |
| 190 def testCheckStart_ratelimitingDisabled(self): |
| 191 settings.ratelimiting_enabled = False |
| 192 request, _ = testing_helpers.GetRequestObjects( |
| 193 project=self.project) |
| 194 request.headers['X-AppEngine-Country'] = 'US' |
| 195 request.remote_addr = '192.168.1.0' |
| 196 now = 0.0 |
| 197 |
| 198 # Call CheckStart a lot. Should be ok. |
| 199 for _ in range(ratelimiter.DEFAULT_LIMIT): |
| 200 self.ratelimiter.CheckStart(request, now) |
| 201 now = now + 0.001 |
| 202 |
| 203 def testCheckStart_perCountryLoggedOutLimit(self): |
| 204 ratelimiter.COUNTRY_LIMITS['US'] = 10 |
| 205 |
| 206 request, _ = testing_helpers.GetRequestObjects( |
| 207 project=self.project) |
| 208 request.headers[ratelimiter.COUNTRY_HEADER] = 'US' |
| 209 request.remote_addr = '192.168.1.1' |
| 210 now = 0.0 |
| 211 |
| 212 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 213 for m in range(ratelimiter.DEFAULT_LIMIT + 2): |
| 214 self.ratelimiter.CheckStart(request, now) |
| 215 # Vary remote address to make sure the limit covers |
| 216 # the whole country, regardless of IP. |
| 217 request.remote_addr = '192.168.1.%d' % m |
| 218 now = now + 0.001 |
| 219 |
| 220 # CheckStart for a country that isn't covered by a country-specific limit. |
| 221 request.headers['X-AppEngine-Country'] = 'UK' |
| 222 for m in range(11): |
| 223 self.ratelimiter.CheckStart(request, now) |
| 224 # Vary remote address to make sure the limit covers |
| 225 # the whole country, regardless of IP. |
| 226 request.remote_addr = '192.168.1.%d' % m |
| 227 now = now + 0.001 |
| 228 |
| 229 # And regular rate limits work per-IP. |
| 230 request.remote_addr = '192.168.1.1' |
| 231 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 232 for m in range(ratelimiter.DEFAULT_LIMIT): |
| 233 self.ratelimiter.CheckStart(request, now) |
| 234 # Vary remote address to make sure the limit covers |
| 235 # the whole country, regardless of IP. |
| 236 now = now + 0.001 |
| 237 |
| 238 def testCheckEnd_overCostThresh(self): |
| 239 request, _ = testing_helpers.GetRequestObjects( |
| 240 project=self.project) |
| 241 request.headers[ratelimiter.COUNTRY_HEADER] = 'US' |
| 242 request.remote_addr = '192.168.1.1' |
| 243 start_time = 0.0 |
| 244 |
| 245 # Send some requests, all under the limit. |
| 246 for _ in range(ratelimiter.DEFAULT_LIMIT-1): |
| 247 start_time = start_time + 0.001 |
| 248 self.ratelimiter.CheckStart(request, start_time) |
| 249 now = start_time + 0.010 |
| 250 self.ratelimiter.CheckEnd(request, now, start_time) |
| 251 |
| 252 # Now issue some more request, this time taking long |
| 253 # enough to get the cost threshold penalty. |
| 254 # Fast forward enough to impact a later bucket than the |
| 255 # previous requests. |
| 256 start_time = now + 120.0 |
| 257 self.ratelimiter.CheckStart(request, start_time) |
| 258 |
| 259 # Take longer than the threshold to process the request. |
| 260 now = start_time + (settings.ratelimiting_cost_thresh_ms + 1) / 1000 |
| 261 |
| 262 # The request finished, taking longer than the cost |
| 263 # threshold. |
| 264 self.ratelimiter.CheckEnd(request, now, start_time) |
| 265 |
| 266 with self.assertRaises(ratelimiter.RateLimitExceeded): |
| 267 # One more request after the expensive query should |
| 268 # throw an excpetion. |
| 269 self.ratelimiter.CheckStart(request, start_time) |
| 270 |
| 271 def testCheckEnd_overCostThreshButDisabled(self): |
| 272 request, _ = testing_helpers.GetRequestObjects( |
| 273 project=self.project) |
| 274 request.headers[ratelimiter.COUNTRY_HEADER] = 'US' |
| 275 request.remote_addr = '192.168.1.1' |
| 276 start_time = 0.0 |
| 277 settings.ratelimiting_cost_enabled = False |
| 278 |
| 279 # Send some requests, all under the limit. |
| 280 for _ in range(ratelimiter.DEFAULT_LIMIT-1): |
| 281 start_time = start_time + 0.001 |
| 282 self.ratelimiter.CheckStart(request, start_time) |
| 283 now = start_time + 0.010 |
| 284 self.ratelimiter.CheckEnd(request, now, start_time) |
| 285 |
| 286 # Now issue some more request, this time taking long |
| 287 # enough to get the cost threshold penalty. |
| 288 # Fast forward enough to impact a later bucket than the |
| 289 # previous requests. |
| 290 start_time = now + 120.0 |
| 291 self.ratelimiter.CheckStart(request, start_time) |
| 292 |
| 293 # Take longer than the threshold to process the request. |
| 294 now = start_time + (settings.ratelimiting_cost_thresh_ms + 10)/1000 |
| 295 |
| 296 # The request finished, taking longer than the cost |
| 297 # threshold. |
| 298 self.ratelimiter.CheckEnd(request, now, start_time) |
| 299 |
| 300 # One more request after the expensive query should |
| 301 # throw an excpetion, but cost thresholds are disabled. |
| 302 self.ratelimiter.CheckStart(request, start_time) |
| 303 |
| 304 def testChekcEnd_underCostThresh(self): |
| 305 request, _ = testing_helpers.GetRequestObjects( |
| 306 project=self.project) |
| 307 request.headers[ratelimiter.COUNTRY_HEADER] = 'asdasd' |
| 308 request.remote_addr = '192.168.1.1' |
| 309 start_time = 0.0 |
| 310 |
| 311 # Send some requests, all under the limit. |
| 312 for _ in range(ratelimiter.DEFAULT_LIMIT): |
| 313 self.ratelimiter.CheckStart(request, start_time) |
| 314 now = start_time + 0.010 |
| 315 self.ratelimiter.CheckEnd(request, now, start_time) |
| 316 start_time = now + 0.010 |
| 317 |
| 318 def testChekcEnd_underCostThresh(self): |
| 319 request, _ = testing_helpers.GetRequestObjects( |
| 320 project=self.project) |
| 321 request.headers[ratelimiter.COUNTRY_HEADER] = 'asdasd' |
| 322 request.remote_addr = '192.168.1.1' |
| 323 start_time = 0.0 |
| 324 |
| 325 # Send some requests, all under the limit. |
| 326 for _ in range(ratelimiter.DEFAULT_LIMIT): |
| 327 self.ratelimiter.CheckStart(request, start_time) |
| 328 now = start_time + 0.01 |
| 329 self.ratelimiter.CheckEnd(request, now, start_time) |
| 330 start_time = now + 0.01 |
OLD | NEW |