framework/ratelimiter.py - monorail-avm99963 - Gitiles

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed 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.

 """
 from __future__ import print_function
 from __future__ import division
 from __future__ import absolute_import

 import datetime
 import logging
 import os
 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

 from services import client_config_svc


 N_MINUTES = 5
 EXPIRE_AFTER_SECS = 60 * 60
 DEFAULT_LIMIT = 60 * N_MINUTES  # 300 page requests in 5 minutes is 1 QPS.
 DEFAULT_API_QPM = 1000  # For example, chromiumdash uses ~64 per page, 8s each.

 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_ALLOWLIST = ['default', 'api']


 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


 def _CreateApiCacheKeys(client_id, client_email, now_sec):
   country = os.environ.get('HTTP_X_APPENGINE_COUNTRY')
   ip = os.environ.get('REMOTE_ADDR')
   now = datetime.datetime.fromtimestamp(now_sec)
   minute_buckets = [now - datetime.timedelta(minutes=m) for m in
                     range(N_MINUTES)]
   minute_strs = [str(minute_bucket.replace(second=0, microsecond=0))
                  for minute_bucket in minute_buckets]
   keys = []

   if client_id and client_id != 'anonymous':
     keys.append(['apiratelimit-%s-%s' % (client_id, minute_str)
                  for minute_str in minute_strs])
   if client_email:
     keys.append(['apiratelimit-%s-%s' % (client_email, minute_str)
                  for minute_str in minute_strs])
   else:
     keys.append(['apiratelimit-%s-%s' % (ip, minute_str)
                  for minute_str in minute_strs])
     if country in COUNTRY_LIMITS:
       keys.append(['apiratelimit-%s-%s' % (country, minute_str)
                    for minute_str in minute_strs])

   return keys


 class RateLimiter(object):

   blocked_requests = ts_mon.CounterMetric(
       'monorail/ratelimiter/blocked_request',
       'Count of requests that exceeded the rate limit and were blocked.',
       None)
   limit_exceeded = ts_mon.CounterMetric(
       'monorail/ratelimiter/rate_exceeded',
       'Count of requests that exceeded the rate limit.',
       None)
   cost_thresh_exceeded = ts_mon.CounterMetric(
       'monorail/ratelimiter/cost_thresh_exceeded',
       'Count of requests that were expensive to process',
       None)
   checks = ts_mon.CounterMetric(
       'monorail/ratelimiter/check',
       'Count of checks done, by fail/success type.',
       [ts_mon.StringField('type')])

   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_ALLOWLIST 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.
     self._AuxCheckStart(
         keysets, COUNTRY_LIMITS.get(country, DEFAULT_LIMIT),
         settings.ratelimiting_enabled,
         RateLimitExceeded(country=country, ip=ip, user_email=user_email))

   def _AuxCheckStart(self, keysets, limit, ratelimiting_enabled,
                      exception_obj):
     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_closed'})
           raise exception_obj
         self.checks.increment({'type': 'fail_open'})

       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('%s, %d' % (exception_obj.message, count))

         self.limit_exceeded.increment()

         if ratelimiting_enabled:
           self.blocked_requests.increment()
           raise exception_obj

       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_ALLOWLIST):
       return

     elapsed_ms = int((now - start_time) * 1000)
     # Would it kill the python lib maintainers to have timedelta.total_ms()?
     penalty = elapsed_ms // settings.ratelimiting_ms_per_count - 1
     if penalty >= 1:
       # TODO: Look into caching the keys instead of generating them twice
       # for every request. Say, return them from CheckStart so they can
       # be passed back in here later.
       keysets, country, ip, user_email = _CacheKeys(request, start_time)

       self._AuxCheckEnd(
           keysets,
           'Rate Limit Cost Threshold Exceeded: %s, %s, %s' % (
               country, ip, user_email),
           penalty)

   def _AuxCheckEnd(self, keysets, log_str, penalty):
     self.cost_thresh_exceeded.increment()
     for keys in keysets:
       logging.info(log_str)

       # 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, delta=penalty, initial_value=0)


 class ApiRateLimiter(RateLimiter):

   blocked_requests = ts_mon.CounterMetric(
       'monorail/apiratelimiter/blocked_request',
       'Count of requests that exceeded the rate limit and were blocked.',
       None)
   limit_exceeded = ts_mon.CounterMetric(
       'monorail/apiratelimiter/rate_exceeded',
       'Count of requests that exceeded the rate limit.',
       None)
   cost_thresh_exceeded = ts_mon.CounterMetric(
       'monorail/apiratelimiter/cost_thresh_exceeded',
       'Count of requests that were expensive to process',
       None)
   checks = ts_mon.CounterMetric(
       'monorail/apiratelimiter/check',
       'Count of checks done, by fail/success type.',
       [ts_mon.StringField('type')])

   #pylint: disable=arguments-differ
   def CheckStart(self, client_id, client_email, now=None):
     if now is None:
       now = time.time()

     keysets = _CreateApiCacheKeys(client_id, client_email, now)
     qpm_limit = client_config_svc.GetQPMDict().get(
         client_email, DEFAULT_API_QPM)
     if qpm_limit < DEFAULT_API_QPM:
       qpm_limit = DEFAULT_API_QPM
     window_limit = qpm_limit * N_MINUTES
     self._AuxCheckStart(
         keysets, window_limit,
         settings.api_ratelimiting_enabled,
         ApiRateLimitExceeded(client_id, client_email))

   #pylint: disable=arguments-differ
   def CheckEnd(self, client_id, client_email, now, start_time):

     elapsed_ms = int((now - start_time) * 1000)
     penalty = elapsed_ms // settings.ratelimiting_ms_per_count - 1

     if penalty >= 1:
       keysets = _CreateApiCacheKeys(client_id, client_email, start_time)
       self._AuxCheckEnd(
           keysets,
           'API Rate Limit Cost Threshold Exceeded: %s, %s' % (
               client_id, client_email),
           penalty)


 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
     message = 'RateLimitExceeded: %s, %s, %s' % (
         self.country, self.ip, self.user_email)
     super(RateLimitExceeded, self).__init__(message)


 class ApiRateLimitExceeded(Exception):
   def __init__(self, client_id, client_email):
     self.client_id = client_id
     self.client_email = client_email
     message = 'RateLimitExceeded: %s, %s' % (
         self.client_id, self.client_email)
     super(ApiRateLimitExceeded, self).__init__(message)
	# Copyright 2016 The Chromium Authors. All rights reserved.
	# Use of this source code is governed 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.

	"""
	from __future__ import print_function
	from __future__ import division
	from __future__ import absolute_import

	import datetime
	import logging
	import os
	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

	from services import client_config_svc


	N_MINUTES = 5
	EXPIRE_AFTER_SECS = 60 * 60
	DEFAULT_LIMIT = 60 * N_MINUTES # 300 page requests in 5 minutes is 1 QPS.
	DEFAULT_API_QPM = 1000 # For example, chromiumdash uses ~64 per page, 8s each.

	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_ALLOWLIST = ['default', 'api']


	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


	def _CreateApiCacheKeys(client_id, client_email, now_sec):
	country = os.environ.get('HTTP_X_APPENGINE_COUNTRY')
	ip = os.environ.get('REMOTE_ADDR')
	now = datetime.datetime.fromtimestamp(now_sec)
	minute_buckets = [now - datetime.timedelta(minutes=m) for m in
	range(N_MINUTES)]
	minute_strs = [str(minute_bucket.replace(second=0, microsecond=0))
	for minute_bucket in minute_buckets]
	keys = []

	if client_id and client_id != 'anonymous':
	keys.append(['apiratelimit-%s-%s' % (client_id, minute_str)
	for minute_str in minute_strs])
	if client_email:
	keys.append(['apiratelimit-%s-%s' % (client_email, minute_str)
	for minute_str in minute_strs])
	else:
	keys.append(['apiratelimit-%s-%s' % (ip, minute_str)
	for minute_str in minute_strs])
	if country in COUNTRY_LIMITS:
	keys.append(['apiratelimit-%s-%s' % (country, minute_str)
	for minute_str in minute_strs])

	return keys


	class RateLimiter(object):

	blocked_requests = ts_mon.CounterMetric(
	'monorail/ratelimiter/blocked_request',
	'Count of requests that exceeded the rate limit and were blocked.',
	None)
	limit_exceeded = ts_mon.CounterMetric(
	'monorail/ratelimiter/rate_exceeded',
	'Count of requests that exceeded the rate limit.',
	None)
	cost_thresh_exceeded = ts_mon.CounterMetric(
	'monorail/ratelimiter/cost_thresh_exceeded',
	'Count of requests that were expensive to process',
	None)
	checks = ts_mon.CounterMetric(
	'monorail/ratelimiter/check',
	'Count of checks done, by fail/success type.',
	[ts_mon.StringField('type')])

	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_ALLOWLIST 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.
	self._AuxCheckStart(
	keysets, COUNTRY_LIMITS.get(country, DEFAULT_LIMIT),
	settings.ratelimiting_enabled,
	RateLimitExceeded(country=country, ip=ip, user_email=user_email))

	def _AuxCheckStart(self, keysets, limit, ratelimiting_enabled,
	exception_obj):
	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_closed'})
	raise exception_obj
	self.checks.increment({'type': 'fail_open'})

	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('%s, %d' % (exception_obj.message, count))

	self.limit_exceeded.increment()

	if ratelimiting_enabled:
	self.blocked_requests.increment()
	raise exception_obj

	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_ALLOWLIST):
	return

	elapsed_ms = int((now - start_time) * 1000)
	# Would it kill the python lib maintainers to have timedelta.total_ms()?
	penalty = elapsed_ms // settings.ratelimiting_ms_per_count - 1
	if penalty >= 1:
	# TODO: Look into caching the keys instead of generating them twice
	# for every request. Say, return them from CheckStart so they can
	# be passed back in here later.
	keysets, country, ip, user_email = _CacheKeys(request, start_time)

	self._AuxCheckEnd(
	keysets,
	'Rate Limit Cost Threshold Exceeded: %s, %s, %s' % (
	country, ip, user_email),
	penalty)

	def _AuxCheckEnd(self, keysets, log_str, penalty):
	self.cost_thresh_exceeded.increment()
	for keys in keysets:
	logging.info(log_str)

	# 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, delta=penalty, initial_value=0)


	class ApiRateLimiter(RateLimiter):

	blocked_requests = ts_mon.CounterMetric(
	'monorail/apiratelimiter/blocked_request',
	'Count of requests that exceeded the rate limit and were blocked.',
	None)
	limit_exceeded = ts_mon.CounterMetric(
	'monorail/apiratelimiter/rate_exceeded',
	'Count of requests that exceeded the rate limit.',
	None)
	cost_thresh_exceeded = ts_mon.CounterMetric(
	'monorail/apiratelimiter/cost_thresh_exceeded',
	'Count of requests that were expensive to process',
	None)
	checks = ts_mon.CounterMetric(
	'monorail/apiratelimiter/check',
	'Count of checks done, by fail/success type.',
	[ts_mon.StringField('type')])

	#pylint: disable=arguments-differ
	def CheckStart(self, client_id, client_email, now=None):
	if now is None:
	now = time.time()

	keysets = _CreateApiCacheKeys(client_id, client_email, now)
	qpm_limit = client_config_svc.GetQPMDict().get(
	client_email, DEFAULT_API_QPM)
	if qpm_limit < DEFAULT_API_QPM:
	qpm_limit = DEFAULT_API_QPM
	window_limit = qpm_limit * N_MINUTES
	self._AuxCheckStart(
	keysets, window_limit,
	settings.api_ratelimiting_enabled,
	ApiRateLimitExceeded(client_id, client_email))

	#pylint: disable=arguments-differ
	def CheckEnd(self, client_id, client_email, now, start_time):

	elapsed_ms = int((now - start_time) * 1000)
	penalty = elapsed_ms // settings.ratelimiting_ms_per_count - 1

	if penalty >= 1:
	keysets = _CreateApiCacheKeys(client_id, client_email, start_time)
	self._AuxCheckEnd(
	keysets,
	'API Rate Limit Cost Threshold Exceeded: %s, %s' % (
	client_id, client_email),
	penalty)


	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
	message = 'RateLimitExceeded: %s, %s, %s' % (
	self.country, self.ip, self.user_email)
	super(RateLimitExceeded, self).__init__(message)


	class ApiRateLimitExceeded(Exception):
	def __init__(self, client_id, client_email):
	self.client_id = client_id
	self.client_email = client_email
	message = 'RateLimitExceeded: %s, %s' % (
	self.client_id, self.client_email)
	super(ApiRateLimitExceeded, self).__init__(message)