def set_rate_limits(per_ip, window_size):
"""Set rate limit parameters for the AcousticBrainz webserver. If no arguments
are provided, print the current limits. To set limits, specify PER_IP and WINDOW_SIZE
\b
PER_IP: the number of requests allowed per IP address
WINDOW_SIZE: the window in number of seconds for how long the limit is applied
"""
current_limit_per_ip = cache.get(ratelimit.ratelimit_per_ip_key)
current_limit_window = cache.get(ratelimit.ratelimit_window_key)
click.echo("Current values:")
if current_limit_per_ip is None and current_limit_window is None:
click.echo("No values set, showing limit defaults")
current_limit_per_ip = ratelimit.ratelimit_per_ip_default
current_limit_window = ratelimit.ratelimit_window_default
click.echo("Requests per IP: %s" % current_limit_per_ip)
click.echo("Window size (s): %s" % current_limit_window)
if per_ip is not None and window_size is not None:
if per_ip / float(window_size) < 1:
click.echo(
"Warning: Effective rate limit is less than 1 query per second"
)
ratelimit.set_rate_limits(per_ip, per_ip, window_size)
print("New ratelimit parameters set:")
click.echo("Requests per IP: %s" % per_ip)
click.echo("Window size (s): %s" % window_size)
def test_ratelimit(self):
""" Tests that the ratelimit decorator works
"""
# Set the limits as per defines in this class
set_rate_limits(self.max_token_requests, self.max_ip_requests,
self.ratelimit_window)
# create an app
app = flask.CustomFlask(__name__)
self.assertIsNotNone(app)
app.debug = True
app.config['SECRET_KEY'] = 'this is a totally secret key btw'
app.init_debug_toolbar()
@app.after_request
def after_request_callbacks(response):
return inject_x_rate_headers(response)
# add a dummy route
@app.route('/')
@ratelimit()
def index():
return '<html><body>test</body></html>'
def print_headers(response):
print("X-RateLimit-Remaining",
response.headers['X-RateLimit-Remaining'])
print("X-RateLimit-Limit", response.headers['X-RateLimit-Limit'])
print("X-RateLimit-Reset", response.headers['X-RateLimit-Reset'])
print("X-RateLimit-Reset-In",
response.headers['X-RateLimit-Reset-In'])
print()
def make_requests(client, nominal_num_requests, token=None):
print("===== make %d requests" % nominal_num_requests)
# make one more than the allowed number of requests to catch the 429
num_requests = nominal_num_requests + 1
# make a specified number of requests
while True:
reset_time = 0
restart = False
for i in range(num_requests):
if token:
response = client.get('/',
headers={'Authorization': token})
else:
response = client.get('/')
if reset_time == 0:
reset_time = response.headers['X-RateLimit-Reset']
if reset_time != response.headers['X-RateLimit-Reset']:
# Whoops, we didn't get our tests done before the window expired. start over.
restart = True
# when restarting we need to do one request less, since the current requests counts to the new window
num_requests = nominal_num_requests
break
if i == num_requests - 1:
self.assertEqual(response.status_code, 429)
else:
self.assertEqual(response.status_code, 200)
self.assertEqual(
int(response.headers['X-RateLimit-Remaining']),
num_requests - i - 2)
print_headers(response)
sleep(1.1)
if not restart:
break
client = app.test_client()
# Make a pile of requests based on IP address
make_requests(client, self.max_ip_requests)
# Set a user token and make requests based on the token
cache.flush_all()
set_user_validation_function(validate_user)
set_rate_limits(self.max_token_requests, self.max_ip_requests,
self.ratelimit_window)
make_requests(client,
self.max_token_requests,
token="Token %s" % valid_user)
def create_app(debug=None):
app = CustomFlask(
import_name=__name__,
use_flask_uuid=True,
)
# Configuration
load_config(app)
if debug is not None:
app.debug = debug
if app.debug and app.config['SECRET_KEY']:
app.init_debug_toolbar()
# Logging
app.init_loggers(file_config=app.config.get('LOG_FILE'),
email_config=app.config.get('LOG_EMAIL'),
sentry_config=app.config.get('LOG_SENTRY')
)
# Database connection
from db import init_db_engine
init_db_engine(app.config['SQLALCHEMY_DATABASE_URI'])
# Cache
if 'REDIS_HOST' in app.config and\
'REDIS_PORT' in app.config and\
'REDIS_NAMESPACE' in app.config and\
'REDIS_NS_VERSIONS_LOCATION' in app.config:
if not os.path.exists(app.config['REDIS_NS_VERSIONS_LOCATION']):
os.makedirs(app.config['REDIS_NS_VERSIONS_LOCATION'])
from brainzutils import cache
cache.init(
host=app.config['REDIS_HOST'],
port=app.config['REDIS_PORT'],
namespace=app.config['REDIS_NAMESPACE'],
ns_versions_loc=app.config['REDIS_NS_VERSIONS_LOCATION'])
else:
raise Exception('One or more redis cache configuration options are missing from config.py')
# Add rate limiting support
@app.after_request
def after_request_callbacks(response):
return inject_x_rate_headers(response)
# check for ratelimit config values and set them if present
if 'RATELIMIT_PER_IP' in app.config and 'RATELIMIT_WINDOW' in app.config:
set_rate_limits(app.config['RATELIMIT_PER_IP'], app.config['RATELIMIT_PER_IP'], app.config['RATELIMIT_WINDOW'])
# MusicBrainz
import musicbrainzngs
from db import SCHEMA_VERSION
musicbrainzngs.set_useragent(app.config['MUSICBRAINZ_USERAGENT'], SCHEMA_VERSION)
if app.config['MUSICBRAINZ_HOSTNAME']:
musicbrainzngs.set_hostname(app.config['MUSICBRAINZ_HOSTNAME'])
# OAuth
from webserver.login import login_manager, provider
login_manager.init_app(app)
provider.init(app.config['MUSICBRAINZ_CLIENT_ID'],
app.config['MUSICBRAINZ_CLIENT_SECRET'])
# Error handling
from webserver.errors import init_error_handlers
init_error_handlers(app)
# Static files
import static_manager
# Template utilities
app.jinja_env.add_extension('jinja2.ext.do')
from webserver import utils
app.jinja_env.filters['date'] = utils.reformat_date
app.jinja_env.filters['datetime'] = utils.reformat_datetime
# During development, built js and css assets don't have a hash, but in production we use
# a manifest to map a name to name.hash.extension for caching/cache busting
if app.debug:
app.context_processor(lambda: dict(get_static_path=static_manager.development_get_static_path))
else:
static_manager.read_manifest()
app.context_processor(lambda: dict(get_static_path=static_manager.manifest_get_static_path))
_register_blueprints(app)
# Admin section
from flask_admin import Admin
from webserver.admin import views as admin_views
admin = Admin(app, index_view=admin_views.HomeView(name='Admin'))
admin.add_view(admin_views.AdminsView(name='Admins'))
@app.before_request
def prod_https_login_redirect():
""" Redirect to HTTPS in production except for the API endpoints
"""
if urlparse.urlsplit(request.url).scheme == 'http' \
and app.config['DEBUG'] == False \
and app.config['TESTING'] == False \
and request.blueprint not in ('api', 'api_v1_core', 'api_v1_datasets', 'api_v1_dataset_eval'):
url = request.url[7:] # remove http:// from url
return redirect('https://{}'.format(url), 301)
@app.before_request
def before_request_gdpr_check():
# skip certain pages, static content and the API
if request.path == url_for('index.gdpr_notice') \
or request.path == url_for('login.logout') \
or request.path.startswith('/_debug') \
or request.path.startswith('/static') \
or request.path.startswith(API_PREFIX):
return
# otherwise if user is logged in and hasn't agreed to gdpr,
# redirect them to agree to terms page.
elif current_user.is_authenticated and current_user.gdpr_agreed is None:
return redirect(url_for('index.gdpr_notice', next=request.full_path))
return app
def setUp(self):
super(ServerTestCase, self).setUp()
#TODO: https://tickets.metabrainz.org/browse/BU-27
set_rate_limits(1000, 1000, 10000)
def set_rate_limits(per_token_limit, per_ip_limit, window_size):
from brainzutils.ratelimit import set_rate_limits
application = webserver.create_app()
with application.app_context():
set_rate_limits(per_token_limit, per_ip_limit, window_size)
def setUp(self):
self.reset_db()
# TODO: https://tickets.metabrainz.org/browse/BU-27
set_rate_limits(1000, 1000, 10000)