Adds methods to the Http2Adapter API to get and set per-stream user data.

This fills one of the few remaining API gaps.

PiperOrigin-RevId: 378197532
diff --git a/http2/adapter/http2_adapter.h b/http2/adapter/http2_adapter.h
index 293e2fb..fc64c0b 100644
--- a/http2/adapter/http2_adapter.h
+++ b/http2/adapter/http2_adapter.h
@@ -102,6 +102,15 @@
   virtual int SubmitTrailer(Http2StreamId stream_id,
                             absl::Span<const Header> trailers) = 0;
 
+  // Sets a user data pointer for the given stream. Can be called after
+  // SubmitRequest/SubmitResponse, or after receiving any frame for a given
+  // stream.
+  virtual void SetStreamUserData(Http2StreamId stream_id, void* user_data) = 0;
+
+  // Returns nullptr if the stream does not exist, or if stream user data has
+  // not been set.
+  virtual void* GetStreamUserData(Http2StreamId stream_id) = 0;
+
  protected:
   // Subclasses should expose a public factory method for constructing and
   // initializing (via Initialize()) adapter instances.
diff --git a/http2/adapter/nghttp2_adapter.cc b/http2/adapter/nghttp2_adapter.cc
index deb4468..31b215c 100644
--- a/http2/adapter/nghttp2_adapter.cc
+++ b/http2/adapter/nghttp2_adapter.cc
@@ -157,6 +157,16 @@
                                 nvs.size());
 }
 
+void NgHttp2Adapter::SetStreamUserData(Http2StreamId stream_id,
+                                       void* user_data) {
+  nghttp2_session_set_stream_user_data(session_->raw_ptr(), stream_id,
+                                       user_data);
+}
+
+void* NgHttp2Adapter::GetStreamUserData(Http2StreamId stream_id) {
+  return nghttp2_session_get_stream_user_data(session_->raw_ptr(), stream_id);
+}
+
 NgHttp2Adapter::NgHttp2Adapter(Http2VisitorInterface& visitor,
                                Perspective perspective)
     : Http2Adapter(visitor), visitor_(visitor), perspective_(perspective) {}
diff --git a/http2/adapter/nghttp2_adapter.h b/http2/adapter/nghttp2_adapter.h
index d37e3c9..b8dfb3b 100644
--- a/http2/adapter/nghttp2_adapter.h
+++ b/http2/adapter/nghttp2_adapter.h
@@ -67,6 +67,9 @@
   int SubmitTrailer(Http2StreamId stream_id,
                     absl::Span<const Header> trailers) override;
 
+  void SetStreamUserData(Http2StreamId stream_id, void* user_data) override;
+  void* GetStreamUserData(Http2StreamId stream_id) override;
+
   // TODO(b/181586191): Temporary accessor until equivalent functionality is
   // available in this adapter class.
   NgHttp2Session& session() { return *session_; }
diff --git a/http2/adapter/nghttp2_adapter_test.cc b/http2/adapter/nghttp2_adapter_test.cc
index 2e2a1bf..ded57cc 100644
--- a/http2/adapter/nghttp2_adapter_test.cc
+++ b/http2/adapter/nghttp2_adapter_test.cc
@@ -92,9 +92,11 @@
                  {":path", "/this/is/request/three"}});
   const auto nvs3 = GetNghttp2Nvs(headers3);
 
-  const int32_t stream_id1 =
-      nghttp2_submit_request(adapter->session().raw_ptr(), nullptr, nvs1.data(),
-                             nvs1.size(), nullptr, nullptr);
+  const char* kSentinel1 = "arbitrary pointer 1";
+  const char* kSentinel3 = "arbitrary pointer 3";
+  const int32_t stream_id1 = nghttp2_submit_request(
+      adapter->session().raw_ptr(), nullptr, nvs1.data(), nvs1.size(), nullptr,
+      const_cast<char*>(kSentinel1));
   ASSERT_GT(stream_id1, 0);
   QUICHE_LOG(INFO) << "Created stream: " << stream_id1;
 
