Merge pull request #22626 from bocon13/end-stream-fix

Fixing bug with END_STREAM if header has continuations

src/core/ext/transport/chttp2/transport/hpack_encoder.cc

@@ -277,6 +277,7 @@ typedef struct {
   /* maximum size of a frame */
   size_t max_frame_size;
   bool use_true_binary_metadata;
+  bool is_end_of_stream;
 } framer_state;
 
 /* fills p (which is expected to be kDataFrameHeaderSize bytes long)
@@ -315,17 +316,29 @@ static size_t current_frame_size(framer_state* st) {
 }
 
 /* finish a frame - fill in the previously reserved header */
-static void finish_frame(framer_state* st, int is_header_boundary,
-                         int is_last_in_stream) {
+static void finish_frame(framer_state* st, int is_header_boundary) {
   uint8_t type = 0xff;
-  type = st->is_first_frame ? GRPC_CHTTP2_FRAME_HEADER
-                            : GRPC_CHTTP2_FRAME_CONTINUATION;
-  fill_header(
-      GRPC_SLICE_START_PTR(st->output->slices[st->header_idx]), type,
-      st->stream_id, current_frame_size(st),
-      static_cast<uint8_t>(
-          (is_last_in_stream ? GRPC_CHTTP2_DATA_FLAG_END_STREAM : 0) |
-          (is_header_boundary ? GRPC_CHTTP2_DATA_FLAG_END_HEADERS : 0)));
+  type =
+      static_cast<uint8_t>(st->is_first_frame ? GRPC_CHTTP2_FRAME_HEADER
+                                              : GRPC_CHTTP2_FRAME_CONTINUATION);
+  uint8_t flags = 0xff;
+  /* per the HTTP/2 spec:
+       A HEADERS frame carries the END_STREAM flag that signals the end of a
+       stream. However, a HEADERS frame with the END_STREAM flag set can be
+       followed by CONTINUATION frames on the same stream. Logically, the
+       CONTINUATION frames are part of the HEADERS frame.
+     Thus, we add the END_STREAM flag to the HEADER frame (the first frame). */
+  flags = static_cast<uint8_t>(st->is_first_frame && st->is_end_of_stream
+                                   ? GRPC_CHTTP2_DATA_FLAG_END_STREAM
+                                   : 0);
+  /* per the HTTP/2 spec:
+       A HEADERS frame without the END_HEADERS flag set MUST be followed by
+       a CONTINUATION frame for the same stream.
+     Thus, we add the END_HEADER flag to the last frame. */
+  flags |= static_cast<uint8_t>(
+      is_header_boundary ? GRPC_CHTTP2_DATA_FLAG_END_HEADERS : 0);
+  fill_header(GRPC_SLICE_START_PTR(st->output->slices[st->header_idx]), type,
+              st->stream_id, current_frame_size(st), flags);
   st->stats->framing_bytes += kDataFrameHeaderSize;
   st->is_first_frame = 0;
 }
@@ -347,7 +360,7 @@ static void ensure_space(framer_state* st, size_t need_bytes) {
   if (GPR_LIKELY(current_frame_size(st) + need_bytes <= st->max_frame_size)) {
     return;
   }
-  finish_frame(st, 0, 0);
+  finish_frame(st, 0);
   begin_frame(st);
 }
 
@@ -362,7 +375,7 @@ static void add_header_data(framer_state* st, grpc_slice slice) {
   } else {
     st->stats->header_bytes += remaining;
     grpc_slice_buffer_add(st->output, grpc_slice_split_head(&slice, remaining));
-    finish_frame(st, 0, 0);
+    finish_frame(st, 0);
     begin_frame(st);
     add_header_data(st, slice);
   }
@@ -841,6 +854,7 @@ void grpc_chttp2_encode_header(grpc_chttp2_hpack_compressor* c,
   st.stats = options->stats;
   st.max_frame_size = options->max_frame_size;
   st.use_true_binary_metadata = options->use_true_binary_metadata;
+  st.is_end_of_stream = options->is_eof;
 
   /* Encode a metadata batch; store the returned values, representing
      a metadata element that needs to be unreffed back into the metadata
@@ -883,5 +897,5 @@ void grpc_chttp2_encode_header(grpc_chttp2_hpack_compressor* c,
     deadline_enc(c, deadline, &st);
   }
 
-  finish_frame(&st, 1, options->is_eof);
+  finish_frame(&st, 1);
 }

test/core/transport/chttp2/hpack_encoder_test.cc

@@ -49,6 +49,102 @@ typedef struct {
   bool only_intern_key;
 } verify_params;
 
+/* verify that the output frames that are generated by encoding the stream
+   have sensible type and flags values */
+static void verify_frames(grpc_slice_buffer& output, bool header_is_eof) {
+  /* per the HTTP/2 spec:
+       All frames begin with a fixed 9-octet header followed by a
+       variable-length payload.
+
+       +-----------------------------------------------+
+       |                 Length (24)                   |
+       +---------------+---------------+---------------+
+       |   Type (8)    |   Flags (8)   |
+       +-+-------------+---------------+-------------------------------+
+       |R|                 Stream Identifier (31)                      |
+       +=+=============================================================+
+       |                   Frame Payload (0...)                      ...
+       +---------------------------------------------------------------+
+   */
+  uint8_t type = 0xff, flags = 0xff;
+  size_t i, merged_length, frame_size;
+  bool first_frame = false;
+  bool in_header = false;
+  bool end_header = false;
+  bool is_closed = false;
+  for (i = 0; i < output.count;) {
+    first_frame = i == 0;
+    grpc_slice* slice = &output.slices[i++];
+
+    // Read gRPC frame header
+    uint8_t* p = GRPC_SLICE_START_PTR(*slice);
+    frame_size = 0;
+    frame_size |= static_cast<uint32_t>(p[0]) << 16;
+    frame_size |= static_cast<uint32_t>(p[1]) << 8;
+    frame_size |= static_cast<uint32_t>(p[2]);
+    type = p[3];
+    flags = p[4];
+
+    // Read remainder of the gRPC frame
+    merged_length = GRPC_SLICE_LENGTH(*slice);
+    while (merged_length < frame_size + 9) {  // including 9 byte frame header
+      grpc_slice* slice = &output.slices[i++];
+      merged_length += GRPC_SLICE_LENGTH(*slice);
+    }
+
+    // Verifications
+    if (first_frame && type != GRPC_CHTTP2_FRAME_HEADER) {
+      gpr_log(GPR_ERROR, "expected first frame to be of type header");
+      gpr_log(GPR_ERROR, "EXPECT: 0x%x", GRPC_CHTTP2_FRAME_HEADER);
+      gpr_log(GPR_ERROR, "GOT:    0x%x", type);
+      g_failure = 1;
+    } else if (first_frame && header_is_eof &&
+               !(flags & GRPC_CHTTP2_DATA_FLAG_END_STREAM)) {
+      gpr_log(GPR_ERROR, "missing END_STREAM flag in HEADER frame");
+      g_failure = 1;
+    }
+    if (is_closed &&
+        (type == GRPC_CHTTP2_FRAME_DATA || type == GRPC_CHTTP2_FRAME_HEADER)) {
+      gpr_log(GPR_ERROR,
+              "stream is closed; new frame headers and data are not allowed");
+      g_failure = 1;
+    }
+    if (end_header && (type == GRPC_CHTTP2_FRAME_HEADER ||
+                       type == GRPC_CHTTP2_FRAME_CONTINUATION)) {
+      gpr_log(GPR_ERROR,
+              "frame header is ended; new headers and continuations are not "
+              "allowed");
+      g_failure = 1;
+    }
+    if (in_header &&
+        (type == GRPC_CHTTP2_FRAME_DATA || type == GRPC_CHTTP2_FRAME_HEADER)) {
+      gpr_log(GPR_ERROR,
+              "parsing frame header; new headers and data are not allowed");
+      g_failure = 1;
+    }
+    if (flags & ~(GRPC_CHTTP2_DATA_FLAG_END_STREAM |
+                  GRPC_CHTTP2_DATA_FLAG_END_HEADERS)) {
+      gpr_log(GPR_ERROR, "unexpected frame flags: 0x%x", flags);
+      g_failure = 1;
+    }
+
+    // Update state
+    if (flags & GRPC_CHTTP2_DATA_FLAG_END_HEADERS) {
+      in_header = false;
+      end_header = true;
+    } else if (type == GRPC_CHTTP2_DATA_FLAG_END_HEADERS) {
+      in_header = true;
+    }
+    if (flags & GRPC_CHTTP2_DATA_FLAG_END_STREAM) {
+      is_closed = true;
+      if (type == GRPC_CHTTP2_FRAME_CONTINUATION) {
+        gpr_log(GPR_ERROR, "unexpected END_STREAM flag in CONTINUATION frame");
+        g_failure = 1;
+      }
+    }
+  }
+}
+
 /* verify that the output generated by encoding the stream matches the
    hexstring passed in */
 static void verify(const verify_params params, const char* expected,
@@ -106,6 +202,7 @@ static void verify(const verify_params params, const char* expected,
       &stats                           /* stats */
   };
   grpc_chttp2_encode_header(&g_compressor, nullptr, 0, &b, &hopt, &output);
+  verify_frames(output, params.eof);
   merged = grpc_slice_merge(output.slices, output.count);
   grpc_slice_buffer_destroy_internal(&output);
   grpc_metadata_batch_destroy(&b);
@@ -151,6 +248,50 @@ static void test_basic_headers() {
   verify(params, "000004 0104 deadbeef 0f 2f 0176", 1, "a", "v");
 }
 
+static void verify_continuation_headers(const char* key, const char* value,
+                                        bool is_eof) {
+  grpc_slice_buffer output;
+  grpc_mdelem elem = grpc_mdelem_from_slices(
+      grpc_slice_intern(grpc_slice_from_static_string(key)),
+      grpc_slice_intern(grpc_slice_from_static_string(value)));
+  grpc_linked_mdelem* e =
+      static_cast<grpc_linked_mdelem*>(gpr_malloc(sizeof(*e)));
+  grpc_metadata_batch b;
+  grpc_metadata_batch_init(&b);
+  e[0].md = elem;
+  e[0].prev = nullptr;
+  e[0].next = nullptr;
+  b.list.head = &e[0];
+  b.list.tail = &e[0];
+  b.list.count = 1;
+  grpc_slice_buffer_init(&output);
+
+  grpc_transport_one_way_stats stats;
+  stats = {};
+  grpc_encode_header_options hopt = {0xdeadbeef, /* stream_id */
+                                     is_eof,     /* is_eof */
+                                     false,      /* use_true_binary_metadata */
+                                     150,        /* max_frame_size */
+                                     &stats /* stats */};
+  grpc_chttp2_encode_header(&g_compressor, nullptr, 0, &b, &hopt, &output);
+  verify_frames(output, is_eof);
+  grpc_slice_buffer_destroy_internal(&output);
+  grpc_metadata_batch_destroy(&b);
+  gpr_free(e);
+}
+
+static void test_continuation_headers() {
+  char value[200];
+  memset(value, 'a', 200);
+  value[199] = 0;  // null terminator
+  verify_continuation_headers("key", value, true);
+
+  char value2[400];
+  memset(value2, 'b', 400);
+  value2[399] = 0;  // null terminator
+  verify_continuation_headers("key2", value2, true);
+}
+
 static void encode_int_to_str(int i, char* p) {
   p[0] = static_cast<char>('a' + i % 26);
   i /= 26;
@@ -225,6 +366,7 @@ static void verify_table_size_change_match_elem_size(const char* key,
       16384,           /* max_frame_size */
       &stats /* stats */};
   grpc_chttp2_encode_header(&g_compressor, nullptr, 0, &b, &hopt, &output);
+  verify_frames(output, false);
   grpc_slice_buffer_destroy_internal(&output);
   grpc_metadata_batch_destroy(&b);
 
@@ -267,6 +409,7 @@ int main(int argc, char** argv) {
   TEST(test_decode_table_overflow);
   TEST(test_encode_header_size);
   TEST(test_interned_key_indexed);
+  TEST(test_continuation_headers);
   grpc_shutdown();
   for (i = 0; i < num_to_delete; i++) {
     gpr_free(to_delete[i]);