Merge pull request #468 from apolcyn/add_send_with_padding_option

Allow sending data frames with padding

Author: Lukasa

HISTORY.rst

@@ -22,6 +22,9 @@ API Changes (Backward-Compatible)
   a subclass of ``int``, this is non-breaking.
 - Deprecated the other fields in ``h2.settings``. These will be removed in
   3.0.0.
+- Added an optional ``pad_length`` parameter to ``H2Connection.send_data``
+  to allow the user to include padding on a data frame.
+
 
 Bugfixes
 ~~~~~~~~

h2/connection.py

@@ -833,7 +833,7 @@ def send_headers(self, stream_id, headers, end_stream=False,
 
         self._prepare_for_sending(frames)
 
-    def send_data(self, stream_id, data, end_stream=False):
+    def send_data(self, stream_id, data, end_stream=False, pad_length=None):
         """
         Send data on a given stream.
 
@@ -857,24 +857,44 @@ def send_data(self, stream_id, data, end_stream=False):
         :param end_stream: (optional) Whether this is the last data to be sent
             on the stream. Defaults to ``False``.
         :type end_stream: ``bool``
+        :param pad_length: (optional) Length of the padding to apply to the
+            data frame. Defaults to ``None`` for no use of padding. Note that
+            a value of ``0`` results in padding of length ``0``
+            (with the "padding" flag set on the frame).
+
+            .. versionadded:: 2.6.0
+
+        :type pad_length: ``int``
         :returns: Nothing
         """
-        if len(data) > self.local_flow_control_window(stream_id):
+        frame_size = len(data)
+        if pad_length is not None:
+            if not isinstance(pad_length, int):
+                raise TypeError("pad_length must be an int")
+            if pad_length < 0 or pad_length > 255:
+                raise ValueError("pad_length must be within range: [0, 255]")
+            # Account for padding bytes plus the 1-byte padding length field.
+            frame_size += pad_length + 1
+
+        if frame_size > self.local_flow_control_window(stream_id):
             raise FlowControlError(
                 "Cannot send %d bytes, flow control window is %d." %
-                (len(data), self.local_flow_control_window(stream_id))
+                (frame_size, self.local_flow_control_window(stream_id))
             )
-        elif len(data) > self.max_outbound_frame_size:
+        elif frame_size > self.max_outbound_frame_size:
             raise FrameTooLargeError(
                 "Cannot send frame size %d, max frame size is %d" %
-                (len(data), self.max_outbound_frame_size)
+                (frame_size, self.max_outbound_frame_size)
             )
 
         self.state_machine.process_input(ConnectionInputs.SEND_DATA)
-        frames = self.streams[stream_id].send_data(data, end_stream)
+        frames = self.streams[stream_id].send_data(
+            data, end_stream, pad_length=pad_length
+        )
+
         self._prepare_for_sending(frames)
 
-        self.outbound_flow_control_window -= len(data)
+        self.outbound_flow_control_window -= frame_size
         assert self.outbound_flow_control_window >= 0
 
     def end_stream(self, stream_id):

h2/stream.py

@@ -841,7 +841,7 @@ def locally_pushed(self):
         assert not events
         return []
 
-    def send_data(self, data, end_stream=False):
+    def send_data(self, data, end_stream=False, pad_length=None):
         """
         Prepare some data frames. Optionally end the stream.
 
@@ -854,8 +854,12 @@ def send_data(self, data, end_stream=False):
         if end_stream:
             self.state_machine.process_input(StreamInputs.SEND_END_STREAM)
             df.flags.add('END_STREAM')
+        if pad_length is not None:
+            df.flags.add('PADDED')
+            df.pad_length = pad_length
 
-        self.outbound_flow_control_window -= len(data)
+        # Subtract flow_controlled_length to account for possible padding
+        self.outbound_flow_control_window -= df.flow_controlled_length
         assert self.outbound_flow_control_window >= 0
 
         return [df]

test/test_basic_logic.py

@@ -109,10 +109,117 @@ def test_sending_data(self):
         c.clear_outbound_data_buffer()
         events = c.send_data(1, b'some data')
         assert not events
+        data_to_send = c.data_to_send()
         assert (
-            c.data_to_send() == b'\x00\x00\t\x00\x00\x00\x00\x00\x01some data'
+            data_to_send == b'\x00\x00\t\x00\x00\x00\x00\x00\x01some data'
         )
 
+        buffer = h2.frame_buffer.FrameBuffer(server=False)
+        buffer.max_frame_size = 65535
+        buffer.add_data(data_to_send)
+        data_frame = list(buffer)[0]
+        sanity_check_data_frame(
+            data_frame=data_frame,
+            expected_flow_controlled_length=len(b'some data'),
+            expect_padded_flag=False,
+            expected_data_frame_pad_length=0
+        )
+
+    def test_sending_data_with_padding(self):
+        """
+        Single data frames with padding are encoded correctly.
+        """
+        c = h2.connection.H2Connection()
+        c.initiate_connection()
+        c.send_headers(1, self.example_request_headers)
+
+        # Clear the data, then send some data.
+        c.clear_outbound_data_buffer()
+        events = c.send_data(1, b'some data', pad_length=5)
+        assert not events
+        data_to_send = c.data_to_send()
+        assert data_to_send == (
+            b'\x00\x00\x0f\x00\x08\x00\x00\x00\x01'
+            b'\x05some data\x00\x00\x00\x00\x00'
+        )
+
+        buffer = h2.frame_buffer.FrameBuffer(server=False)
+        buffer.max_frame_size = 65535
+        buffer.add_data(data_to_send)
+        data_frame = list(buffer)[0]
+        sanity_check_data_frame(
+            data_frame=data_frame,
+            expected_flow_controlled_length=len(b'some data') + 1 + 5,
+            expect_padded_flag=True,
+            expected_data_frame_pad_length=5
+        )
+
+    def test_sending_data_with_zero_length_padding(self):
+        """
+        Single data frames with zero-length padding are encoded
+        correctly.
+        """
+        c = h2.connection.H2Connection()
+        c.initiate_connection()
+        c.send_headers(1, self.example_request_headers)
+
+        # Clear the data, then send some data.
+        c.clear_outbound_data_buffer()
+        events = c.send_data(1, b'some data', pad_length=0)
+        assert not events
+        data_to_send = c.data_to_send()
+        assert data_to_send == (
+            b'\x00\x00\x0a\x00\x08\x00\x00\x00\x01'
+            b'\x00some data'
+        )
+
+        buffer = h2.frame_buffer.FrameBuffer(server=False)
+        buffer.max_frame_size = 65535
+        buffer.add_data(data_to_send)
+        data_frame = list(buffer)[0]
+        sanity_check_data_frame(
+            data_frame=data_frame,
+            expected_flow_controlled_length=len(b'some data') + 1,
+            expect_padded_flag=True,
+            expected_data_frame_pad_length=0
+        )
+
+    @pytest.mark.parametrize("expected_error,pad_length", [
+        (None,  0),
+        (None, 255),
+        (None, None),
+        (ValueError, -1),
+        (ValueError, 256),
+        (TypeError, 'invalid'),
+        (TypeError, ''),
+        (TypeError, '10'),
+        (TypeError, {}),
+        (TypeError, ['1', '2', '3']),
+        (TypeError, []),
+        (TypeError, 1.5),
+        (TypeError, 1.0),
+        (TypeError, -1.0),
+    ])
+    def test_sending_data_with_invalid_padding_length(self,
+                                                      expected_error,
+                                                      pad_length):
+        """
+        ``send_data`` with a ``pad_length`` parameter that is an integer
+        outside the range of [0, 255] throws a ``ValueError``, and a
+        ``pad_length`` parameter which is not an ``integer`` type
+        throws a ``TypeError``.
+        """
+        c = h2.connection.H2Connection()
+        c.initiate_connection()
+        c.send_headers(1, self.example_request_headers)
+
+        c.clear_outbound_data_buffer()
+        if expected_error is not None:
+            with pytest.raises(expected_error):
+                c.send_data(1, b'some data', pad_length=pad_length)
+        else:
+            c.send_data(1, b'some data', pad_length=pad_length)
+
     def test_closing_stream_sending_data(self, frame_factory):
         """
         We can close a stream with a data frame.
@@ -1576,3 +1683,27 @@ def test_receiving_goaway_frame_with_unknown_error(self, frame_factory):
         assert c.state_machine.state == h2.connection.ConnectionState.CLOSED
 
         assert not c.data_to_send()
+
+
+def sanity_check_data_frame(data_frame,
+                            expected_flow_controlled_length,
+                            expect_padded_flag,
+                            expected_data_frame_pad_length):
+        """
+        ``data_frame`` is a frame of type ``hyperframe.frame.DataFrame``,
+        and the ``flags`` and ``flow_controlled_length`` of ``data_frame``
+        match expectations.
+        """
+
+        assert isinstance(data_frame, hyperframe.frame.DataFrame)
+
+        assert(
+            (data_frame.flow_controlled_length ==
+             expected_flow_controlled_length)
+        )
+        if expect_padded_flag:
+            assert 'PADDED' in data_frame.flags
+        else:
+            assert 'PADDED' not in data_frame.flags
+
+        assert data_frame.pad_length == expected_data_frame_pad_length

test/test_flow_control_window.py

@@ -52,6 +52,23 @@ def test_flow_control_decreases_with_sent_data(self):
         remaining_length = self.DEFAULT_FLOW_WINDOW - len(b'some data')
         assert (c.local_flow_control_window(1) == remaining_length)
 
+    @pytest.mark.parametrize("pad_length", [5, 0])
+    def test_flow_control_decreases_with_sent_data_with_padding(self,
+                                                                pad_length):
+        """
+        When padded data is sent on a stream, the flow control window drops
+        by the length of the padding plus 1 for the 1-byte padding length
+        field.
+        """
+        c = h2.connection.H2Connection()
+        c.send_headers(1, self.example_request_headers)
+
+        c.send_data(1, b'some data', pad_length=pad_length)
+        remaining_length = (
+            self.DEFAULT_FLOW_WINDOW - len(b'some data') - pad_length - 1
+        )
+        assert c.local_flow_control_window(1) == remaining_length
+
     def test_flow_control_decreases_with_received_data(self, frame_factory):
         """
         When data is received on a stream, the remote flow control window
@@ -70,7 +87,7 @@ def test_flow_control_decreases_with_received_data(self, frame_factory):
     def test_flow_control_decreases_with_padded_data(self, frame_factory):
         """
         When padded data is received on a stream, the remote flow control
-        window should drop by an amount that includes the padding.
+        window drops by an amount that includes the padding.
         """
         c = h2.connection.H2Connection(client_side=False)
         c.receive_data(frame_factory.preamble())