@@ -104,12 +106,16 @@
   ASSERT_GT(stream_id2, 0);
   QUICHE_LOG(INFO) << "Created stream: " << stream_id2;
 
-  const int32_t stream_id3 =
-      nghttp2_submit_request(adapter->session().raw_ptr(), nullptr, nvs3.data(),
-                             nvs3.size(), nullptr, nullptr);
+  const int32_t stream_id3 = nghttp2_submit_request(
+      adapter->session().raw_ptr(), nullptr, nvs3.data(), nvs3.size(), nullptr,
+      const_cast<char*>(kSentinel3));
   ASSERT_GT(stream_id3, 0);
   QUICHE_LOG(INFO) << "Created stream: " << stream_id3;
 
+  const char* kSentinel2 = "arbitrary pointer 2";
+  adapter->SetStreamUserData(stream_id2, const_cast<char*>(kSentinel2));
+  adapter->SetStreamUserData(stream_id3, nullptr);
+
   adapter->Send();
   EXPECT_THAT(visitor.data(), EqualsFrames({spdy::SpdyFrameType::HEADERS,
                                             spdy::SpdyFrameType::HEADERS,
@@ -118,6 +124,10 @@
 
   EXPECT_EQ(0, adapter->GetHighestReceivedStreamId());
 
+  EXPECT_EQ(kSentinel1, adapter->GetStreamUserData(stream_id1));
+  EXPECT_EQ(kSentinel2, adapter->GetStreamUserData(stream_id2));
+  EXPECT_EQ(nullptr, adapter->GetStreamUserData(stream_id3));
+
   const std::string stream_frames =
       TestFrameSequence()
           .Headers(1,
@@ -332,6 +342,8 @@
                                  .Serialize();
   testing::InSequence s;
 
+  const char* kSentinel1 = "arbitrary pointer 1";
+
   // Client preface (empty SETTINGS)
   EXPECT_CALL(visitor, OnFrameHeader(0, 0, SETTINGS, 0));
   EXPECT_CALL(visitor, OnSettingsStart());
@@ -347,7 +359,10 @@
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":scheme", "https"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":authority", "example.com"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":path", "/this/is/request/one"));
-  EXPECT_CALL(visitor, OnEndHeadersForStream(1));
+  EXPECT_CALL(visitor, OnEndHeadersForStream(1))
+      .WillOnce(testing::InvokeWithoutArgs([&adapter, kSentinel1]() {
+        adapter->SetStreamUserData(1, const_cast<char*>(kSentinel1));
+      }));
   EXPECT_CALL(visitor, OnFrameHeader(1, 4, WINDOW_UPDATE, 0));
   EXPECT_CALL(visitor, OnWindowUpdate(1, 2000));
   EXPECT_CALL(visitor, OnFrameHeader(1, 25, DATA, 0));
@@ -370,6 +385,14 @@
   const ssize_t result = adapter->ProcessBytes(frames);
   EXPECT_EQ(frames.size(), result);
 
+  EXPECT_EQ(kSentinel1, adapter->GetStreamUserData(1));
+
+  // Because stream 3 has already been closed, it's not possible to set user
+  // data.
+  const char* kSentinel3 = "another arbitrary pointer";
+  adapter->SetStreamUserData(3, const_cast<char*>(kSentinel3));
+  EXPECT_EQ(nullptr, adapter->GetStreamUserData(3));
+
   EXPECT_EQ(3, adapter->GetHighestReceivedStreamId());
 
   EXPECT_EQ(adapter->GetPeerConnectionWindow(),
@@ -400,6 +423,8 @@
                                  .Serialize();
   testing::InSequence s;
 
+  const char* kSentinel1 = "arbitrary pointer 1";
+
   // Client preface (empty SETTINGS)
   EXPECT_CALL(visitor, OnFrameHeader(0, 0, SETTINGS, 0));
   EXPECT_CALL(visitor, OnSettingsStart());
@@ -411,7 +436,10 @@
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":scheme", "https"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":authority", "example.com"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":path", "/this/is/request/one"));
-  EXPECT_CALL(visitor, OnEndHeadersForStream(1));
+  EXPECT_CALL(visitor, OnEndHeadersForStream(1))
+      .WillOnce(testing::InvokeWithoutArgs([&adapter, kSentinel1]() {
+        adapter->SetStreamUserData(1, const_cast<char*>(kSentinel1));
+      }));
   EXPECT_CALL(visitor, OnEndStream(1));
 
   const ssize_t result = adapter->ProcessBytes(frames);
@@ -435,8 +463,15 @@
       &body1);
   EXPECT_EQ(submit_result, 0);
   EXPECT_TRUE(adapter->session().want_write());
