diff --git a/README.md b/README.md index 0e08a2a..a1dbf40 100644 --- a/README.md +++ b/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 diff --git a/engine/src/rec_mlx/http_body.py b/engine/src/rec_mlx/http_body.py new file mode 100644 index 0000000..b240e92 --- /dev/null +++ b/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 diff --git a/engine/src/rec_mlx/server.py b/engine/src/rec_mlx/server.py index 4b81ab0..f46816d 100644 --- a/engine/src/rec_mlx/server.py +++ b/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): diff --git a/engine/tests/test_http_body.py b/engine/tests/test_http_body.py new file mode 100644 index 0000000..8107a62 --- /dev/null +++ b/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")) diff --git a/engine/tests/test_serving.py b/engine/tests/test_serving.py index f3d9cbf..6b040bc 100644 --- a/engine/tests/test_serving.py +++ b/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()