Merge pull request #1211 from nghttp2/stream-user-data

Tweak nghttp2_session_set_stream_user_data

lib/includes/nghttp2/nghttp2.h

@@ -3802,10 +3802,13 @@ nghttp2_priority_spec_check_default(const nghttp2_priority_spec *pri_spec);
  * .. warning::
  *
  *   This function returns assigned stream ID if it succeeds.  But
- *   that stream is not opened yet.  The application must not submit
+ *   that stream is not created yet.  The application must not submit
  *   frame to that stream ID before
  *   :type:`nghttp2_before_frame_send_callback` is called for this
- *   frame.
+ *   frame.  This means `nghttp2_session_get_stream_user_data()` does
+ *   not work before the callback.  But
+ *   `nghttp2_session_set_stream_user_data()` handles this situation
+ *   specially, and it can set data to a stream during this period.
  *
  */
 NGHTTP2_EXTERN int32_t nghttp2_submit_request(

lib/nghttp2_session.c

@@ -7208,14 +7208,44 @@ int nghttp2_session_set_stream_user_data(nghttp2_session *session,
                                          int32_t stream_id,
                                          void *stream_user_data) {
   nghttp2_stream *stream;
+  nghttp2_frame *frame;
+  nghttp2_outbound_item *item;
+
   stream = nghttp2_session_get_stream(session, stream_id);
-  if (!stream) {
-    return NGHTTP2_ERR_INVALID_ARGUMENT;
-  }
+  if (stream) {
     stream->stream_user_data = stream_user_data;
     return 0;
   }
 
+  if (session->server || !nghttp2_session_is_my_stream_id(session, stream_id) ||
+      !nghttp2_outbound_queue_top(&session->ob_syn)) {
+    return NGHTTP2_ERR_INVALID_ARGUMENT;
+  }
+
+  frame = &nghttp2_outbound_queue_top(&session->ob_syn)->frame;
+  assert(frame->hd.type == NGHTTP2_HEADERS);
+
+  if (frame->hd.stream_id > stream_id ||
+      (uint32_t)stream_id >= session->next_stream_id) {
+    return NGHTTP2_ERR_INVALID_ARGUMENT;
+  }
+
+  for (item = session->ob_syn.head; item; item = item->qnext) {
+    if (item->frame.hd.stream_id < stream_id) {
+      continue;
+    }
+
+    if (item->frame.hd.stream_id > stream_id) {
+      break;
+    }
+
+    item->aux_data.headers.stream_user_data = stream_user_data;
+    return 0;
+  }
+
+  return NGHTTP2_ERR_INVALID_ARGUMENT;
+}
+
 int nghttp2_session_resume_data(nghttp2_session *session, int32_t stream_id) {
   int rv;
   nghttp2_stream *stream;

tests/main.c

@@ -319,6 +319,8 @@ int main() {
                    test_nghttp2_session_pause_data) ||
       !CU_add_test(pSuite, "session_no_closed_streams",
                    test_nghttp2_session_no_closed_streams) ||
+      !CU_add_test(pSuite, "session_set_stream_user_data",
+                   test_nghttp2_session_set_stream_user_data) ||
       !CU_add_test(pSuite, "http_mandatory_headers",
                    test_nghttp2_http_mandatory_headers) ||
       !CU_add_test(pSuite, "http_content_length",

tests/nghttp2_session_test.c

@@ -10696,6 +10696,39 @@ void test_nghttp2_session_no_closed_streams(void) {
   nghttp2_option_del(option);
 }
 
+void test_nghttp2_session_set_stream_user_data(void) {
+  nghttp2_session *session;
+  nghttp2_session_callbacks callbacks;
+  int32_t stream_id;
+  int user_data1, user_data2;
+  int rv;
+  const uint8_t *datap;
+  ssize_t datalen;
+
+  memset(&callbacks, 0, sizeof(nghttp2_session_callbacks));
+
+  nghttp2_session_client_new(&session, &callbacks, NULL);
+
+  stream_id = nghttp2_submit_request(session, NULL, reqnv, ARRLEN(reqnv), NULL,
+                                     &user_data1);
+
+  rv = nghttp2_session_set_stream_user_data(session, stream_id, &user_data2);
+
+  CU_ASSERT(0 == rv);
+
+  datalen = nghttp2_session_mem_send(session, &datap);
+
+  CU_ASSERT(datalen > 0);
+
+  CU_ASSERT(&user_data2 ==
+            nghttp2_session_get_stream_user_data(session, stream_id));
+
+  CU_ASSERT(NGHTTP2_ERR_INVALID_ARGUMENT ==
+            nghttp2_session_set_stream_user_data(session, 2, NULL));
+
+  nghttp2_session_del(session);
+}
+
 static void check_nghttp2_http_recv_headers_fail(
     nghttp2_session *session, nghttp2_hd_deflater *deflater, int32_t stream_id,
     int stream_state, const nghttp2_nv *nva, size_t nvlen) {

tests/nghttp2_session_test.h

@@ -158,6 +158,7 @@ void test_nghttp2_session_cancel_from_before_frame_send(void);
 void test_nghttp2_session_removed_closed_stream(void);
 void test_nghttp2_session_pause_data(void);
 void test_nghttp2_session_no_closed_streams(void);
+void test_nghttp2_session_set_stream_user_data(void);
 void test_nghttp2_http_mandatory_headers(void);
 void test_nghttp2_http_content_length(void);
 void test_nghttp2_http_content_length_mismatch(void);