From f652be5e688556a183b8b868ad205889221ea79a Mon Sep 17 00:00:00 2001 From: Kanishk Pachauri Date: Fri, 29 Dec 2023 02:49:49 +0530 Subject: [PATCH 1/3] :sparkles: Feat: added middleware for Limit Upload Size --- src/paste/middleware.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 src/paste/middleware.py diff --git a/src/paste/middleware.py b/src/paste/middleware.py new file mode 100644 index 0000000..caf6bd5 --- /dev/null +++ b/src/paste/middleware.py @@ -0,0 +1,25 @@ +from starlette import status +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint +from starlette.requests import Request +from starlette.responses import Response +from starlette.types import ASGIApp + + +class LimitUploadSize(BaseHTTPMiddleware): + def __init__(self, app: ASGIApp, max_upload_size: int) -> None: + super().__init__(app) + self.max_upload_size = max_upload_size + + async def dispatch( + self, request: Request, call_next: RequestResponseEndpoint + ) -> Response: + if request.method == "POST": + if "content-length" not in request.headers: + return Response(status_code=status.HTTP_411_LENGTH_REQUIRED) + content_length = int(request.headers["content-length"]) + if content_length > self.max_upload_size: + return Response( + "File is too large", + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + ) + return await call_next(request) From 56a432cdc71f2ed045993129d4e92d3e6b09ceed Mon Sep 17 00:00:00 2001 From: Kanishk Pachauri Date: Fri, 29 Dec 2023 02:50:45 +0530 Subject: [PATCH 2/3] :heavy_plus_sign: Added middleware for fileuploadsize --- src/paste/main.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/paste/main.py b/src/paste/main.py index b60bb1b..0661b15 100644 --- a/src/paste/main.py +++ b/src/paste/main.py @@ -10,6 +10,7 @@ from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address from .utils import generate_uuid +from .middleware import LimitUploadSize limiter = Limiter(key_func=get_remote_address) app = FastAPI(title="paste.py 🐍") @@ -26,6 +27,8 @@ allow_headers=["*"], ) +app.add_middleware(LimitUploadSize, max_upload_size=20_000_000) # ~20MB + large_uuid_storage = [] BASE_DIR = Path(__file__).resolve().parent @@ -58,7 +61,7 @@ async def post_as_a_file(request: Request, file: UploadFile = File(...)): @app.get("/paste/{uuid}") -async def post_as_a_text(uuid): +async def get_paste_data(uuid): path = f"data/{uuid}" try: with open(path, "rb") as f: From 0a0b3dad2fc5b9de13cf39686feff6d529c7d2af Mon Sep 17 00:00:00 2001 From: Kanishk Pachauri Date: Fri, 29 Dec 2023 02:51:37 +0530 Subject: [PATCH 3/3] :white_check_mark: Test: Added test for file-size-upload --- tests/test_api.py | 46 +++++++++++++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 17 deletions(-) diff --git a/tests/test_api.py b/tests/test_api.py index 924e77d..38a968a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,5 +1,6 @@ from fastapi.testclient import TestClient from src.paste.main import app +import os client = TestClient(app) @@ -14,31 +15,29 @@ def test_get_health_route(): def test_get_homepage_route(): - response_expected_headers = 'text/html; charset=utf-8' + response_expected_headers = "text/html; charset=utf-8" response = client.get("/") assert response.status_code == 200 - assert response.headers.get( - 'Content-Type', '') == response_expected_headers + assert response.headers.get("Content-Type", "") == response_expected_headers def test_get_web_route(): - response_expected_headers = 'text/html; charset=utf-8' + response_expected_headers = "text/html; charset=utf-8" response = client.get("/web") assert response.status_code == 200 - assert response.headers.get( - 'Content-Type', '') == response_expected_headers + assert response.headers.get("Content-Type", "") == response_expected_headers -def test_get_paste_route(): - data = 'This is a test file.' +def test_get_paste_data_route(): + data = "This is a test file." response = client.get("/paste/test") assert response.status_code == 200 assert response.text == data def test_post_web_route(): - data = 'This is a test data' - form_data = {'content': data} + data = "This is a test data" + form_data = {"content": data} response = client.post("/web", data=form_data) global file file = str(response.url).split("/")[-1] @@ -54,8 +53,7 @@ def test_delete_paste_route(): def test_post_file_route(): - response = client.post( - "/file", files={"file": ("test.txt", b"test file content")}) + response = client.post("/file", files={"file": ("test.txt", b"test file content")}) assert response.status_code == 201 response_file_uuid = response.text response = client.get(f"/paste/{response_file_uuid}") @@ -73,13 +71,27 @@ def test_post_file_route_failure(): "detail": [ { "type": "missing", - "loc": [ - "body", - "file" - ], + "loc": ["body", "file"], "msg": "Field required", "input": None, - "url": "https://errors.pydantic.dev/2.5/v/missing" + "url": "https://errors.pydantic.dev/2.5/v/missing", } ] } + + +def test_post_file_route_size_limit(): + large_file_name = "large_file.txt" + file_size = 20 * 1024 * 1024 # 20 MB in bytes + additional_bytes = 100 # Adding some extra bytes to exceed 20 MB + content = b"This is a line in the file.\n" + with open(large_file_name, "wb") as file: + while file.tell() < file_size: + file.write(content) + file.write(b"Extra bytes to exceed 20 MB\n" * additional_bytes) + files = {"file": open(large_file_name, "rb")} + response = client.post("/file", files=files) + # cleanup + os.remove(large_file_name) + assert response.status_code == 413 + assert response.text == "File is too large"