Merge pull request #31 from Mr-Sunglasses/testss

feat: file upload limit

Author: Mr-Sunglasses

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:

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)

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"