blob: 0c968dfa722700dbd8c5bd8527254c2f78384aa8 [file] [log] [blame]
# 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)