+
+  // Stream user data should have been set successfully after receiving headers.
+  EXPECT_EQ(kSentinel1, adapter->GetStreamUserData(1));
+  adapter->SetStreamUserData(1, nullptr);
+  EXPECT_EQ(nullptr, adapter->GetStreamUserData(1));
+
   EXPECT_CALL(visitor, OnCloseStream(1, Http2ErrorCode::NO_ERROR));
   adapter->Send();
+
   EXPECT_THAT(visitor.data(), EqualsFrames({spdy::SpdyFrameType::HEADERS,
                                             spdy::SpdyFrameType::DATA}));
   EXPECT_THAT(visitor.data(), testing::HasSubstr(kBody));
diff --git a/http2/adapter/oghttp2_adapter.cc b/http2/adapter/oghttp2_adapter.cc
index a71a71e..1ba8d37 100644
--- a/http2/adapter/oghttp2_adapter.cc
+++ b/http2/adapter/oghttp2_adapter.cc
@@ -119,6 +119,15 @@
   return session_->SubmitTrailer(stream_id, trailers);
 }
 
+void OgHttp2Adapter::SetStreamUserData(Http2StreamId stream_id,
+                                       void* user_data) {
+  session_->SetStreamUserData(stream_id, user_data);
+}
+
+void* OgHttp2Adapter::GetStreamUserData(Http2StreamId stream_id) {
+  return session_->GetStreamUserData(stream_id);
+}
+
 const Http2Session& OgHttp2Adapter::session() const {
   return *session_;
 }
diff --git a/http2/adapter/oghttp2_adapter.h b/http2/adapter/oghttp2_adapter.h
index 1b9d90d..7ef07f6 100644
--- a/http2/adapter/oghttp2_adapter.h
+++ b/http2/adapter/oghttp2_adapter.h
@@ -50,6 +50,9 @@
   int SubmitTrailer(Http2StreamId stream_id,
                     absl::Span<const Header> trailers) override;
 
+  void SetStreamUserData(Http2StreamId stream_id, void* user_data) override;
+  void* GetStreamUserData(Http2StreamId stream_id) override;
+
   const Http2Session& session() const;
 
  private:
diff --git a/http2/adapter/oghttp2_session.cc b/http2/adapter/oghttp2_session.cc
index aa135ff..330476c 100644
--- a/http2/adapter/oghttp2_session.cc
+++ b/http2/adapter/oghttp2_session.cc
@@ -33,6 +33,22 @@
 
 OgHttp2Session::~OgHttp2Session() {}
 
