From 1327ed9ef665909b0cdd1446a39d0b8d9dc9f626 Mon Sep 17 00:00:00 2001
From: Nate Prewitt <Nate.Prewitt@gmail.com>
Date: Wed, 31 Aug 2016 11:51:57 -0600
Subject: [PATCH] check for pseudo header in trailer

---
 h2/utilities.py              |  7 ++++++
 test/test_complex_logic.py   |  4 ++--
 test/test_invalid_headers.py | 43 +++++++++++++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/h2/utilities.py b/h2/utilities.py
index 4991d9fb1..3b3ff26d0 100644
--- a/h2/utilities.py
+++ b/h2/utilities.py
@@ -322,6 +322,13 @@ def _reject_pseudo_header_fields(headers, hdr_validation_flags):
 
         yield header
 
+    # Pseudo-header fields MUST NOT appear in trailers - RFC 7540 § 8.1.2.1
+    if hdr_validation_flags.is_trailer and seen_pseudo_header_fields:
+        raise ProtocolError(
+            "Received pseudo-header in trailer %s" %
+            seen_pseudo_header_fields
+        )
+
 
 def _validate_host_authority_header(headers):
     """
diff --git a/test/test_complex_logic.py b/test/test_complex_logic.py
index 0d58f7d8f..d4b42b895 100644
--- a/test/test_complex_logic.py
+++ b/test/test_complex_logic.py
@@ -90,9 +90,9 @@ def test_correctly_count_server_streams(self, frame_factory):
             assert c.open_outbound_streams == expected_outbound_streams
 
             # Pushed streams can only be closed remotely.
-            f = frame_factory.build_headers_frame(
+            f = frame_factory.build_data_frame(
                 stream_id=stream_id+1,
-                headers=self.example_response_headers,
+                data=b'the content',
                 flags=['END_STREAM'],
             )
             c.receive_data(f.serialize())
diff --git a/test/test_invalid_headers.py b/test/test_invalid_headers.py
index f9f4a8cce..4d323810b 100644
--- a/test/test_invalid_headers.py
+++ b/test/test_invalid_headers.py
@@ -113,6 +113,37 @@ def test_transfer_encoding_trailers_is_valid(self, frame_factory):
         request_event = events[0]
         assert request_event.headers == headers
 
+    def test_pseudo_headers_rejected_in_trailer(self, frame_factory):
+        """
+        Ensure we reject pseudo headers included in trailers
+        """
+        trailers = [(':path', '/'), ('extra', 'value')]
+
+        c = h2.connection.H2Connection(client_side=False)
+        c.receive_data(frame_factory.preamble())
+        c.clear_outbound_data_buffer()
+
+        header_frame = frame_factory.build_headers_frame(
+            self.base_request_headers
+        )
+        trailer_frame = frame_factory.build_headers_frame(
+            trailers, flags=["END_STREAM"]
+        )
+        head = header_frame.serialize()
+        trailer = trailer_frame.serialize()
+
+        c.receive_data(head)
+        # Raise exception if pseudo header in trailer
+        with pytest.raises(h2.exceptions.ProtocolError) as e:
+            c.receive_data(trailer)
+        assert "pseudo-header in trailer" in str(e)
+
+        # Test appropriate response frame is generated
+        expected_frame = frame_factory.build_goaway_frame(
+            last_stream_id=1, error_code=h2.errors.ErrorCodes.PROTOCOL_ERROR
+        )
+        assert c.data_to_send() == expected_frame.serialize()
+
 
 class TestSendingInvalidFrameSequences(object):
     """
@@ -217,6 +248,14 @@ class TestFilter(object):
         ]
     ]
 
+    hdr_validation_no_trailers = [
+        h2.utilities.HeaderValidationFlags(is_client, is_trailer)
+        for is_client, is_trailer in [
+            (True, False),
+            (False, False)
+        ]
+    ]
+
     @pytest.mark.parametrize('validation_function', validation_functions)
     @pytest.mark.parametrize('hdr_validation_flags', hdr_validation_combos)
     @given(headers=HEADERS_STRATEGY)
@@ -241,7 +280,9 @@ def test_invalid_pseudo_headers(self, hdr_validation_flags):
             h2.utilities.validate_headers(headers, hdr_validation_flags)
 
     @pytest.mark.parametrize('validation_function', validation_functions)
-    @pytest.mark.parametrize('hdr_validation_flags', hdr_validation_combos)
+    @pytest.mark.parametrize(
+        'hdr_validation_flags', hdr_validation_no_trailers
+    )
     def test_matching_authority_host_headers(self,
                                              validation_function,
                                              hdr_validation_flags):