feat: add authenticated attachment uploads
This commit is contained in:
nac
@@ -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)
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user