fix: stricter limits

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:

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
 

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