diff --git a/lnbits/core/wasm_ext/routes/api.py b/lnbits/core/wasm_ext/routes/api.py index 48d7abfb6..715e65fdf 100644 --- a/lnbits/core/wasm_ext/routes/api.py +++ b/lnbits/core/wasm_ext/routes/api.py @@ -2,18 +2,31 @@ from __future__ import annotations import json import re +from dataclasses import dataclass from typing import Annotated, Any from fastapi import Depends, FastAPI, HTTPException, Request from lnbits.core.models import Account +from lnbits.core.services.extensions import get_wasm_runtime_limits_for_extension from lnbits.decorators import check_access_token, check_account_exists +from lnbits.settings import settings from ..wasm.config import WasmAPIRouteConfig from ..wasm.invoke import invoke_wasm_extension_export from ..wasm.loader import WasmExtension +class WasmRequestBodyTooLargeError(ValueError): + pass + + +@dataclass(frozen=True) +class WasmRoutePayload: + data: dict[str, Any] + request_bytes: int | None + + def register_wasm_extension_api_routes(app: FastAPI, extension: WasmExtension) -> None: for route_config in extension.config.api_routes: _add_wasm_extension_api_route(app, extension, route_config) @@ -39,20 +52,27 @@ def _add_wasm_extension_api_route( access_token: str | None = None, ) -> dict[str, Any]: try: - payload = await _read_api_payload(request, path_params) + limits = await get_wasm_runtime_limits_for_extension(extension.id) + payload = await _read_api_payload( + request, + path_params, + max_body_bytes=limits["wasm_runtime_max_request_bytes"], + ) return await invoke_wasm_extension_export( extension.id, export_name, - payload, + payload.data, user=account, access_token=access_token, trigger_type="http", method=request.method, path=request.url.path, request_id=request.headers.get("x-request-id"), - request_bytes=_request_body_bytes(request), + request_bytes=payload.request_bytes, context_data={"origin": _request_origin(request)}, ) + except WasmRequestBodyTooLargeError as exc: + raise HTTPException(status_code=413, detail=str(exc)) from exc except KeyError as exc: raise HTTPException(status_code=404, detail=str(exc)) from exc except PermissionError as exc: @@ -86,22 +106,70 @@ def _add_wasm_extension_api_route( async def _read_api_payload( request: Request, path_params: dict[str, str], -) -> dict[str, Any]: + *, + max_body_bytes: int, +) -> WasmRoutePayload: payload = _read_api_path_params(request, path_params) payload.update(_read_api_query_params(request)) + request_bytes: int | None = None if request.method in {"POST", "PUT", "PATCH"}: - payload.update(await _read_json_object(request)) - return payload + body, request_bytes = await _read_json_object_with_size( + request, + max_body_bytes=max_body_bytes, + ) + payload.update(body) + return WasmRoutePayload(payload, request_bytes) -async def _read_json_object(request: Request) -> dict[str, Any]: - body = await request.body() +async def _read_json_object( + request: Request, + *, + max_body_bytes: int | None = None, +) -> dict[str, Any]: + body, _ = await _read_json_object_with_size( + request, + max_body_bytes=( + settings.wasm_runtime_max_request_bytes + if max_body_bytes is None + else max_body_bytes + ), + ) + return body + + +async def _read_json_object_with_size( + request: Request, + *, + max_body_bytes: int, +) -> tuple[dict[str, Any], int]: + body = await _read_limited_body(request, max_body_bytes=max_body_bytes) if not body: - return {} + return {}, 0 value = json.loads(body) if not isinstance(value, dict): raise TypeError("WASM extension API payload must be a JSON object.") - return value + return value, len(body) + + +async def _read_limited_body(request: Request, *, max_body_bytes: int) -> bytes: + content_length = _request_content_length(request) + if _wasm_request_too_large(content_length, max_body_bytes): + raise WasmRequestBodyTooLargeError( + f"WASM extension request is too large: {content_length} bytes." + ) + + chunks: list[bytes] = [] + size = 0 + async for chunk in request.stream(): + if not chunk: + continue + size += len(chunk) + if _wasm_request_too_large(size, max_body_bytes): + raise WasmRequestBodyTooLargeError( + f"WASM extension request is too large: {size} bytes." + ) + chunks.append(chunk) + return b"".join(chunks) def _read_api_path_params( @@ -119,15 +187,17 @@ def _read_api_query_params(request: Request) -> dict[str, Any]: return {_snake_to_camel(key): value for key, value in request.query_params.items()} -def _request_body_bytes(request: Request) -> int | None: - if request.method not in {"POST", "PUT", "PATCH"}: - return None +def _request_content_length(request: Request) -> int | None: content_length = request.headers.get("content-length") if content_length and content_length.isdigit(): return int(content_length) return None +def _wasm_request_too_large(size: int | None, max_body_bytes: int) -> bool: + return size is not None and max_body_bytes > 0 and size > max_body_bytes + + def _request_origin(request: Request) -> str | None: origin = request.headers.get("origin") if not origin: diff --git a/lnbits/core/wasm_ext/routes/ui.py b/lnbits/core/wasm_ext/routes/ui.py index 63029a96e..d0d4362f1 100644 --- a/lnbits/core/wasm_ext/routes/ui.py +++ b/lnbits/core/wasm_ext/routes/ui.py @@ -17,6 +17,7 @@ from lnbits.decorators import ( from ..wasm.loader import WasmExtension from .api import ( + WasmRequestBodyTooLargeError, _has_route, _path_template_pattern, _read_json_object, @@ -66,6 +67,8 @@ def _add_wasm_extension_frame_config_route( ) -> dict[str, Any]: try: body = await _read_json_object(request) + except WasmRequestBodyTooLargeError as exc: + raise HTTPException(status_code=413, detail=str(exc)) from exc except (TypeError, ValueError) as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc diff --git a/tests/unit/test_wasm_extension_routes.py b/tests/unit/test_wasm_extension_routes.py new file mode 100644 index 000000000..cc34b5644 --- /dev/null +++ b/tests/unit/test_wasm_extension_routes.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import cast + +import pytest +from fastapi import Request + +from lnbits.core.wasm_ext.routes.api import ( + WasmRequestBodyTooLargeError, + _read_api_payload, + _read_json_object_with_size, +) + + +@pytest.mark.anyio +async def test_wasm_json_reader_rejects_large_content_length_without_reading(): + request = _FakeRequest([b"{}"], content_length="11") + + with pytest.raises(WasmRequestBodyTooLargeError, match="11 bytes"): + await _read_json_object_with_size(cast(Request, request), max_body_bytes=10) + + assert request.stream_started is False + + +@pytest.mark.anyio +async def test_wasm_json_reader_rejects_large_stream_without_content_length(): + request = _FakeRequest([b'{"value":"', b"x" * 20, b'"}']) + + with pytest.raises(WasmRequestBodyTooLargeError): + await _read_json_object_with_size(cast(Request, request), max_body_bytes=16) + + assert request.stream_started is True + + +@pytest.mark.anyio +async def test_wasm_api_payload_records_actual_body_bytes(): + body = b'{"amount":21}' + request = _FakeRequest( + [body], + path_params={"invoice_id": "abc"}, + query_params={"include_paid": "true"}, + ) + + payload = await _read_api_payload( + cast(Request, request), + {"invoice_id": "invoiceId"}, + max_body_bytes=100, + ) + + assert payload.data == { + "invoiceId": "abc", + "includePaid": "true", + "amount": 21, + } + assert payload.request_bytes == len(body) + + +class _FakeRequest: + method = "POST" + + def __init__( + self, + chunks: list[bytes], + *, + content_length: str | None = None, + path_params: dict[str, str] | None = None, + query_params: dict[str, str] | None = None, + ) -> None: + self._chunks = chunks + self.headers: dict[str, str] = {} + if content_length is not None: + self.headers["content-length"] = content_length + self.path_params = path_params or {} + self.query_params = query_params or {} + self.stream_started = False + + async def stream(self) -> AsyncIterator[bytes]: + self.stream_started = True + for chunk in self._chunks: + yield chunk