| # 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 |
| |
| """Request rate limiting implementation. |
| |
| This is intented to be used for automatic DDoS protection. |
| |
| """ |
| |
| import datetime |
| import logging |
| import settings |
| import time |
| |
| from infra_libs import ts_mon |
| |
| from google.appengine.api import memcache |
| from google.appengine.api.modules import modules |
| from google.appengine.api import users |
| |
| N_MINUTES = 5 |
| EXPIRE_AFTER_SECS = 60 * 60 |
| DEFAULT_LIMIT = 300 # 300 requests in 5 minutes is 1 QPS. |
| |
| ANON_USER = 'anon' |
| |
| COUNTRY_HEADER = 'X-AppEngine-Country' |
| |
| COUNTRY_LIMITS = { |
| # Two-letter country code: max requests per N_MINUTES |
| # This limit will apply to all requests coming |
| # from this country. |
| # To add a country code, see GAE logs and use the |
| # appropriate code from https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2 |
| # E.g., 'cn': 300, # Limit to 1 QPS. |
| } |
| |
| # Modules not in this list will not have rate limiting applied by this |
| # class. |
| MODULE_WHITELIST = ['default'] |
| |
| def _CacheKeys(request, now_sec): |
| """ Returns an array of arrays. Each array contains strings with |
| the same prefix and a timestamp suffix, starting with the most |
| recent and decrementing by 1 minute each time. |
| """ |
| now = datetime.datetime.fromtimestamp(now_sec) |
| country = request.headers.get(COUNTRY_HEADER, 'ZZ') |
| ip = request.remote_addr |
| minute_buckets = [now - datetime.timedelta(minutes=m) for m in |
| range(N_MINUTES)] |
| user = users.get_current_user() |
| user_email = user.email() if user else ANON_USER |
| |
| # <IP, country, user_email> to be rendered into each key prefix. |
| prefixes = [] |
| |
| # All logged-in users get a per-user rate limit, regardless of IP and country. |
| if user: |
| prefixes.append(['ALL', 'ALL', user.email()]) |
| else: |
| # All anon requests get a per-IP ratelimit. |
| prefixes.append([ip, 'ALL', 'ALL']) |
| |
| # All requests from a problematic country get a per-country rate limit, |
| # regardless of the user (even a non-logged-in one) or IP. |
| if country in COUNTRY_LIMITS: |
| prefixes.append(['ALL', country, 'ALL']) |
| |
| keysets = [] |
| for prefix in prefixes: |
| keysets.append(['ratelimit-%s-%s' % ('-'.join(prefix), |
| str(minute_bucket.replace(second=0, microsecond=0))) |
| for minute_bucket in minute_buckets]) |
| |
| return keysets, country, ip, user_email |
| |
| class RateLimiter: |
| blocked_requests = ts_mon.CounterMetric( |
| 'monorail/ratelimiter/blocked_request') |
| limit_exceeded = ts_mon.CounterMetric( |
| 'monorail/ratelimiter/rate_exceeded') |
| cost_thresh_exceeded = ts_mon.CounterMetric( |
| 'monorail/ratelimiter/cost_thresh_exceeded') |
| checks = ts_mon.CounterMetric( |
| 'monorail/ratelimiter/check') |
| |
| def __init__(self, _cache=memcache, fail_open=True, **_kwargs): |
| self.fail_open = fail_open |
| |
| def CheckStart(self, request, now=None): |
| if (modules.get_current_module_name() not in MODULE_WHITELIST or |
| users.is_current_user_admin()): |
| return |
| logging.info('X-AppEngine-Country: %s' % |
| request.headers.get(COUNTRY_HEADER, 'ZZ')) |
| |
| if now is None: |
| now = time.time() |
| |
| keysets, country, ip, user_email = _CacheKeys(request, now) |
| # There are either two or three sets of keys in keysets. |
| # Three if the user's country is in COUNTRY_LIMITS, otherwise two. |
| for keys in keysets: |
| count = 0 |
| try: |
| counters = memcache.get_multi(keys) |
| count = sum(counters.values()) |
| self.checks.increment({'type': 'success'}) |
| except Exception as e: |
| logging.error(e) |
| if not self.fail_open: |
| self.checks.increment({'type': 'fail_open'}) |
| raise RateLimitExceeded(country=country, ip=ip, user_email=user_email) |
| self.checks.increment({'type': 'fail_closed'}) |
| |
| limit = COUNTRY_LIMITS.get(country, DEFAULT_LIMIT) |
| if count > limit: |
| # Since webapp2 won't let us return a 429 error code |
| # <http://tools.ietf.org/html/rfc6585#section-4>, we can't |
| # monitor rate limit exceeded events with our standard tools. |
| # We return a 400 with a custom error message to the client, |
| # and this logging is so we can monitor it internally. |
| logging.info('Rate Limit Exceeded: %s, %s, %s, %d' % ( |
| country, ip, user_email, count)) |
| |
| self.limit_exceeded.increment() |
| |
| if settings.ratelimiting_enabled: |
| self.blocked_requests.increment() |
| raise RateLimitExceeded(country=country, ip=ip, user_email=user_email) |
| |
| k = keys[0] |
| # Only update the latest *time* bucket for each prefix (reverse chron). |
| memcache.add(k, 0, time=EXPIRE_AFTER_SECS) |
| memcache.incr(k, initial_value=0) |
| |
| def CheckEnd(self, request, now, start_time): |
| """If a request was expensive to process, charge some extra points |
| against this set of buckets. |
| We pass in both now and start_time so we can update the buckets |
| based on keys created from start_time instead of now. |
| now and start_time are float seconds. |
| """ |
| if (modules.get_current_module_name() not in MODULE_WHITELIST or |
| not settings.ratelimiting_cost_enabled): |
| return |
| |
| elapsed_ms = (now - start_time) * 1000 |
| # Would it kill the python lib maintainers to have timedelta.total_ms()? |
| if elapsed_ms < settings.ratelimiting_cost_thresh_ms: |
| return |
| |
| # TODO: Look into caching the keys instead of generating them twice |
| # for every request. Say, return them from CheckStart so they can |
| # be bassed back in here later. |
| keysets, country, ip, user_email = _CacheKeys(request, start_time) |
| for keys in keysets: |
| logging.info('Rate Limit Cost Threshold Exceeded: %s, %s, %s' % ( |
| country, ip, user_email)) |
| self.cost_thresh_exceeded.increment_by(settings.ratelimiting_cost_penalty) |
| |
| # Only update the latest *time* bucket for each prefix (reverse chron). |
| k = keys[0] |
| memcache.add(k, 0, time=EXPIRE_AFTER_SECS) |
| memcache.incr(k, initial_value=0) |
| |
| class RateLimitExceeded(Exception): |
| def __init__(self, country=None, ip=None, user_email=None, **_kwargs): |
| self.country = country |
| self.ip = ip |
| self.user_email = user_email |
| |
| def __str__(self): |
| return 'RateLimitExceeded: %s, %s, %s' % ( |
| self.country, self.ip, self.user_email) |