Added ability to use LF, not only CRLF delimiter

Author: cdeler

h11/_readers.py

@@ -153,7 +153,7 @@ def __call__(self, buf):
         assert self._bytes_to_discard == 0
         if self._bytes_in_chunk == 0:
             # We need to refill our chunk count
-            chunk_header = buf.maybe_extract_until_next(b"\r\n")
+            chunk_header = buf.maybe_extract_until_delimiter(b"\r?\n")
             if chunk_header is None:
                 return None
             matches = validate(

h11/_receivebuffer.py

@@ -1,3 +1,4 @@
+import re
 import sys
 
 __all__ = ["ReceiveBuffer"]
@@ -38,6 +39,12 @@
 # slightly clever thing where we delay calling compress() until we've
 # processed a whole event, which could in theory be slightly more efficient
 # than the internal bytearray support.)
+
+default_delimiter = b"\n\r?\n"
+delimiter_regex = re.compile(b"\n\r?\n", re.MULTILINE)
+line_delimiter_regex = re.compile(b"\r?\n", re.MULTILINE)
+
+
 class ReceiveBuffer(object):
     def __init__(self):
         self._data = bytearray()
@@ -46,6 +53,9 @@ def __init__(self):
         self._looked_at = 0
         self._looked_for = b""
 
+        self._delimiter = b"\n\r?\n"
+        self._delimiter_regex = delimiter_regex
+
     def __bool__(self):
         return bool(len(self))
 
@@ -79,21 +89,34 @@ def maybe_extract_at_most(self, count):
         self._start += len(out)
         return out
 
-    def maybe_extract_until_next(self, needle):
+    def maybe_extract_until_delimiter(self, delimiter=b"\n\r?\n"):
         # Returns extracted bytes on success (advancing offset), or None on
         # failure
-        if self._looked_for == needle:
-            search_start = max(self._start, self._looked_at - len(needle) + 1)
+        if delimiter == self._delimiter:
+            looked_at = max(self._start, self._looked_at - len(delimiter) + 1)
         else:
-            search_start = self._start
-        offset = self._data.find(needle, search_start)
-        if offset == -1:
+            looked_at = self._start
+            self._delimiter = delimiter
+            # re.compile operation is more expensive than just byte compare
+            if delimiter == default_delimiter:
+                self._delimiter_regex = delimiter_regex
+            else:
+                self._delimiter_regex = re.compile(delimiter, re.MULTILINE)
+
+        delimiter_match = next(
+            self._delimiter_regex.finditer(self._data, looked_at), None
+        )
+
+        if delimiter_match is None:
             self._looked_at = len(self._data)
-            self._looked_for = needle
             return None
-        new_start = offset + len(needle)
-        out = self._data[self._start : new_start]
-        self._start = new_start
+
+        _, end = delimiter_match.span(0)
+
+        out = self._data[self._start : end]
+
+        self._start = end
+
         return out
 
     # HTTP/1.1 has a number of constructs where you keep reading lines until
@@ -102,11 +125,19 @@ def maybe_extract_lines(self):
         if self._data[self._start : self._start + 2] == b"\r\n":
             self._start += 2
             return []
+        elif self._start < len(self._data) and self._data[self._start] == b"\n":
+            self._start += 1
+            return []
         else:
-            data = self.maybe_extract_until_next(b"\r\n\r\n")
+            data = self.maybe_extract_until_delimiter(b"\n\r?\n")
+
             if data is None:
                 return None
-            lines = data.split(b"\r\n")
+
+            lines = line_delimiter_regex.split(data)
+
             assert lines[-2] == lines[-1] == b""
+
             del lines[-2:]
+
             return lines

h11/tests/test_receivebuffer.py

@@ -1,3 +1,5 @@
+import pytest
+
 from .._receivebuffer import ReceiveBuffer
 
 
@@ -30,28 +32,28 @@ def test_receivebuffer():
     assert not b
 
     ################################################################
-    # maybe_extract_until_next
+    # maybe_extract_until_delimiter
     ################################################################
 
     b += b"12345a6789aa"
 
-    assert b.maybe_extract_until_next(b"a") == b"12345a"
+    assert b.maybe_extract_until_delimiter(b"a") == b"12345a"
     assert bytes(b) == b"6789aa"
 
-    assert b.maybe_extract_until_next(b"aaa") is None
+    assert b.maybe_extract_until_delimiter(b"aaa") is None
     assert bytes(b) == b"6789aa"
 
     b += b"a12"
-    assert b.maybe_extract_until_next(b"aaa") == b"6789aaa"
+    assert b.maybe_extract_until_delimiter(b"aaa") == b"6789aaa"
     assert bytes(b) == b"12"
 
     # check repeated searches for the same needle, triggering the
     # pickup-where-we-left-off logic
     b += b"345"
-    assert b.maybe_extract_until_next(b"aaa") is None
+    assert b.maybe_extract_until_delimiter(b"aaa") is None
 
     b += b"6789aaa123"
-    assert b.maybe_extract_until_next(b"aaa") == b"123456789aaa"
+    assert b.maybe_extract_until_delimiter(b"aaa") == b"123456789aaa"
     assert bytes(b) == b"123"
 
     ################################################################
@@ -76,3 +78,28 @@ def test_receivebuffer():
     b += b"\r\ntrailing"
     assert b.maybe_extract_lines() == []
     assert bytes(b) == b"trailing"
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        (
+            b"HTTP/1.1 200 OK\r\n",
+            b"Content-type: text/plain\r\n",
+            b"\r\n",
+            b"Some body",
+        ),
+        (b"HTTP/1.1 200 OK\n", b"Content-type: text/plain\n", b"\n", b"Some body"),
+        (b"HTTP/1.1 200 OK\r\n", b"Content-type: text/plain\n", b"\n", b"Some body"),
+    ],
+)
+def test_receivebuffer_for_invalid_delimiter(data):
+    b = ReceiveBuffer()
+
+    for line in data:
+        b += line
+
+    lines = b.maybe_extract_lines()
+
+    assert lines == [b"HTTP/1.1 200 OK", b"Content-type: text/plain"]
+    assert bytes(b) == b"Some body"