From 2ef2a55a95b350672d72cf5effeb0e33e0442c09 Mon Sep 17 00:00:00 2001
From: ghandic <andy.challis@capgemini.com>
Date: Sat, 12 Sep 2020 14:48:22 +1000
Subject: [PATCH] Added in exception handling into router to allow execptions
 raised by middleware to surface up through the use of HTTPException, tests to
 complement it as added additional exception branching

---
 fastapi/routing.py                        |  2 +
 tests/test_custom_middleware_exception.py | 95 +++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 tests/test_custom_middleware_exception.py

diff --git a/fastapi/routing.py b/fastapi/routing.py
index e455f81c95387..9680b72b3dc03 100644
--- a/fastapi/routing.py
+++ b/fastapi/routing.py
@@ -165,6 +165,8 @@ async def app(request: Request) -> Response:
                         body = await request.json()
         except json.JSONDecodeError as e:
             raise RequestValidationError([ErrorWrapper(e, ("body", e.pos))], body=e.doc)
+        except HTTPException:
+            raise
         except Exception as e:
             raise HTTPException(
                 status_code=400, detail="There was an error parsing the body"
diff --git a/tests/test_custom_middleware_exception.py b/tests/test_custom_middleware_exception.py
new file mode 100644
index 0000000000000..9b4a7a9a99ceb
--- /dev/null
+++ b/tests/test_custom_middleware_exception.py
@@ -0,0 +1,95 @@
+import os
+from typing import Optional
+
+from fastapi import APIRouter, FastAPI, File, UploadFile
+from fastapi.exceptions import HTTPException
+from fastapi.testclient import TestClient
+
+app = FastAPI()
+
+router = APIRouter()
+
+
+class ContentSizeLimitMiddleware:
+    """ Content size limiting middleware for ASGI applications
+    Args:
+      app (ASGI application): ASGI application
+      max_content_size (optional): the maximum content size allowed in bytes, None for no limit
+    """
+
+    def __init__(self, app: APIRouter, max_content_size: Optional[int] = None):
+        self.app = app
+        self.max_content_size = max_content_size
+
+    def receive_wrapper(self, receive):
+        received = 0
+
+        async def inner():
+            nonlocal received
+            message = await receive()
+            if message["type"] != "http.request":
+                return message  # pragma: no cover
+
+            body_len = len(message.get("body", b""))
+            received += body_len
+            if received > self.max_content_size:
+                raise HTTPException(
+                    422,
+                    detail={
+                        "name": "ContentSizeLimitExceeded",
+                        "code": 999,
+                        "message": "File limit exceeded",
+                    },
+                )
+            return message
+
+        return inner
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http" or self.max_content_size is None:
+            await self.app(scope, receive, send)
+            return
+
+        wrapper = self.receive_wrapper(receive)
+        await self.app(scope, wrapper, send)
+
+
+@router.post("/middleware")
+def run_middleware(file: UploadFile = File(..., description="Big File")):
+    return {"message": "OK"}
+
+
+app.include_router(router)
+app.add_middleware(ContentSizeLimitMiddleware, max_content_size=2 ** 8)
+
+
+client = TestClient(app)
+
+
+def test_custom_middleware_exception(tmpdir):
+    default_pydantic_max_size = 2 ** 16
+    path = os.path.join(tmpdir, "test.txt")
+    with open(path, "wb") as file:
+        file.write(b"x" * (default_pydantic_max_size + 1))
+
+    with client:
+        response = client.post("/middleware", files={"file": open(path, "rb")})
+        assert response.status_code == 422, response.text
+        assert response.json() == {
+            "detail": {
+                "name": "ContentSizeLimitExceeded",
+                "code": 999,
+                "message": "File limit exceeded",
+            }
+        }
+
+
+def test_custom_middleware_exception_not_raised(tmpdir):
+    path = os.path.join(tmpdir, "test.txt")
+    with open(path, "wb") as file:
+        file.write(b"<file content>")
+
+    with client:
+        response = client.post("/middleware", files={"file": open(path, "rb")})
+        assert response.status_code == 200, response.text
+        assert response.json() == {"message": "OK"}