Merge branch 'add-unknown-frame' of https://github.com/haikuginger/hyperframe into haikuginger-add-unknown-frame

Lukasa · Lukasa · commit f7173957f86f · 2017-03-07T08:22:59.000Z
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -63,6 +63,10 @@ you need is not present!
    :show-inheritance:
    :members:
 
+.. autoclass:: hyperframe.frame.ExtensionFrame
+   :show-inheritance:
+   :members:
+
 .. autodata:: hyperframe.frame.FRAMES
 
 Exceptions
diff --git a/hyperframe/frame.py b/hyperframe/frame.py
@@ -88,15 +88,21 @@ def __repr__(self):
         )
 
     @staticmethod
-    def parse_frame_header(header):
+    def parse_frame_header(header, strict=False):
         """
         Takes a 9-byte frame header and returns a tuple of the appropriate
         Frame object and the length that needs to be read from the socket.
 
         This populates the flags field, and determines how long the body is.
 
+        :param strict: Whether to raise an exception when encountering a frame
+            not defined by spec and implemented by hyperframe.
+
         :raises hyperframe.exceptions.UnknownFrameError: If a frame of unknown
             type is received.
+
+        .. versionchanged:: 5.0.0
+            Added :param:`strict` to accommodate :class:`ExtensionFrame`
         """
         try:
             fields = _STRUCT_HBBBL.unpack(header)
@@ -110,7 +116,15 @@ def parse_frame_header(header):
         stream_id = fields[4] & 0x7FFFFFFF
 
         if type not in FRAMES:
-            raise UnknownFrameError(type, length)
+            if strict:
+                raise UnknownFrameError(type, length)
+
+            frame = ExtensionFrame(
+                type=type,
+                flag_byte=flags,
+                stream_id=stream_id,
+            )
+            return (frame, length)
 
         frame = FRAMES[type](stream_id)
         frame.parse_flags(flags)
@@ -736,6 +750,46 @@ def parse_body(self, data):
         self.body_len = len(data)
 
 
+class ExtensionFrame(Frame):
+    """
+    ExtensionFrame is used to wrap frames which are not natively interpretable
+    by hyperframe.
+
+    Although certain byte prefixes are ordained by specification to have
+    certain contextual meanings, frames with other prefixes are not prohibited,
+    and may be used to communicate arbitrary meaning between HTTP/2 peers.
+
+    Thus, hyperframe, rather than raising an exception when such a frame is
+    encountered, wraps it in a generic frame to be properly acted upon by
+    upstream consumers which might have additional context on how to use it.
+
+    .. versionadded:: 5.0.0
+    """
+
+    stream_association = _STREAM_ASSOC_EITHER
+
+    def __init__(self, type, flag_byte, stream_id, **kwargs):
+        super(ExtensionFrame, self).__init__(stream_id, **kwargs)
+        self.type = type
+        self.flag_byte = flag_byte
+
+    def assign_flag_mapping(self, flags):
+        """
+        When initially created, an ExtensionFrame has no concept of what flags
+        might be relevant to it, since a frame's flags are defined in the
+        specification for that frame. This method allows a frame to be assigned
+        a flag mapping after the fact, allowing downstream consumers to modify
+        ExtensionFrame to fit their needs.
+        """
+        self.defined_flags = flags
+        self.flags = Flags(self.defined_flags)
+        return self.parse_flags(self.flag_byte)
+
+    def parse_body(self, data):
+        self.body = data.tobytes()
+        self.body_len = len(data)
+
+
 _FRAME_CLASSES = [
     DataFrame,
     HeadersFrame,
diff --git a/test/test_frames.py b/test/test_frames.py
@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
+from hyperframe.flags import Flag
 from hyperframe.frame import (
     Frame, Flags, DataFrame, PriorityFrame, RstStreamFrame, SettingsFrame,
     PushPromiseFrame, PingFrame, GoAwayFrame, WindowUpdateFrame, HeadersFrame,
-    ContinuationFrame, AltSvcFrame
+    ContinuationFrame, AltSvcFrame, ExtensionFrame
 )
 from hyperframe.exceptions import (
     UnknownFrameError, InvalidPaddingError, InvalidFrameError
@@ -35,10 +36,12 @@ def test_base_frame_cant_parse_body(self):
         with pytest.raises(NotImplementedError):
             f.parse_body(data)
 
-    def test_parse_frame_header_unknown_type(self):
+    def test_parse_frame_header_unknown_type_strict(self):
         with pytest.raises(UnknownFrameError) as excinfo:
-            Frame.parse_frame_header(b'\x00\x00\x59\xFF\x00\x00\x00\x00\x01')
-
+            Frame.parse_frame_header(
+                b'\x00\x00\x59\xFF\x00\x00\x00\x00\x01',
+                strict=True
+            )
         exception = excinfo.value
         assert exception.frame_type == 0xFF
         assert exception.length == 0x59
@@ -53,6 +56,41 @@ def test_parse_frame_header_ignore_first_bit_of_stream_id(self):
 
         assert f.stream_id == 0
 
+    def test_parse_frame_header_unknown_type(self):
+        f, l = Frame.parse_frame_header(
+            b'\x00\x00\x59\xFF\x00\x00\x00\x00\x01'
+        )
+        assert f.type == 0xFF
+        assert l == 0x59
+        assert isinstance(f, ExtensionFrame)
+        assert f.stream_id == 1
+
+    def test_add_flag_options_later_unknown_type(self):
+        f, l = Frame.parse_frame_header(
+            b'\x00\x00\x59\xFF\x09\x00\x00\x00\x01'
+        )
+        assert f.type == 0xFF
+        assert l == 0x59
+        new_flags = [
+            Flag('FANCY_FLAG', 0x01),
+            Flag('REAL_THING', 0x04),
+            Flag('HUUUUUUGE', 0x08),
+        ]
+        flags = f.assign_flag_mapping(new_flags)
+        assert f.defined_flags == new_flags
+        assert flags
+        assert 'FANCY_FLAG' in flags
+        assert 'REAL_THING' not in flags
+        assert 'HUUUUUUGE' in flags
+
+    def test_parse_body_unknown_type(self):
+        f = decode_frame(
+            b'\x00\x00\x0C\xFF\x00\x00\x00\x00\x01hello world!'
+        )
+        assert f.body == b'hello world!'
+        assert f.body_len == 12
+        assert f.stream_id == 1
+
     def test_repr(self, monkeypatch):
         f = Frame(stream_id=0)
         monkeypatch.setattr(Frame, "serialize_body", lambda _: b"body")
