From 019b0477511f0ecc791b86a9c35f575661561ca0 Mon Sep 17 00:00:00 2001 From: Pat Patterson Date: Thu, 8 Aug 2024 12:43:13 -0700 Subject: [PATCH] Add random rate limit responses to embedded webserver --- b2listen/b2listen.py | 11 +++++++++-- b2listen/server.py | 24 ++++++++++++++++++------ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/b2listen/b2listen.py b/b2listen/b2listen.py index b31c3c4..994fe57 100644 --- a/b2listen/b2listen.py +++ b/b2listen/b2listen.py @@ -193,7 +193,13 @@ def parse_args() -> argparse.Namespace: group.add_argument('--url', type=str, help=f'Local webserver URL, for example: "http://localhost:8080")') # noqa group.add_argument('--run-server', action='store_true', - help=f'Run the embedded HTTP server') + help=f'Run the embedded webserver') + + run_server = parser_listen.add_argument_group(description='To simulate rate limiting in the embedded webserver:') + run_server.add_argument('--rate-limit-frequency', type=int, required=False, + help='Respond with "429 Too Many Requests" for this percentage of notifications.') + run_server.add_argument('--retry-after', type=int, required=False, + help='Value for the "Retry-After" header in a rate limit response.') use_existing = parser_listen.add_argument_group(description='To use an existing Event Notification rule:') use_existing.add_argument('--rule-name', type=str, required=False, @@ -320,7 +326,8 @@ def run_cloudflared(command: str, loglevel: str, service_url: str, label: str, u def listen(args: argparse.Namespace): if args.run_server: - http_server = Server(interface='localhost', port=0, daemon=True) + http_server = Server(interface='localhost', port=0, daemon=True, + rate_limit_frequency=args.rate_limit_frequency, retry_after=args.retry_after) http_server.start() service_url = f'http://{http_server.interface}:{http_server.port}' # noqa else: diff --git a/b2listen/server.py b/b2listen/server.py index b3eecf2..15d0adc 100644 --- a/b2listen/server.py +++ b/b2listen/server.py @@ -9,11 +9,15 @@ From https://gist.github.com/mdonkers/63e115cc0c79b4f6b8b3a6b797e485c7 """ +import random +from http import HTTPStatus from http.server import BaseHTTPRequestHandler, HTTPServer from sys import argv import logging from threading import Thread +RETRY_AFTER = 'Retry-After' + logging.basicConfig() logger = logging.getLogger('b2listen.server') @@ -22,15 +26,19 @@ class S(BaseHTTPRequestHandler): - def _set_response(self): - self.send_response(200) - self.send_header('Content-type', 'text/html') + rate_limit_frequency = 0 + retry_after = 0 + + def _set_response(self, status_code): + self.send_response(status_code) + if status_code == HTTPStatus.TOO_MANY_REQUESTS: + self.send_header(RETRY_AFTER, str(self.retry_after)) self.end_headers() # noinspection PyPep8Naming def do_GET(self): logger.info("GET request,\nPath: %s\nHeaders:\n%s\n", str(self.path), str(self.headers)) - self._set_response() + self._set_response(HTTPStatus.OK) self.wfile.write("GET request for {}".format(self.path).encode('utf-8')) # noinspection PyPep8Naming @@ -40,19 +48,23 @@ def do_POST(self): logger.info("POST request,\nPath: %s\nHeaders:\n%s\nBody:\n%s\n", str(self.path), str(self.headers), post_data.decode('utf-8')) - self._set_response() + status_code = HTTPStatus.OK if random.random() > (self.rate_limit_frequency / 100) \ + else HTTPStatus.TOO_MANY_REQUESTS + self._set_response(status_code) self.wfile.write("POST request for {}".format(self.path).encode('utf-8')) class Server(Thread): def __init__(self, server_class=HTTPServer, handler_class=S, interface=DEFAULT_INTERFACE, port=DEFAULT_PORT, - daemon=False): + daemon=False, rate_limit_frequency=0, retry_after=0): super().__init__(daemon=daemon) server_address = (interface, port) # noinspection PyTypeChecker self.httpd = server_class(server_address, handler_class) self.interface = self.httpd.server_address[0] self.port = self.httpd.server_address[1] + S.rate_limit_frequency = rate_limit_frequency + S.retry_after = retry_after def run(self): logger.info(f'Starting HTTP server on {self.interface}:{self.port}')