From b99ea03635aa83aab88c20083aa202f27d2d162c Mon Sep 17 00:00:00 2001 From: Kumi Date: Mon, 22 Jan 2024 08:36:34 +0100 Subject: [PATCH] Implement 5 MB file size limit enforcement Enhanced both FicheServer and LinesServer to reject data transfers exceeding a 5MB limit by introducing a max_size parameter across server classes, request handlers, and CLI arguments. This change improves system stability and security by preventing excessive resource use. Additionally, customized HTTP responses are provided for missing or invalid Content-Length headers and oversized files, creating a clearer client-server communication and ensuring better request validation logic. --- src/pyfiche/classes/fiche.py | 7 +++++++ src/pyfiche/classes/lines.py | 29 +++++++++++++++++++++++++---- src/pyfiche/fiche_server.py | 1 + src/pyfiche/lines_server.py | 1 + 4 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/pyfiche/classes/fiche.py b/src/pyfiche/classes/fiche.py index fc06355..5154c5f 100644 --- a/src/pyfiche/classes/fiche.py +++ b/src/pyfiche/classes/fiche.py @@ -21,6 +21,7 @@ class FicheServer: slug_size: int = 8 https: bool = False buffer_size: int = 4096 + max_size: int = 5242880 # 5 MB by default _output_dir: pathlib.Path = pathlib.Path('data/') _log_file: Optional[pathlib.Path] = None _banlist: Optional[pathlib.Path] = None @@ -92,6 +93,7 @@ class FicheServer: fiche.banlist = args.banlist or fiche.banlist fiche.allowlist = args.allowlist or fiche.allowlist fiche.buffer_size = args.buffer_size or fiche.buffer_size + fiche.max_size = args.max_size or fiche.max_size fiche.logger = logging.getLogger('pyfiche') fiche.logger.setLevel(logging.INFO if not args.debug else logging.DEBUG) @@ -181,6 +183,11 @@ class FicheServer: received_data.extend(data) + if len(received_data) > self.max_size: + self.logger.error(f"Received data exceeds maximum size ({self.max_size} bytes), terminating connection.") + conn.sendall(b"Data exceeds maximum size.\n") + return + except socket.timeout: pass diff --git a/src/pyfiche/classes/lines.py b/src/pyfiche/classes/lines.py index 28817d6..59cf4d1 100644 --- a/src/pyfiche/classes/lines.py +++ b/src/pyfiche/classes/lines.py @@ -67,7 +67,13 @@ body {{ # Reject any POST requests that don't have a Content-Length header if "Content-Length" not in self.headers: - return self.not_found() + return self.invalid_request() + + if not self.headers["Content-Length"].isdigit(): + return self.invalid_request() + + if int(self.headers["Content-Length"]) > self.max_size: + return self.file_too_large() # Upload the file @@ -91,7 +97,17 @@ body {{ self.send_response(303) self.send_header("Location", f"/{slug}") self.end_headers() - + + def invalid_request(self): + self.send_response(400) + self.end_headers() + self.wfile.write(b"Invalid request") + + def file_too_large(self): + self.send_response(413) + self.end_headers() + self.wfile.write(b"File too large") + def not_found(self): self.send_response(404) self.end_headers() @@ -211,13 +227,16 @@ body {{ self.wfile.write(full_html.encode("utf-8")) -def make_lines_handler(data_dir, logger, banlist=None, allowlist=None): +def make_lines_handler( + data_dir, logger, banlist=None, allowlist=None, max_size=5242880 +): class CustomHandler(LinesHTTPRequestHandler): def __init__(self, *args, **kwargs): self.data_dir: pathlib.Path = data_dir self.logger: logging.Logger = logger self.banlist: Optional[pathlib.Path] = banlist self.allowlist: Optional[pathlib.Path] = allowlist + self.max_size: int = max_size super().__init__(*args, **kwargs) @@ -227,6 +246,7 @@ def make_lines_handler(data_dir, logger, banlist=None, allowlist=None): class LinesServer: port: int = 9997 listen_addr: str = "0.0.0.0" + max_size: int = 5242880 # 5 MB by default _data_dir: pathlib.Path = pathlib.Path("data/") _log_file: Optional[pathlib.Path] = None _banlist: Optional[pathlib.Path] = None @@ -295,6 +315,7 @@ class LinesServer: lines.port = args.port or lines.port lines.listen_addr = args.listen_addr or lines.listen_addr + lines.max_size = args.max_size or lines.max_size lines.data_dir = args.data_dir or lines.data_dir lines.log_file = args.log_file or lines.log_file lines.banlist = args.banlist or lines.banlist @@ -316,7 +337,7 @@ class LinesServer: def run(self): handler_class = make_lines_handler( - self.data_dir, self.logger, self.banlist, self.allowlist + self.data_dir, self.logger, self.banlist, self.allowlist, self.max_size ) with HTTPServer((self.listen_addr, self.port), handler_class) as httpd: diff --git a/src/pyfiche/fiche_server.py b/src/pyfiche/fiche_server.py index e799eda..5eebf07 100644 --- a/src/pyfiche/fiche_server.py +++ b/src/pyfiche/fiche_server.py @@ -17,6 +17,7 @@ def main(): parser.add_argument('-S', '--https', action='store_true', help='HTTPS (requires reverse proxy)') parser.add_argument('-o', '--output_dir', help='Output directory path (default: data/)') parser.add_argument('-B', '--buffer_size', type=int, help='Buffer size (default: 4096)') + parser.add_argument('-M', '--max_size', type=int, help='Maximum file size (in bytes) (default: 5242880)') parser.add_argument('-l', '--log_file', help='Log file path (default: None - log to stdout)') parser.add_argument('-b', '--banlist', help='Banlist file path') parser.add_argument('-w', '--allowlist', help='Allowlist file path') diff --git a/src/pyfiche/lines_server.py b/src/pyfiche/lines_server.py index dbfa84b..669927a 100644 --- a/src/pyfiche/lines_server.py +++ b/src/pyfiche/lines_server.py @@ -16,6 +16,7 @@ def main(): parser.add_argument('-l', '--log_file', help='Log file path (default: None - log to stdout)') parser.add_argument('-b', '--banlist', help='Banlist file path') parser.add_argument('-w', '--allowlist', help='Allowlist file path') + parser.add_argument('-M', '--max_size', type=int, help='Maximum file size (in bytes) (default: 5242880)') parser.add_argument('-D', '--debug', action='store_true', help='Debug mode') # Parse the arguments