From 731a41c652088b7a9b30ad4af0762370c3f1e0fc Mon Sep 17 00:00:00 2001 From: nac Date: Mon, 13 Jul 2026 01:55:59 +0000 Subject: [PATCH] feat: add authenticated attachment uploads --- nocobase_crm_mcp/uploads.py | 112 ++++++++++++++++++++++++++ tests/test_uploads.py | 155 ++++++++++++++++++++++++++++++++++++ 2 files changed, 267 insertions(+) create mode 100644 nocobase_crm_mcp/uploads.py create mode 100644 tests/test_uploads.py diff --git a/nocobase_crm_mcp/uploads.py b/nocobase_crm_mcp/uploads.py new file mode 100644 index 0000000..f6a4ee6 --- /dev/null +++ b/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) diff --git a/tests/test_uploads.py b/tests/test_uploads.py new file mode 100644 index 0000000..5f196e6 --- /dev/null +++ b/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()