+void OgHttp2Session::SetStreamUserData(Http2StreamId stream_id,
+                                       void* user_data) {
+  auto it = stream_map_.find(stream_id);
+  if (it != stream_map_.end()) {
+    it->second.user_data = user_data;
+  }
+}
+
+void* OgHttp2Session::GetStreamUserData(Http2StreamId stream_id) {
+  auto it = stream_map_.find(stream_id);
+  if (it != stream_map_.end()) {
+    return it->second.user_data;
+  }
+  return nullptr;
+}
+
 ssize_t OgHttp2Session::ProcessBytes(absl::string_view bytes) {
   ssize_t preface_consumed = 0;
   if (!remaining_preface_.empty()) {
diff --git a/http2/adapter/oghttp2_session.h b/http2/adapter/oghttp2_session.h
index 8655820..30bc9d8 100644
--- a/http2/adapter/oghttp2_session.h
+++ b/http2/adapter/oghttp2_session.h
@@ -51,6 +51,8 @@
   Http2StreamId GetHighestReceivedStreamId() const {
     return highest_received_stream_id_;
   }
+  void SetStreamUserData(Http2StreamId stream_id, void* user_data);
+  void* GetStreamUserData(Http2StreamId stream_id);
 
   // From Http2Session.
   ssize_t ProcessBytes(absl::string_view bytes) override;
diff --git a/http2/adapter/oghttp2_session_test.cc b/http2/adapter/oghttp2_session_test.cc
index c6525fa..185fd3a 100644
--- a/http2/adapter/oghttp2_session_test.cc
+++ b/http2/adapter/oghttp2_session_test.cc
@@ -187,20 +187,38 @@
   EXPECT_THAT(visitor.data(), EqualsFrames({SpdyFrameType::SETTINGS}));
   visitor.Clear();
 
-  const char* kSentinel = "";
+  const char* kSentinel1 = "arbitrary pointer 1";
   TestDataFrameSource body1(visitor, "This is an example request body.");
   int stream_id =
       session.SubmitRequest(ToHeaders({{":method", "POST"},
                                        {":scheme", "http"},
                                        {":authority", "example.com"},
                                        {":path", "/this/is/request/one"}}),
-                            &body1, const_cast<char*>(kSentinel));
+                            &body1, const_cast<char*>(kSentinel1));
   EXPECT_GT(stream_id, 0);
   EXPECT_TRUE(session.want_write());
+  EXPECT_EQ(kSentinel1, session.GetStreamUserData(stream_id));
   session.Send();
   EXPECT_THAT(visitor.data(), EqualsFrames({spdy::SpdyFrameType::HEADERS,
                                             spdy::SpdyFrameType::DATA}));
+  visitor.Clear();
   EXPECT_FALSE(session.want_write());
+
+  stream_id =
+      session.SubmitRequest(ToHeaders({{":method", "POST"},
+                                       {":scheme", "http"},
+                                       {":authority", "example.com"},
+                                       {":path", "/this/is/request/one"}}),
+                            nullptr, nullptr);
+  EXPECT_GT(stream_id, 0);
+  EXPECT_TRUE(session.want_write());
+  const char* kSentinel2 = "arbitrary pointer 2";
+  EXPECT_EQ(nullptr, session.GetStreamUserData(stream_id));
+  session.SetStreamUserData(stream_id, const_cast<char*>(kSentinel2));
+  EXPECT_EQ(kSentinel2, session.GetStreamUserData(stream_id));
+
+  session.Send();
+  EXPECT_THAT(visitor.data(), EqualsFrames({spdy::SpdyFrameType::HEADERS}));
 }
 
 TEST(OgHttp2SessionTest, ClientStartShutdown) {
@@ -234,7 +252,7 @@
 }
 
 TEST(OgHttp2SessionTest, ServerHandlesFrames) {
-  testing::StrictMock<MockHttp2Visitor> visitor;
+  DataSavingVisitor visitor;
   OgHttp2Session session(
       visitor, OgHttp2Session::Options{.perspective = Perspective::kServer});
 
@@ -261,6 +279,8 @@
                                  .Serialize();
   testing::InSequence s;
 
+  const char* kSentinel1 = "arbitrary pointer 1";
+
   // Client preface (empty SETTINGS)
   EXPECT_CALL(visitor, OnFrameHeader(0, 0, SETTINGS, 0));
   EXPECT_CALL(visitor, OnSettingsStart());
@@ -276,7 +296,10 @@
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":scheme", "https"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":authority", "example.com"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":path", "/this/is/request/one"));
-  EXPECT_CALL(visitor, OnEndHeadersForStream(1));
+  EXPECT_CALL(visitor, OnEndHeadersForStream(1))
+      .WillOnce(testing::InvokeWithoutArgs([&session, kSentinel1]() {
+        session.SetStreamUserData(1, const_cast<char*>(kSentinel1));
+      }));
   EXPECT_CALL(visitor, OnFrameHeader(1, 4, WINDOW_UPDATE, 0));
   EXPECT_CALL(visitor, OnWindowUpdate(1, 2000));
   EXPECT_CALL(visitor, OnFrameHeader(1, 25, DATA, 0));
