Fix parsing stream chunks that do not contain frames (#23)

Author: vrslev

Justfile

@@ -5,7 +5,7 @@ install:
     uv -q sync
 
 test *args:
-    uv -q run pytest {{args}}
+    uv -q run pytest -- {{args}}
 
 lint:
     uv -q run ruff check .

stompman/protocol.py

@@ -31,6 +31,13 @@
 CARRIAGE_NEWLINE_CARRIAGE_NEWLINE = (CARRIAGE, NEWLINE, CARRIAGE, NEWLINE)
 
 
+def iter_bytes(bytes_: bytes) -> tuple[bytes, ...]:
+    return struct.unpack(f"{len(bytes_)!s}c", bytes_)
+
+
+VALID_COMMANDS = [list(iter_bytes(command)) for command in COMMANDS_TO_FRAMES]
+
+
 def dump_header(key: str, value: str) -> bytes:
     escaped_key = "".join(ESCAPE_CHARS.get(char, char) for char in key)
     escaped_value = "".join(ESCAPE_CHARS.get(char, char) for char in value)
@@ -79,7 +86,7 @@ def parse_headers(buffer: list[bytes]) -> tuple[str, str] | None:
     return (b"".join(key_buffer).decode(), b"".join(value_buffer).decode()) if key_parsed else None
 
 
-def parse_lines_into_frame(lines: deque[list[bytes]]) -> AnyClientFrame | AnyServerFrame | None:
+def parse_lines_into_frame(lines: deque[list[bytes]]) -> AnyClientFrame | AnyServerFrame:
     command = b"".join(lines.popleft())
     headers = {}
 
@@ -89,9 +96,7 @@ def parse_lines_into_frame(lines: deque[list[bytes]]) -> AnyClientFrame | AnySer
             headers[header[0]] = header[1]
     body = b"".join(lines.popleft()) if lines else b""
 
-    if known_frame_type := COMMANDS_TO_FRAMES.get(command):
-        return known_frame_type(headers=cast(Any, headers), body=body)
-    return None
+    return COMMANDS_TO_FRAMES[command](headers=cast(Any, headers), body=body)
 
 
 @dataclass
@@ -101,29 +106,33 @@ class Parser:
     _previous_byte: bytes = field(default=b"", init=False)
     _headers_processed: bool = field(default=False, init=False)
 
+    def _reset(self) -> None:
+        self._headers_processed = False
+        self._lines.clear()
+        self._current_line = []
+
     def load_frames(self, raw_frames: bytes) -> Iterator[AnyClientFrame | AnyServerFrame | HeartbeatFrame]:
-        buffer = deque(struct.unpack(f"{len(raw_frames)!s}c", raw_frames))
+        buffer = deque(iter_bytes(raw_frames))
         while buffer:
             byte = buffer.popleft()
 
-            if self._headers_processed and byte == NULL:
-                self._lines.append(self._current_line)
-                if parsed_frame := parse_lines_into_frame(self._lines):
-                    yield parsed_frame
-                self._headers_processed = False
-                self._lines.clear()
-                self._current_line = []
+            if byte == NULL:
+                if self._headers_processed:
+                    self._lines.append(self._current_line)
+                    yield parse_lines_into_frame(self._lines)
+                self._reset()
 
             elif not self._headers_processed and byte == NEWLINE:
                 if self._current_line or self._lines:
-                    if not self._current_line:  # extra empty line after headers
-                        self._headers_processed = True
-
                     if self._previous_byte == b"\r":
                         self._current_line.pop()
+                    self._headers_processed = not self._current_line  # extra empty line after headers
 
-                    self._lines.append(self._current_line)
-                    self._current_line = []
+                    if not self._lines and self._current_line not in VALID_COMMANDS:
+                        self._reset()
+                    else:
+                        self._lines.append(self._current_line)
+                        self._current_line = []
                 else:
                     yield HeartbeatFrame()
 

tests/test_protocol.py

@@ -187,19 +187,16 @@ def test_dump_frame(frame: AnyClientFrame, dumped_frame: bytes) -> None:
                 HeartbeatFrame(),
             ],
         ),
-        (
-            b"\n",
-            [HeartbeatFrame()],
-        ),
+        (b"\n", [HeartbeatFrame()]),
         # Two headers: only first should be accepted
         (
             b"CONNECTED\naccept-version:1.0\naccept-version:1.1\n\n\x00",
             [ConnectedFrame(headers={"accept-version": "1.0"})],
         ),
         # no end of line after command
-        (b"SOME_COMMAND", []),
-        (b"SOME_COMMAND\n", []),
-        (b"SOME_COMMAND\x00", []),
+        (b"CONNECTED", []),
+        (b"CONNECTED\n", []),
+        (b"CONNECTED\x00", []),
         # \r\n after command
         (b"CONNECTED\r\n\n\n\x00", [ConnectedFrame(headers={}, body=b"\n")]),
         (b"CONNECTED\r\nheader:1.0\n\n\x00", [ConnectedFrame(headers={"header": "1.0"})]),
@@ -213,14 +210,24 @@ def test_dump_frame(frame: AnyClientFrame, dumped_frame: bytes) -> None:
         # header value with :
         (b"CONNECTED\nheader:what:?\n\n\x00", [ConnectedFrame(headers={})]),
         # no NULL
-        (b"SOME_COMMAND\nheader:what:?\n\nhello", []),
+        (b"CONNECTED\nheader:what:?\n\nhello", []),
         # header never end
-        (b"SOME_COMMAND\nheader:hello", []),
-        (b"SOME_COMMAND\nheader:hello\n", []),
-        (b"SOME_COMMAND\nheader:hello\n\x00", []),
-        (b"SOME_COMMAND\nn", []),
+        (b"CONNECTED\nheader:hello", []),
+        (b"CONNECTED\nheader:hello\n", []),
+        (b"CONNECTED\nheader:hello\n\x00", []),
+        (b"CONNECTED\nn", []),
+        # unknown command
+        (b"SOME_COMMAND\nhead:\nheader:1.1\n\n\x00", [HeartbeatFrame()]),
         # unknown command
-        (b"SOME_COMMAND\nhead:\nheader:1.1\n\n\x00", []),
+        (
+            b"whatever\nWHATEVER\nheader:1.1\n\n\x00CONNECTED\nheader:1.1\n\n\x00\nwhatever\nCONNECTED\nheader:1.2\n\n\x00",
+            [
+                HeartbeatFrame(),
+                ConnectedFrame(headers={"header": "1.1"}, body=b""),
+                HeartbeatFrame(),
+                ConnectedFrame(headers={"header": "1.2"}, body=b""),
+            ],
+        ),
     ],
 )
 def test_load_frames(raw_frames: bytes, loaded_frames: list[AnyServerFrame]) -> None: