Harden order-model HTTP transport against unsafe deserialization

.github/workflows/linting.yml

@@ -38,6 +38,8 @@ jobs:
 
       - name: Run Ruff
         uses: chartboost/ruff-action@v1
+        with:
+          version: "0.11.2"
 
       - name: Run ShellCheck
         uses: ludeeus/action-shellcheck@master

README.md

@@ -133,6 +133,11 @@ To start the order model Ray server:
 bash scripts/start-order-model.sh
 ```
 
+The order-model HTTP endpoint listens on localhost by default. Expose it to
+remote clients only behind network access controls and authentication. The
+client and server must use the same version: requests and responses use a
+fixed-size little-endian int32 binary format.
+
 > **Prerequisites:**
 > - The Ray server must be running and accessible at the configured IP and port
 > - Sufficient computational resources are required to run the model

market_simulation/rollout/model_client.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import logging
-import pickle
 from typing import TYPE_CHECKING
 
 import ray
 import requests
 
+from market_simulation.conf import C
+from market_simulation.rollout.wire_format import decode_int32, encode_int32
+
 if TYPE_CHECKING:
     import numpy as np
     import numpy.typing as npt
@@ -15,10 +17,11 @@
 @ray.remote
 def send_query(url: str, state: npt.NDArray[np.int32]) -> npt.NDArray[np.int32]:
     """Send query to serving model."""
-    pdata = pickle.dumps(state)
-    response = requests.post(url, data=pdata)
-    output = pickle.loads(response.content)
-    return output
+    expected_elements = C.order_model.seq_len * C.order_model.token_dim
+    pdata = encode_int32(state, expected_elements)
+    response = requests.post(url, data=pdata, headers={"Content-Type": "application/octet-stream"})
+    response.raise_for_status()
+    return decode_int32(response.content, 1)
 
 
 class ModelClient:

market_simulation/rollout/order_model_serving.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 import logging
-import pickle
 
 import numpy as np
 import numpy.typing as npt
 import torch
 from ray import serve
+from starlette.responses import Response
 
 from market_simulation.conf import C
 from market_simulation.models.order_model import OrderModel
+from market_simulation.rollout.wire_format import PayloadTooLargeError, decode_int32, encode_int32, read_fixed_body
 
 
 @serve.deployment(
@@ -35,7 +36,7 @@ def _load_model(self) -> OrderModel:
         return order_model
 
     @serve.batch(max_batch_size=C.model_serving.max_batch_size)  # type: ignore
-    async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[npt.NDArray[np.int32]]:
+    async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[bytes]:
         """Batch inference."""
         batch_size = len(requests)
         input_tensor = torch.from_numpy(np.asarray(requests)).cuda()
@@ -45,18 +46,27 @@ async def batch_inference(self, requests: list[npt.NDArray[np.int32]]) -> list[n
             output_tensor: np.ndarray = self.model.sample(input_tensor, self.temperature).int().cpu().reshape((batch_size, -1)).numpy()
         logging.info(f"output shape: {output_tensor.shape}")
 
-        results: list[npt.NDArray[np.int32]] = []
+        results: list[bytes] = []
         for i in range(batch_size):
             output = output_tensor[i]
-            arr = np.array([output], dtype=np.int32)
-            results.append(pickle.dumps(arr))  # type: ignore
+            results.append(encode_int32(output, 1))
         return results
 
-    async def __call__(self, request) -> list[npt.NDArray[np.int32]]:  # noqa: ANN001
+    async def __call__(self, request) -> Response:  # noqa: ANN001
         """Handle request."""
-        request_bytes = await request.body()
-        arr = pickle.loads(request_bytes)
-        return await self.batch_inference(arr)  # type: ignore
+        expected_elements = C.order_model.seq_len * C.order_model.token_dim
+        expected_bytes = expected_elements * np.dtype(np.int32).itemsize
+        if request.headers.get("content-type") != "application/octet-stream":
+            return Response("Expected application/octet-stream", status_code=415)
+        try:
+            request_bytes = await read_fixed_body(request.stream(), expected_bytes)
+            arr = decode_int32(request_bytes, expected_elements)
+        except PayloadTooLargeError as error:
+            return Response(str(error), status_code=413)
+        except ValueError as error:
+            return Response(str(error), status_code=400)
+        result = await self.batch_inference(arr)  # type: ignore
+        return Response(result, media_type="application/octet-stream")
 
 
 order_model_app = OrderModelServing.bind()  # type: ignore

market_simulation/rollout/wire_format.py

@@ -0,0 +1,39 @@
+"""Fixed-size int32 wire format for order-model requests and responses."""
+
+from collections.abc import AsyncIterable
+
+import numpy as np
+import numpy.typing as npt
+
+WIRE_DTYPE = np.dtype("<i4")
+
+
+class PayloadTooLargeError(ValueError):
+    """Raised when a streamed request exceeds the fixed wire size."""
+
+
+async def read_fixed_body(chunks: AsyncIterable[bytes], expected_size: int) -> bytes:
+    """Read a request body without buffering more than the expected size."""
+    body = bytearray()
+    async for chunk in chunks:
+        if len(body) + len(chunk) > expected_size:
+            raise PayloadTooLargeError("Order-model request is too large")
+        body.extend(chunk)
+    if len(body) != expected_size:
+        raise ValueError("Invalid int32 payload length")
+    return bytes(body)
+
+
+def encode_int32(values: npt.NDArray[np.int32], expected_elements: int) -> bytes:
+    """Encode exactly the expected number of int32 values in little-endian order."""
+    array = np.asarray(values)
+    if array.dtype != np.dtype(np.int32) or array.size != expected_elements:
+        raise ValueError("Expected a fixed-size int32 array")
+    return array.astype(WIRE_DTYPE, copy=False).tobytes(order="C")
+
+
+def decode_int32(payload: bytes, expected_elements: int) -> npt.NDArray[np.int32]:
+    """Decode a fixed-size int32 payload without object deserialization."""
+    if len(payload) != expected_elements * WIRE_DTYPE.itemsize:
+        raise ValueError("Invalid int32 payload length")
+    return np.frombuffer(payload, dtype=WIRE_DTYPE).astype(np.int32, copy=True)

ray_serving.yaml

@@ -1,7 +1,7 @@
 proxy_location: EveryNode
 
 http_options:
-  host: 0.0.0.0
+  host: 127.0.0.1
 
   port: 8000
 

tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for MarS."""

tests/test_wire_format.py

@@ -0,0 +1,71 @@
+"""Checks for the order-model network format and bounded request reader."""
+
+import asyncio
+import pickle
+import unittest
+from collections.abc import AsyncIterator, Callable
+
+import numpy as np
+
+from market_simulation.rollout.wire_format import PayloadTooLargeError, decode_int32, encode_int32, read_fixed_body
+
+
+async def chunks(*parts: bytes) -> AsyncIterator[bytes]:
+    """Yield a request body in arbitrary network chunks."""
+    for part in parts:
+        yield part
+
+
+def fail_if_unpickled() -> None:
+    """Make an unsafe deserialization observable to the test."""
+    raise AssertionError("Untrusted pickle was executed")
+
+
+class PickleTrap:
+    """A payload that must remain inert on the wire."""
+
+    def __reduce__(self) -> tuple[Callable[[], None], tuple[()]]:
+        """Return a callable that fails if deserialized."""
+        return fail_if_unpickled, ()
+
+
+class WireFormatTests(unittest.TestCase):
+    """Exercise fixed-size int32 transport at the network boundary."""
+
+    def test_round_trip(self) -> None:
+        """Encode and decode fixed-width little-endian int32 payloads."""
+        values = np.array([1, -2, 123456], dtype=np.int32)
+        payload = encode_int32(values, 3)
+        assert payload == b"\x01\x00\x00\x00\xfe\xff\xff\xff\x40\xe2\x01\x00"
+        np.testing.assert_array_equal(decode_int32(payload, 3), values)
+
+    def test_rejects_wrong_type_and_size(self) -> None:
+        """Reject payloads with an unexpected dtype, element count, or byte size."""
+        expect_raises(ValueError, encode_int32, np.array([1], dtype=object), 1)
+        expect_raises(ValueError, encode_int32, np.array([1], dtype=np.int32), 2)
+        expect_raises(ValueError, decode_int32, b"\x00" * 3, 1)
+
+    def test_stream_is_bounded_and_exact(self) -> None:
+        """Read exactly the expected stream length and reject over/underflow."""
+        assert asyncio.run(read_fixed_body(chunks(b"ab", b"cd"), 4)) == b"abcd"
+        expect_raises(PayloadTooLargeError, asyncio.run, read_fixed_body(chunks(b"abcd", b"e" * 1000000), 4))
+        expect_raises(ValueError, asyncio.run, read_fixed_body(chunks(b"abc"), 4))
+
+    def test_pickle_payload_is_never_deserialized(self) -> None:
+        """Treat pickle bytes as inert int32 wire data."""
+        payload = pickle.dumps(PickleTrap()).ljust(64, b"\x00")
+        decoded = decode_int32(payload, 16)
+        assert decoded.size == 16
+
+
+def expect_raises(exception_type: type[Exception], function: Callable[..., object], *args: object) -> None:
+    """Assert that a callable raises the expected exception type."""
+    try:
+        function(*args)
+    except exception_type:
+        return
+    raise AssertionError(f"Expected {exception_type.__name__}")
+
+
+if __name__ == "__main__":
+    unittest.main()