@@ -299,9 +322,25 @@
   const ssize_t result = session.ProcessBytes(frames);
   EXPECT_EQ(frames.size(), result);
 
+  EXPECT_EQ(kSentinel1, session.GetStreamUserData(1));
+
+  // TODO(birenroy): drop stream state when streams are closed. It should no
+  // longer be possible to set user data.
+  const char* kSentinel3 = "another arbitrary pointer";
+  session.SetStreamUserData(3, const_cast<char*>(kSentinel3));
+  EXPECT_EQ(kSentinel3, session.GetStreamUserData(3));
+
   EXPECT_EQ(session.GetRemoteWindowSize(),
             kDefaultInitialStreamWindowSize + 1000);
   EXPECT_EQ(3, session.GetHighestReceivedStreamId());
+
+  EXPECT_TRUE(session.want_write());
+  // Some bytes should have been serialized.
+  session.Send();
+  // Initial SETTINGS, SETTINGS ack.
+  // TODO(birenroy): automatically queue PING acks.
+  EXPECT_THAT(visitor.data(), EqualsFrames({spdy::SpdyFrameType::SETTINGS,
+                                            spdy::SpdyFrameType::SETTINGS}));
 }
 
 // Verifies that a server session enqueues initial SETTINGS before whatever
@@ -349,6 +388,8 @@
                                  .Serialize();
   testing::InSequence s;
 
+  const char* kSentinel1 = "arbitrary pointer 1";
+
   // Client preface (empty SETTINGS)
   EXPECT_CALL(visitor, OnFrameHeader(0, 0, SETTINGS, 0));
   EXPECT_CALL(visitor, OnSettingsStart());
@@ -360,7 +401,10 @@
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":scheme", "https"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":authority", "example.com"));
   EXPECT_CALL(visitor, OnHeaderForStream(1, ":path", "/this/is/request/one"));
-  EXPECT_CALL(visitor, OnEndHeadersForStream(1));
+  EXPECT_CALL(visitor, OnEndHeadersForStream(1))
+      .WillOnce(testing::InvokeWithoutArgs([&session, kSentinel1]() {
+        session.SetStreamUserData(1, const_cast<char*>(kSentinel1));
+      }));
   EXPECT_CALL(visitor, OnEndStream(1));
 
   const ssize_t result = session.ProcessBytes(frames);
@@ -384,6 +428,12 @@
       &body1);
   EXPECT_EQ(submit_result, 0);
   EXPECT_TRUE(session.want_write());
+
+  // Stream user data should have been set successfully after receiving headers.
+  EXPECT_EQ(kSentinel1, session.GetStreamUserData(1));
+  session.SetStreamUserData(1, nullptr);
+  EXPECT_EQ(nullptr, session.GetStreamUserData(1));
+
   EXPECT_CALL(visitor, OnCloseStream(1, Http2ErrorCode::NO_ERROR));
   session.Send();
   EXPECT_THAT(visitor.data(),