Fix chunked request bodies on both model APIs

README.md

@@ -109,6 +109,9 @@ The response contains generated text in `choices`. For other clients, use
 base URL `http://127.0.0.1:11500/v1`, model `flint`, and the API key from
 `rec auth token`.
 
+Requests may use `Content-Length` or `Transfer-Encoding: chunked` with a body
+of at most 8 MiB.
+
 The Anthropic-compatible API is available at
 `http://127.0.0.1:11500/v1/messages`. Point `ANTHROPIC_BASE_URL` at
 `http://127.0.0.1:11500` and send the token from `rec auth token` as a bearer

engine/src/rec_mlx/http_body.py

@@ -0,0 +1,55 @@
+"""Bounded HTTP/1 request-body framing, before handing JSON to mlx-lm."""
+
+MAX_BODY = 8 * 1024 * 1024
+MAX_LINE = 8192
+MAX_OVERHEAD = 64 * 1024
+
+
+def read_body(stream, headers):
+    lengths = headers.get_all("Content-Length", [])
+    encodings = headers.get_all("Transfer-Encoding", [])
+    if encodings:
+        if lengths or len(encodings) != 1 or encodings[0].strip().lower() != "chunked":
+            raise ValueError("Use either Content-Length or Transfer-Encoding: chunked")
+        body = bytearray()
+        overhead = 0
+
+        def line():
+            nonlocal overhead
+            value = stream.readline(MAX_LINE + 1)
+            overhead += len(value)
+            if len(value) > MAX_LINE or overhead > MAX_OVERHEAD or not value.endswith(b"\r\n"):
+                raise ValueError("Invalid or excessive chunk framing")
+            return value[:-2]
+
+        while True:
+            size_text = line().split(b";", 1)[0]
+            if not size_text or any(c not in b"0123456789abcdefABCDEF" for c in size_text):
+                raise ValueError("Invalid chunk size")
+            size = int(size_text, 16)
+            if len(body) + size > MAX_BODY:
+                raise ValueError("Request body must be between 1 byte and 8 MiB")
+            if size == 0:
+                # Trailers are deliberately discarded, never promoted into headers.
+                while trailer := line():
+                    if b":" not in trailer or trailer[:1] in (b" ", b"\t"):
+                        raise ValueError("Invalid chunk trailer")
+                break
+            chunk = stream.read(size)
+            if len(chunk) != size or stream.read(2) != b"\r\n":
+                raise ValueError("Truncated or invalid chunk")
+            overhead += 2
+            body.extend(chunk)
+        if not body:
+            raise ValueError("Request body must be between 1 byte and 8 MiB")
+        return bytes(body)
+    length_text = lengths[0].strip() if len(lengths) == 1 else ""
+    if not length_text.isascii() or not length_text.isdecimal():
+        raise ValueError("A single Content-Length or Transfer-Encoding: chunked is required")
+    length = int(length_text)
+    if not 0 < length <= MAX_BODY:
+        raise ValueError("Request body must be between 1 byte and 8 MiB")
+    body = stream.read(length)
+    if len(body) != length:
+        raise ValueError("Truncated request body")
+    return body

engine/src/rec_mlx/server.py

@@ -16,6 +16,7 @@
     messages_to_chat,
 )
 from .engine import Sampling
+from .http_body import read_body
 
 
 class RequestError(ValueError):
@@ -99,19 +100,31 @@ def do_POST(self):
             if self.route() not in ("/v1/chat/completions", "/v1/messages"):
                 self._reject(404, "Unknown endpoint")
                 return
-            try:
-                if not 0 < int(self.headers["Content-Length"]) <= 8 * 1024 * 1024:
-                    raise ValueError("Request body must be between 1 byte and 8 MiB")
-            except (ValueError, TypeError) as error:
-                self._reject(400, str(error))
-                return
             if generator.stopping.is_set():
                 self._reject(503, "The server is stopping")
                 return
             if not admission.acquire(blocking=False):
                 self._reject(429, "The request queue is full; retry after a generation finishes")
                 return
+            source = self.rfile
             try:
+                timeout = self.connection.gettimeout()
+                try:
+                    self.connection.settimeout(10)
+                    body = read_body(self.rfile, self.headers)
+                except (ValueError, TimeoutError) as error:
+                    self.close_connection = True
+                    self._reject(400, str(error))
+                    return
+                finally:
+                    self.connection.settimeout(timeout)
+                self.rfile = io.BytesIO(body)
+                if "Transfer-Encoding" in self.headers:
+                    del self.headers["Transfer-Encoding"]
+                if "Content-Length" in self.headers:
+                    self.headers.replace_header("Content-Length", str(len(body)))
+                else:
+                    self.headers["Content-Length"] = str(len(body))
                 if self.route() == "/v1/messages":
                     self.anthropic = True
                     self._prepare_messages()
@@ -121,6 +134,7 @@ def do_POST(self):
             except (BrokenPipeError, ConnectionResetError):
                 logging.info("Client disconnected; cancelling its generation")
             finally:
+                self.rfile = source
                 admission.release()
 
         def _prepare_messages(self):

engine/tests/test_http_body.py

@@ -0,0 +1,65 @@
+import io
+import unittest
+from email.message import Message
+
+from rec_mlx.http_body import MAX_BODY, MAX_LINE, read_body
+
+
+class BodyTest(unittest.TestCase):
+    def read(self, body, *headers):
+        message = Message()
+        for name, value in headers:
+            message[name] = value
+        return read_body(io.BytesIO(body), message)
+
+    def test_fixed_length(self):
+        self.assertEqual(self.read(b"abc", ("Content-Length", "3")), b"abc")
+        self.assertEqual(self.read(b"abc", ("Content-Length", " 3\t")), b"abc")
+        self.assertEqual(
+            self.read(b"x" * MAX_BODY, ("Content-Length", str(MAX_BODY))), b"x" * MAX_BODY
+        )
+
+    def test_chunk_extensions_and_trailers(self):
+        self.assertEqual(
+            self.read(
+                b"1;foo=bar\r\na\r\n2\r\nbc\r\n0\r\nContent-Length: 99\r\n\r\n",
+                ("Transfer-Encoding", "CHUNKED"),
+            ),
+            b"abc",
+        )
+
+    def test_rejects_bad_lengths_and_ambiguous_framing(self):
+        cases = [
+            (),
+            (("Content-Length", "-1"),),
+            (("Content-Length", "+1"),),
+            (("Content-Length", "0"),),
+            (("Content-Length", "2"),),  # truncated
+            (("Content-Length", str(MAX_BODY + 1)),),
+            (("Content-Length", "1"), ("Content-Length", "1")),
+            (("Content-Length", "1"), ("Transfer-Encoding", "chunked")),
+            (("Transfer-Encoding", "gzip, chunked"),),
+            (("Transfer-Encoding", "chunked"), ("Transfer-Encoding", "chunked")),
+        ]
+        for headers in cases:
+            with self.subTest(headers=headers), self.assertRaises(ValueError):
+                self.read(b"x", *headers)
+
+    def test_rejects_malformed_or_unbounded_chunks(self):
+        for body in (
+            b"0\r\n\r\n",
+            b"-1\r\n",
+            b"+1\r\n",
+            b" 1\r\n",
+            b"g\r\n",
+            b"1\na\n",
+            b"1\r\n",
+            b"1\r\naXX",
+            b"1\r\na\r\n0\r\n",
+            b"0\r\nbad\r\n\r\n",
+            b"800001\r\n",
+            b"1;" + b"a" * MAX_LINE + b"\r\n",
+            b"1\r\na\r\n" * 14000 + b"0\r\n\r\n",
+        ):
+            with self.subTest(body=body[:40]), self.assertRaises(ValueError):
+                self.read(body, ("Transfer-Encoding", "chunked"))

engine/tests/test_serving.py

@@ -92,6 +92,62 @@ def generate_tokens(self, ids, sampling, *, should_stop, return_logprobs=False,
 
 
 class ServingTest(unittest.TestCase):
+    def test_body_timeout_rejects_without_generation_and_releases_admission(self):
+        with patch("rec_mlx.server.read_body", side_effect=TimeoutError("body read timed out")):
+            self.assertEqual(self.request("B", max_tokens=1)[0], 400)
+        self.assertEqual(self.engine.calls, [])
+        self.assertEqual(self.request("B", max_tokens=1)[0], 200)
+
+    def test_invalid_http_framing_is_rejected_before_generation(self):
+        for headers in (
+            b"",
+            b"Content-Length: 1\r\nContent-Length: 1\r\n",
+            b"Content-Length: 1\r\nTransfer-Encoding: chunked\r\n",
+            b"Transfer-Encoding: gzip\r\n",
+        ):
+            with self.subTest(headers=headers):
+                client = socket.create_connection(self.server.server_address, timeout=5)
+                try:
+                    client.sendall(
+                        b"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
+                        + headers
+                        + b"\r\nx"
+                    )
+                    response = http.client.HTTPResponse(client)
+                    response.begin()
+                    self.assertEqual(response.status, 400, response.read())
+                finally:
+                    client.close()
+        self.assertEqual(self.engine.calls, [])
+
+    def test_chunked_requests_on_both_apis(self):
+        for path in ("/v1/chat/completions", "/v1/messages"):
+            for stream in (False, True):
+                with self.subTest(path=path, stream=stream):
+                    body = json.dumps(
+                        {
+                            "model": "flint",
+                            "max_tokens": 1,
+                            "stream": stream,
+                            "messages": [{"role": "user", "content": "B"}],
+                        }
+                    ).encode()
+                    conn = http.client.HTTPConnection(*self.server.server_address, timeout=5)
+                    try:
+                        conn.request(
+                            "POST",
+                            path,
+                            [body[:10], body[10:]],
+                            {"Content-Type": "application/json"},
+                            encode_chunked=True,
+                        )
+                        response = conn.getresponse()
+                        payload = response.read()
+                        self.assertEqual(response.status, 200, payload)
+                        self.assertTrue(payload)
+                    finally:
+                        conn.close()
+
     def test_idle_preparation_follows_response_and_yields_to_queued_request(self):
         preparing = threading.Event()
         proceed = threading.Event()