Verified Commit 65240c83 authored by Kevin Morris's avatar Kevin Morris
Browse files

feat(rpc): enforce ratelimiting



New configuration options:

- `[ratelimit] cache`
    - A boolean indicating whether we should use configured cache (1)
      or database (0) for ratelimiting.
Signed-off-by: Kevin Morris's avatarKevin Morris <kevr@0cost.org>
parent 66629750
Pipeline #12410 passed with stage
in 7 minutes and 46 seconds
from datetime import datetime
from fastapi import Request
from redis.client import Pipeline
from aurweb import config, db, logging
from aurweb.models import ApiRateLimit
from aurweb.redis import redis_connection
logger = logging.get_logger(__name__)
def _update_ratelimit_redis(request: Request, pipeline: Pipeline):
window_length = config.getint("ratelimit", "window_length")
now = int(datetime.utcnow().timestamp())
time_to_delete = now - window_length
host = request.client.host
window_key = f"ratelimit-ws:{host}"
requests_key = f"ratelimit:{host}"
pipeline.get(window_key)
window = pipeline.execute()[0]
if not window or int(window.decode()) < time_to_delete:
pipeline.set(window_key, now)
pipeline.expire(window_key, window_length)
pipeline.set(requests_key, 1)
pipeline.expire(requests_key, window_length)
pipeline.execute()
else:
pipeline.incr(requests_key)
pipeline.execute()
def _update_ratelimit_db(request: Request):
window_length = config.getint("ratelimit", "window_length")
now = int(datetime.utcnow().timestamp())
time_to_delete = now - window_length
with db.begin():
db.delete(ApiRateLimit, ApiRateLimit.WindowStart < time_to_delete)
host = request.client.host
record = db.query(ApiRateLimit, ApiRateLimit.IP == host).first()
with db.begin():
if not record:
record = db.create(ApiRateLimit,
WindowStart=now,
IP=host, Requests=1)
else:
record.Requests += 1
logger.debug(record.Requests)
return record
def update_ratelimit(request: Request, pipeline: Pipeline):
""" Update the ratelimit stored in Redis or the database depending
on AUR_CONFIG's [options] cache setting.
This Redis-capable function is slightly different than most. If Redis
is not configured to use a real server, this function instead uses
the database to persist tracking of a particular host.
:param request: FastAPI request
:param pipeline: redis.client.Pipeline
:returns: ApiRateLimit record when Redis cache is not configured, else None
"""
if config.getboolean("ratelimit", "cache"):
return _update_ratelimit_redis(request, pipeline)
return _update_ratelimit_db(request)
def check_ratelimit(request: Request):
""" Increment and check to see if request has exceeded their rate limit.
:param request: FastAPI request
:returns: True if the request host has exceeded the rate limit else False
"""
redis = redis_connection()
pipeline = redis.pipeline()
record = update_ratelimit(request, pipeline)
# Get cache value, else None.
host = request.client.host
pipeline.get(f"ratelimit:{host}")
requests = pipeline.execute()[0]
# Take into account the split paths. When Redis is used, a
# valid cache value will be returned which must be converted
# to an int. Otherwise, use the database record returned
# by update_ratelimit.
if not config.getboolean("ratelimit", "cache"):
# If we got nothing from pipeline.get, we did not use
# the Redis path of logic: use the DB record's count.
requests = record.Requests
else:
# Otherwise, just case Redis results over to an int.
requests = int(requests.decode())
limit = config.getint("ratelimit", "request_limit")
exceeded_ratelimit = requests > limit
if exceeded_ratelimit:
logger.debug(f"{host} has exceeded the ratelimit.")
return exceeded_ratelimit
from http import HTTPStatus
from typing import List, Optional
from urllib.parse import unquote
from fastapi import APIRouter, Query, Request
from fastapi.responses import JSONResponse
from aurweb.ratelimit import check_ratelimit
from aurweb.rpc import RPC
router = APIRouter()
......@@ -64,6 +66,11 @@ async def rpc(request: Request,
# Create a handle to our RPC class.
rpc = RPC(version=v, type=type)
# If ratelimit was exceeded, return a 429 Too Many Requests.
if check_ratelimit(request):
return JSONResponse(rpc.error("Rate limit reached"),
status_code=int(HTTPStatus.TOO_MANY_REQUESTS))
# Prepare list of arguments for input. If 'arg' was given, it'll
# be a list with one element.
arguments = parse_args(request)
......
......@@ -48,6 +48,11 @@ redis_address = redis://localhost
[ratelimit]
request_limit = 4000
window_length = 86400
; Force-utilize cache for ratelimiting. In FastAPI, forced cache (1)
; will cause the ratelimit path to use a real or fake Redis instance
; depending on the configured options.cache setting. Otherwise,
; cache will be ignored and the database will be used.
cache = 1
[notifications]
notify-cmd = /usr/bin/aurweb-notify
......
from unittest import mock
import pytest
from redis.client import Pipeline
from aurweb import config, db, logging
from aurweb.models import ApiRateLimit
from aurweb.ratelimit import check_ratelimit
from aurweb.redis import redis_connection
from aurweb.testing import setup_test_db
from aurweb.testing.requests import Request
logger = logging.get_logger(__name__)
@pytest.fixture(autouse=True)
def setup():
setup_test_db(ApiRateLimit.__tablename__)
@pytest.fixture
def pipeline():
redis = redis_connection()
pipeline = redis.pipeline()
pipeline.delete("ratelimit-ws:127.0.0.1")
pipeline.delete("ratelimit:127.0.0.1")
pipeline.execute()
yield pipeline
def mock_config_getint(section: str, key: str):
if key == "request_limit":
return 4
elif key == "window_length":
return 100
return config.getint(section, key)
def mock_config_getboolean(return_value: int = 0):
def fn(section: str, key: str):
if section == "ratelimit" and key == "cache":
return return_value
return config.getboolean(section, key)
return fn
def mock_config_get(return_value: str = "none"):
def fn(section: str, key: str):
if section == "options" and key == "cache":
return return_value
return config.get(section, key)
return fn
@mock.patch("aurweb.config.getint", side_effect=mock_config_getint)
@mock.patch("aurweb.config.getboolean", side_effect=mock_config_getboolean(1))
@mock.patch("aurweb.config.get", side_effect=mock_config_get("none"))
def test_ratelimit_redis(get: mock.MagicMock, getboolean: mock.MagicMock,
getint: mock.MagicMock, pipeline: Pipeline):
""" This test will only cover aurweb.ratelimit's Redis
path if a real Redis server is configured. Otherwise,
it'll use the database. """
# We'll need a Request for everything here.
request = Request()
# Run check_ratelimit for our request_limit. These should succeed.
for i in range(4):
assert not check_ratelimit(request)
# This check_ratelimit should fail, being the 4001th request.
assert check_ratelimit(request)
# Delete the Redis keys.
host = request.client.host
pipeline.delete(f"ratelimit-ws:{host}")
pipeline.delete(f"ratelimit:{host}")
one, two = pipeline.execute()
assert one and two
# Should be good to go again!
assert not check_ratelimit(request)
@mock.patch("aurweb.config.getint", side_effect=mock_config_getint)
@mock.patch("aurweb.config.getboolean", side_effect=mock_config_getboolean(0))
@mock.patch("aurweb.config.get", side_effect=mock_config_get("none"))
def test_ratelimit_db(get: mock.MagicMock, getboolean: mock.MagicMock,
getint: mock.MagicMock, pipeline: Pipeline):
# We'll need a Request for everything here.
request = Request()
# Run check_ratelimit for our request_limit. These should succeed.
for i in range(4):
assert not check_ratelimit(request)
# This check_ratelimit should fail, being the 4001th request.
assert check_ratelimit(request)
# Delete the ApiRateLimit record.
with db.begin():
db.delete(ApiRateLimit)
# Should be good to go again!
assert not check_ratelimit(request)
from http import HTTPStatus
from unittest import mock
import orjson
import pytest
from fastapi.testclient import TestClient
from redis.client import Pipeline
from aurweb import db, scripts
from aurweb import config, db, scripts
from aurweb.asgi import app
from aurweb.db import begin, create, query
from aurweb.models.account_type import AccountType
......@@ -18,6 +22,7 @@ from aurweb.models.package_relation import PackageRelation
from aurweb.models.package_vote import PackageVote
from aurweb.models.relation_type import RelationType
from aurweb.models.user import User
from aurweb.redis import redis_connection
from aurweb.testing import setup_test_db
......@@ -31,7 +36,7 @@ def setup():
# Set up tables.
setup_test_db("Users", "PackageBases", "Packages", "Licenses",
"PackageDepends", "PackageRelations", "PackageLicenses",
"PackageKeywords", "PackageVotes")
"PackageKeywords", "PackageVotes", "ApiRateLimit")
# Create test package details.
with begin():
......@@ -178,6 +183,18 @@ def setup():
scripts.popupdate.run_single(conn, pkgbase1)
@pytest.fixture
def pipeline():
redis = redis_connection()
pipeline = redis.pipeline()
pipeline.delete("ratelimit-ws:testclient")
pipeline.delete("ratelimit:testclient")
one, two = pipeline.execute()
yield pipeline
def test_rpc_singular_info():
# Define expected response.
expected_data = {
......@@ -441,3 +458,33 @@ def test_rpc_unimplemented_types():
data = response.json()
expected = f"Request type '{type}' is not yet implemented."
assert data.get("error") == expected
def mock_config_getint(section: str, key: str):
if key == "request_limit":
return 4
elif key == "window_length":
return 100
return config.getint(section, key)
@mock.patch("aurweb.config.getint", side_effect=mock_config_getint)
def test_rpc_ratelimit(getint: mock.MagicMock, pipeline: Pipeline):
for i in range(4):
# The first 4 requests should be good.
response = make_request("/rpc?v=5&type=suggest-pkgbase&arg=big")
assert response.status_code == int(HTTPStatus.OK)
# The fifth request should be banned.
response = make_request("/rpc?v=5&type=suggest-pkgbase&arg=big")
assert response.status_code == int(HTTPStatus.TOO_MANY_REQUESTS)
# Delete the cached records.
pipeline.delete("ratelimit-ws:testclient")
pipeline.delete("ratelimit:testclient")
one, two = pipeline.execute()
assert one and two
# The new first request should be good.
response = make_request("/rpc?v=5&type=suggest-pkgbase&arg=big")
assert response.status_code == int(HTTPStatus.OK)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment