diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 054fa2c..e2ff2b1 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -11,9 +11,11 @@ jobs: runs-on: self-hosted steps: + - name: Create data Folder + run: mkdir -p data - name: Pull Docker image run: sudo docker pull mrsunglasses/pastepy - name: Delete Old docker container run: sudo docker rm -f pastepyprod || true - name: Run Docker Container - run: sudo docker run -d -p 8080:8080 --name pastepyprod mrsunglasses/pastepy + run: sudo docker run -d -p 8080:8080 -v $(pwd)/data:/project/data --name pastepyprod mrsunglasses/pastepy diff --git a/README.md b/README.md index 4c0c58e..9d7b992 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# paste.py +# paste.py 🐍 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: 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) diff --git a/tests/test_api.py b/tests/test_api.py index c83a236..38a968a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,8 +1,11 @@ from fastapi.testclient import TestClient from src.paste.main import app +import os client = TestClient(app) +file = None + def test_get_health_route(): data = {"status": "ok"} @@ -12,38 +15,83 @@ 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] assert response.status_code == 200 assert response.text == data def test_delete_paste_route(): - expected_response = "File successfully deleted test" - response = client.delete("/paste/test") + expected_response = f"File successfully deleted {file}" + response = client.delete(f"/paste/{file}") assert response.status_code == 200 assert response.text == expected_response + + +def test_post_file_route(): + 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}") + assert response.status_code == 200 + assert response.text == "test file content" + response = client.delete(f"/paste/{response_file_uuid}") + assert response.status_code == 200 + assert response.text == f"File successfully deleted {response_file_uuid}" + + +def test_post_file_route_failure(): + response = client.post("/file") + assert response.status_code == 422 # Unprocessable Entity + assert response.json() == { + "detail": [ + { + "type": "missing", + "loc": ["body", "file"], + "msg": "Field required", + "input": None, + "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"