Limit number of uploaded images

muckraker/main.py

@@ -13,9 +13,11 @@
 from .models import Issue
 from .render import render_issue
 
+MAX_IMAGE_NUM = 4
 MAX_IMAGE_SIZE = 2 * 1024 * 1024  # 2 MB
 IMAGE_BATCH = 1024
 ACCEPTED_FILE_TYPES = ("image/png", "image/jpeg", "image/jpg")
+IMAGE_SUFFIXES = (".png", ".jpeg", ".jpg")
 
 app = FastAPI(root_path="/api")
 origins = ["*"]
@@ -47,6 +49,13 @@ async def patch_s_issue(
     dir_path: Path = Depends(dir_path),
     image: UploadFile = File()
 ):
+    # Check already uploaded images num
+    files = dir_path.glob('**/*')
+    images = [x for x in files if x.is_file() and x.suffix in IMAGE_SUFFIXES]
+    if len(images) + 1 >= MAX_IMAGE_NUM:
+        rmtree(dir_path)
+        raise HTTPException(429, detail="To many uploads")
+
     # Validate image
     if image.content_type not in ACCEPTED_FILE_TYPES:
         detail = f"Invalid file type: {image.filename}"
@@ -62,7 +71,11 @@ async def patch_s_issue(
     async with aiofiles.open(image_path, "wb") as fd:
         while content := await image.read(IMAGE_BATCH):
             await fd.write(content)
-    return JSONResponse(content="Image uploaded")
+
+    files = dir_path.glob('**/*')
+    images = [x for x in files if x.is_file() and x.suffix in IMAGE_SUFFIXES]
+    print(len(images))
+    return JSONResponse(content={"filename": image.filename})
 
 
 @app.get("/issue/{issue_id}")

tests/test_api.py

@@ -53,12 +53,22 @@ def test_thick_image(issue_dict, thick_image):
 
 def test_issue_not_found(issue_dict, good_image):
     resp = client.post("/issue/", json=issue_dict)
-    issue_id = resp.json()["issue_id"]  # Deletes tempdir
-    client.get(f"/issue/{issue_id}")
+    issue_id = resp.json()["issue_id"]
+    client.get(f"/issue/{issue_id}")  # Deletes tempdir
 
     resp = client.get(f"/issue/{issue_id}")
     assert resp.status_code == 404
 
     files = {"image": good_image}
     resp = client.patch(f"/issue/{issue_id}", files=files)
     assert resp.status_code == 404
+
+
+def test_too_many_images(issue_dict, good_image):
+    resp = client.post("/issue/", json=issue_dict)
+    issue_id = resp.json()["issue_id"]
+
+    for no in range(1, 5):
+        files = {"image": (f"{no}.png", good_image)}
+        resp = client.patch(f"/issue/{issue_id}", files=files)
+    assert resp.status_code == 429