feat: add authenticated attachment uploads

nocobase_crm_mcp/uploads.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import fnmatch
+import tempfile
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from starlette.datastructures import UploadFile
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from .nocobase_http import NocoBaseHttpError
+
+
+_COPY_CHUNK_BYTES = 1024 * 1024
+
+
+class UploadEndpoint:
+    """ASGI endpoint that forwards one bounded multipart file to NocoBase."""
+
+    def __init__(
+        self,
+        http_client: Any,
+        max_bytes: int,
+        temp_dir: str | Path,
+        allowed_mime_types: Sequence[str] | None = None,
+    ) -> None:
+        self._http_client = http_client
+        self._max_bytes = max_bytes
+        self._temp_dir = Path(temp_dir)
+        self._allowed_mime_types = tuple(allowed_mime_types) if allowed_mime_types is not None else None
+
+    async def __call__(self, request: Request) -> JSONResponse:
+        if request.method != "POST":
+            return _error_response("invalid_upload", "Uploads require POST", 405)
+        if not _is_multipart(request):
+            return _error_response("invalid_upload", "Uploads require multipart form data", 400)
+
+        try:
+            form = await request.form()
+        except Exception:
+            return _error_response("invalid_upload", "Could not read upload form data", 400)
+
+        uploads = [(name, value) for name, value in form.multi_items() if isinstance(value, UploadFile)]
+        file_values = form.getlist("file")
+        if len(uploads) != 1 or len(file_values) != 1 or uploads[0][0] != "file":
+            await _close_uploads(value for _, value in uploads)
+            return _error_response("invalid_upload", "Provide exactly one file upload", 400)
+
+        upload = file_values[0]
+        if not isinstance(upload, UploadFile):
+            return _error_response("invalid_upload", "Provide exactly one file upload", 400)
+
+        filename = upload.filename or ""
+        mimetype = upload.content_type
+        if self._allowed_mime_types is not None and not any(
+            fnmatch.fnmatch(mimetype or "", allowed) for allowed in self._allowed_mime_types
+        ):
+            await upload.close()
+            return _error_response("unsupported_media_type", "This media type is not allowed", 415)
+
+        temporary_path: Path | None = None
+        try:
+            with tempfile.NamedTemporaryFile(delete=False, dir=self._temp_dir) as temporary:
+                temporary_path = Path(temporary.name)
+                size = 0
+                while chunk := await upload.read(_COPY_CHUNK_BYTES):
+                    size += len(chunk)
+                    if size > self._max_bytes:
+                        return _error_response("upload_too_large", "Upload exceeds the configured size limit", 413)
+                    temporary.write(chunk)
+
+            attachment = self._http_client.upload_attachment(temporary_path, filename, mimetype)
+            attachment_id = _attachment_id(attachment)
+            if attachment_id is None:
+                return _error_response("nb_upload_failed", "NocoBase attachment upload failed", 502)
+            return JSONResponse(
+                {"upload_id": attachment_id, "filename": filename, "size": size, "mimetype": mimetype}
+            )
+        except NocoBaseHttpError as error:
+            code = "nb_upload_unauthorized" if error.status_code in {401, 403} else "nb_upload_failed"
+            return _error_response(code, "NocoBase attachment upload failed", 502)
+        except Exception:
+            return _error_response("nb_upload_failed", "NocoBase attachment upload failed", 502)
+        finally:
+            await upload.close()
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
+
+
+def _is_multipart(request: Request) -> bool:
+    content_type = request.headers.get("content-type", "")
+    return content_type.split(";", 1)[0].strip().lower() == "multipart/form-data"
+
+
+def _attachment_id(attachment: Any) -> int | None:
+    if not isinstance(attachment, Mapping):
+        return None
+    attachment_id = attachment.get("id")
+    if isinstance(attachment_id, int) and not isinstance(attachment_id, bool):
+        return attachment_id
+    return None
+
+
+async def _close_uploads(uploads: Iterable[UploadFile]) -> None:
+    for upload in uploads:
+        await upload.close()
+
+
+def _error_response(code: str, message: str, status_code: int) -> JSONResponse:
+    return JSONResponse({"error": {"code": code, "message": message}}, status_code=status_code)

tests/test_uploads.py

@@ -0,0 +1,155 @@
+import gc
+import unittest
+import warnings
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import httpx
+from starlette.applications import Starlette
+from starlette.routing import Route
+
+from nocobase_crm_mcp.nocobase_http import NocoBaseHttpError
+from nocobase_crm_mcp.uploads import UploadEndpoint
+
+
+AUTH = {"Authorization": "Bearer secret"}
+
+
+class FakeHttpClient:
+    def __init__(self) -> None:
+        self.calls: list[tuple[Path, str, str | None, bytes]] = []
+        self.error: Exception | None = None
+
+    def upload_attachment(self, path: Path, filename: str, mimetype: str | None) -> dict[str, object]:
+        self.calls.append((path, filename, mimetype, path.read_bytes()))
+        if self.error:
+            raise self.error
+        return {"id": 42}
+
+
+class UploadEndpointTests(unittest.IsolatedAsyncioTestCase):
+    async def _client_for(
+        self,
+        http_client: FakeHttpClient,
+        temp_dir: str,
+        *,
+        max_bytes: int = 5,
+        allowed_mime_types: tuple[str, ...] | None = None,
+    ) -> httpx.AsyncClient:
+        endpoint = UploadEndpoint(http_client, max_bytes, temp_dir, allowed_mime_types)
+
+        async def upload_route(request):
+            return await endpoint(request)
+
+        app = Starlette(
+            routes=[
+                Route(
+                    "/uploads",
+                    upload_route,
+                    methods=["GET", "POST"],
+                )
+            ]
+        )
+        return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver")
+
+    async def test_upload_returns_nocobase_attachment_id_and_removes_temp_file(self) -> None:
+        http_client = FakeHttpClient()
+        with TemporaryDirectory() as temp_dir:
+            async with await self._client_for(http_client, temp_dir) as client:
+                response = await client.post(
+                    "/uploads",
+                    headers=AUTH,
+                    files={"file": ("quote.pdf", b"pdf", "application/pdf")},
+                )
+
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(
+                response.json(),
+                {"upload_id": 42, "filename": "quote.pdf", "size": 3, "mimetype": "application/pdf"},
+            )
+            self.assertEqual(http_client.calls[0][1:], ("quote.pdf", "application/pdf", b"pdf"))
+            self.assertEqual(list(Path(temp_dir).iterdir()), [])
+
+    async def test_upload_rejects_missing_multiple_and_over_limit_files(self) -> None:
+        http_client = FakeHttpClient()
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always", ResourceWarning)
+            with TemporaryDirectory() as temp_dir:
+                async with await self._client_for(http_client, temp_dir) as client:
+                    missing = await client.post("/uploads", headers=AUTH, files={})
+                    multiple = await client.post(
+                        "/uploads",
+                        headers=AUTH,
+                        files=[
+                            ("file", ("one.pdf", b"one", "application/pdf")),
+                            ("file", ("two.pdf", b"two", "application/pdf")),
+                        ],
+                    )
+                    over_limit = await client.post(
+                        "/uploads",
+                        headers=AUTH,
+                        files={"file": ("big.bin", b"x" * 6, "application/octet-stream")},
+                    )
+                temp_paths = list(Path(temp_dir).iterdir())
+            gc.collect()
+
+        self.assertEqual(missing.status_code, 400)
+        self.assertEqual(missing.json()["error"]["code"], "invalid_upload")
+        self.assertEqual(multiple.status_code, 400)
+        self.assertEqual(multiple.json()["error"]["code"], "invalid_upload")
+        self.assertEqual(over_limit.status_code, 413)
+        self.assertEqual(over_limit.json()["error"]["code"], "upload_too_large")
+        self.assertEqual(http_client.calls, [])
+        self.assertEqual(temp_paths, [])
+        self.assertFalse(any(issubclass(warning.category, ResourceWarning) for warning in caught))
+
+    async def test_upload_rejects_disallowed_mime_type_before_calling_nocobase(self) -> None:
+        http_client = FakeHttpClient()
+        with TemporaryDirectory() as temp_dir:
+            async with await self._client_for(
+                http_client, temp_dir, allowed_mime_types=("image/*", "application/pdf")
+            ) as client:
+                response = await client.post(
+                    "/uploads",
+                    headers=AUTH,
+                    files={"file": ("notes.txt", b"notes", "text/plain")},
+                )
+
+            self.assertEqual(response.status_code, 415)
+            self.assertEqual(response.json()["error"]["code"], "unsupported_media_type")
+            self.assertEqual(http_client.calls, [])
+            self.assertEqual(list(Path(temp_dir).iterdir()), [])
+
+    async def test_upload_maps_upstream_errors_and_removes_temp_file(self) -> None:
+        cases = ((401, "nb_upload_unauthorized"), (403, "nb_upload_unauthorized"), (500, "nb_upload_failed"))
+
+        for status_code, error_code in cases:
+            with self.subTest(status_code=status_code), TemporaryDirectory() as temp_dir:
+                http_client = FakeHttpClient()
+                http_client.error = NocoBaseHttpError(status_code, "upstream failure")
+                async with await self._client_for(http_client, temp_dir) as client:
+                    response = await client.post(
+                        "/uploads",
+                        headers=AUTH,
+                        files={"file": ("quote.pdf", b"pdf", "application/pdf")},
+                    )
+
+                self.assertEqual(response.status_code, 502)
+                self.assertEqual(response.json()["error"]["code"], error_code)
+                self.assertEqual(list(Path(temp_dir).iterdir()), [])
+
+    async def test_upload_requires_post_and_multipart_form_data(self) -> None:
+        http_client = FakeHttpClient()
+        with TemporaryDirectory() as temp_dir:
+            async with await self._client_for(http_client, temp_dir) as client:
+                get_response = await client.get("/uploads", headers=AUTH)
+                json_response = await client.post("/uploads", headers=AUTH, json={"file": "not-a-file"})
+
+            self.assertEqual(get_response.status_code, 405)
+            self.assertEqual(json_response.status_code, 400)
+            self.assertEqual(json_response.json()["error"]["code"], "invalid_upload")
+            self.assertEqual(list(Path(temp_dir).iterdir()), [])
+
+
+if __name__ == "__main__":
+    unittest.main()