mirror of https://gitlab.com/tildes/tildes.git
Deimos
5 years ago
2 changed files with 146 additions and 134 deletions
@ -0,0 +1,144 @@ |
|||
# Copyright (c) 2019 Tildes contributors <code@tildes.net> |
|||
# SPDX-License-Identifier: AGPL-3.0-or-later |
|||
|
|||
"""Define and attach request methods to the Pyramid request object.""" |
|||
|
|||
from typing import Any, Dict, Optional, Tuple |
|||
|
|||
from pyramid.config import Configurator |
|||
from pyramid.httpexceptions import HTTPTooManyRequests |
|||
from pyramid.request import Request |
|||
from redis import Redis |
|||
|
|||
from tildes.lib.ratelimit import RATE_LIMITED_ACTIONS, RateLimitResult |
|||
|
|||
|
|||
def get_redis_connection(request: Request) -> Redis: |
|||
"""Return a connection to the Redis server.""" |
|||
socket = request.registry.settings["redis.unix_socket_path"] |
|||
return Redis(unix_socket_path=socket) |
|||
|
|||
|
|||
def is_safe_request_method(request: Request) -> bool: |
|||
"""Return whether the request method is "safe" (is GET or HEAD).""" |
|||
return request.method in {"GET", "HEAD"} |
|||
|
|||
|
|||
def check_rate_limit(request: Request, action_name: str) -> RateLimitResult: |
|||
"""Check the rate limit for a particular action on a request.""" |
|||
try: |
|||
action = RATE_LIMITED_ACTIONS[action_name] |
|||
except KeyError: |
|||
raise ValueError("Invalid action name: %s" % action_name) |
|||
|
|||
action.redis = request.redis |
|||
|
|||
results = [] |
|||
|
|||
if action.by_user and request.user: |
|||
results.append(action.check_for_user_id(request.user.user_id)) |
|||
|
|||
if action.by_ip and request.remote_addr: |
|||
results.append(action.check_for_ip(request.remote_addr)) |
|||
|
|||
# no checks were done, return the "not limited" result |
|||
if not results: |
|||
return RateLimitResult.unlimited_result() |
|||
|
|||
return RateLimitResult.merged_result(results) |
|||
|
|||
|
|||
def apply_rate_limit(request: Request, action_name: str) -> None: |
|||
"""Check the rate limit for an action, and raise HTTP 429 if it's exceeded.""" |
|||
result = request.check_rate_limit(action_name) |
|||
if not result.is_allowed: |
|||
raise result.add_headers_to_response(HTTPTooManyRequests()) |
|||
|
|||
|
|||
def current_listing_base_url( |
|||
request: Request, query: Optional[Dict[str, Any]] = None |
|||
) -> str: |
|||
"""Return the "base" url for the current listing route. |
|||
|
|||
The "base" url represents the current listing, including any filtering options (or |
|||
the fact that filters are disabled). |
|||
|
|||
The `query` argument allows adding query variables to the generated url. |
|||
""" |
|||
base_vars_by_route: Dict[str, Tuple[str, ...]] = { |
|||
"bookmarks": ("per_page", "type"), |
|||
"group": ("order", "period", "per_page", "tag", "unfiltered"), |
|||
"group_search": ("order", "period", "per_page", "q"), |
|||
"home": ("order", "period", "per_page", "tag", "unfiltered"), |
|||
"search": ("order", "period", "per_page", "q"), |
|||
"user": ("order", "per_page", "type"), |
|||
} |
|||
|
|||
try: |
|||
base_view_vars = base_vars_by_route[request.matched_route.name] |
|||
except KeyError: |
|||
raise AttributeError("Current route is not supported.") |
|||
|
|||
query_vars = { |
|||
key: val for key, val in request.GET.copy().items() if key in base_view_vars |
|||
} |
|||
if query: |
|||
query_vars.update(query) |
|||
|
|||
return request.current_route_url(_query=query_vars) |
|||
|
|||
|
|||
def current_listing_normal_url( |
|||
request: Request, query: Optional[Dict[str, Any]] = None |
|||
) -> str: |
|||
"""Return the "normal" url for the current listing route. |
|||
|
|||
The "normal" url represents the current listing without any additional |
|||
filtering-related changes (the user's standard view of that listing). |
|||
|
|||
The `query` argument allows adding query variables to the generated url. |
|||
""" |
|||
normal_vars_by_route: Dict[str, Tuple[str, ...]] = { |
|||
"bookmarks": ("order", "period", "per_page"), |
|||
"group": ("order", "period", "per_page"), |
|||
"group_search": ("order", "period", "per_page", "q"), |
|||
"home": ("order", "period", "per_page"), |
|||
"notifications": ("per_page",), |
|||
"search": ("order", "period", "per_page", "q"), |
|||
"user": ("order", "per_page"), |
|||
} |
|||
|
|||
try: |
|||
normal_view_vars = normal_vars_by_route[request.matched_route.name] |
|||
except KeyError: |
|||
raise AttributeError("Current route is not supported.") |
|||
|
|||
query_vars = { |
|||
key: val for key, val in request.GET.copy().items() if key in normal_view_vars |
|||
} |
|||
if query: |
|||
query_vars.update(query) |
|||
|
|||
return request.current_route_url(_query=query_vars) |
|||
|
|||
|
|||
def includeme(config: Configurator) -> None: |
|||
"""Attach the request methods to the Pyramid request object.""" |
|||
config.add_request_method(is_safe_request_method, "is_safe_method", reify=True) |
|||
|
|||
# Add the request.redis request method to access a redis connection. This is done in |
|||
# a bit of a strange way to support being overridden in tests. |
|||
config.registry["redis_connection_factory"] = get_redis_connection |
|||
# pylint: disable=unnecessary-lambda |
|||
config.add_request_method( |
|||
lambda request: config.registry["redis_connection_factory"](request), |
|||
"redis", |
|||
reify=True, |
|||
) |
|||
# pylint: enable=unnecessary-lambda |
|||
|
|||
config.add_request_method(check_rate_limit, "check_rate_limit") |
|||
config.add_request_method(apply_rate_limit, "apply_rate_limit") |
|||
|
|||
config.add_request_method(current_listing_base_url, "current_listing_base_url") |
|||
config.add_request_method(current_listing_normal_url, "current_listing_normal_url") |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue