Implement webtransport::Session::GetStreamById

PiperOrigin-RevId: 513922467
diff --git a/quiche/quic/core/http/end_to_end_test.cc b/quiche/quic/core/http/end_to_end_test.cc
index 1f8d6fe..0ee3cf0 100644
--- a/quiche/quic/core/http/end_to_end_test.cc
+++ b/quiche/quic/core/http/end_to_end_test.cc
@@ -6603,6 +6603,8 @@
   WebTransportStream* outgoing_stream =
       session->OpenOutgoingUnidirectionalStream();
   ASSERT_TRUE(outgoing_stream != nullptr);
+  EXPECT_EQ(outgoing_stream,
+            session->GetStreamById(outgoing_stream->GetStreamId()));
 
   auto stream_visitor =
       std::make_unique<NiceMock<MockWebTransportStreamVisitor>>();
@@ -6622,6 +6624,8 @@
   WebTransportStream* received_stream =
       session->AcceptIncomingUnidirectionalStream();
   ASSERT_TRUE(received_stream != nullptr);
+  EXPECT_EQ(received_stream,
+            session->GetStreamById(received_stream->GetStreamId()));
   std::string received_data;
   WebTransportStream::ReadResult result = received_stream->Read(&received_data);
   EXPECT_EQ(received_data, "test");
@@ -6681,6 +6685,7 @@
 
   WebTransportStream* stream = session->OpenOutgoingBidirectionalStream();
   ASSERT_TRUE(stream != nullptr);
+  EXPECT_EQ(stream, session->GetStreamById(stream->GetStreamId()));
 
   auto stream_visitor_owned =
       std::make_unique<NiceMock<MockWebTransportStreamVisitor>>();
diff --git a/quiche/quic/core/http/quic_spdy_stream.h b/quiche/quic/core/http/quic_spdy_stream.h
index e59afd5..1267f4c 100644
--- a/quiche/quic/core/http/quic_spdy_stream.h
+++ b/quiche/quic/core/http/quic_spdy_stream.h
@@ -27,6 +27,7 @@
 #include "quiche/quic/core/qpack/qpack_decoded_headers_accumulator.h"
 #include "quiche/quic/core/quic_error_codes.h"
 #include "quiche/quic/core/quic_packets.h"
+#include "quiche/quic/core/quic_session.h"
 #include "quiche/quic/core/quic_stream.h"
 #include "quiche/quic/core/quic_stream_priority.h"
 #include "quiche/quic/core/quic_stream_sequencer.h"
diff --git a/quiche/quic/core/http/web_transport_http3.cc b/quiche/quic/core/http/web_transport_http3.cc
index 1896266..dc0e054 100644
--- a/quiche/quic/core/http/web_transport_http3.cc
+++ b/quiche/quic/core/http/web_transport_http3.cc
@@ -267,6 +267,22 @@
   return stream->interface();
 }
 
+webtransport::Stream* WebTransportHttp3::GetStreamById(
+    webtransport::StreamId id) {
+  if (!streams_.contains(id)) {
+    return nullptr;
+  }
+  QuicStream* stream = session_->GetActiveStream(id);
+  const bool bidi = QuicUtils::IsBidirectionalStreamId(
+      id, ParsedQuicVersion::RFCv1());  // Assume IETF QUIC for WebTransport
+  if (bidi) {
+    return static_cast<QuicSpdyStream*>(stream)->web_transport_stream();
+  } else {
+    return static_cast<WebTransportHttp3UnidirectionalStream*>(stream)
+        ->interface();
+  }
+}
+
 webtransport::DatagramStatus WebTransportHttp3::SendOrQueueDatagram(
     absl::string_view datagram) {
   return MessageStatusToWebTransportStatus(
diff --git a/quiche/quic/core/http/web_transport_http3.h b/quiche/quic/core/http/web_transport_http3.h
index 8744449..6eb7239 100644
--- a/quiche/quic/core/http/web_transport_http3.h
+++ b/quiche/quic/core/http/web_transport_http3.h
@@ -82,6 +82,8 @@
   WebTransportStream* OpenOutgoingBidirectionalStream() override;
   WebTransportStream* OpenOutgoingUnidirectionalStream() override;
 
+  webtransport::Stream* GetStreamById(webtransport::StreamId id) override;
+
   webtransport::DatagramStatus SendOrQueueDatagram(
       absl::string_view datagram) override;
   QuicByteCount GetMaxDatagramSize() const override;
diff --git a/quiche/quic/core/quic_session.h b/quiche/quic/core/quic_session.h
index 5fa690b..c9931fe 100644
--- a/quiche/quic/core/quic_session.h
+++ b/quiche/quic/core/quic_session.h
@@ -649,6 +649,11 @@
     datagram_queue_.SetForceFlush(force_flush);
   }
 
+  // Find stream with |id|, returns nullptr if the stream does not exist or
+  // closed. static streams and zombie streams are not considered active
+  // streams.
+  QuicStream* GetActiveStream(QuicStreamId id) const;
+
  protected:
   using StreamMap =
       absl::flat_hash_map<QuicStreamId, std::unique_ptr<QuicStream>>;
@@ -810,11 +815,6 @@
     connection()->SetLossDetectionTuner(std::move(tuner));
   }
 
-  // Find stream with |id|, returns nullptr if the stream does not exist or
-  // closed. static streams and zombie streams are not considered active
-  // streams.
-  QuicStream* GetActiveStream(QuicStreamId id) const;
-
   const UberQuicStreamIdManager& ietf_streamid_manager() const {
     QUICHE_DCHECK(VersionHasIetfQuicFrames(transport_version()));
     return ietf_streamid_manager_;
diff --git a/quiche/web_transport/test_tools/mock_web_transport.h b/quiche/web_transport/test_tools/mock_web_transport.h
index a24cc7a..9fdecfc 100644
--- a/quiche/web_transport/test_tools/mock_web_transport.h
+++ b/quiche/web_transport/test_tools/mock_web_transport.h
@@ -66,6 +66,7 @@
   MOCK_METHOD(bool, CanOpenNextOutgoingUnidirectionalStream, (), (override));
   MOCK_METHOD(Stream*, OpenOutgoingBidirectionalStream, (), (override));
   MOCK_METHOD(Stream*, OpenOutgoingUnidirectionalStream, (), (override));
+  MOCK_METHOD(Stream*, GetStreamById, (StreamId), (override));
   MOCK_METHOD(DatagramStatus, SendOrQueueDatagram, (absl::string_view datagram),
               (override));
   MOCK_METHOD(uint64_t, GetMaxDatagramSize, (), (const, override));
diff --git a/quiche/web_transport/web_transport.h b/quiche/web_transport/web_transport.h
index 6bf68f4..7ea0085 100644
--- a/quiche/web_transport/web_transport.h
+++ b/quiche/web_transport/web_transport.h
@@ -205,8 +205,7 @@
   //
   // IMPORTANT: See the class note regarding the lifetime of the returned stream
   // object.
-  // TODO(vasilvv): implement this in a follow-up CL.
-  // virtual Stream* GetStreamById(StreamId id) = 0;
+  virtual Stream* GetStreamById(StreamId id) = 0;
 
   virtual DatagramStatus SendOrQueueDatagram(absl::string_view datagram) = 0;
   // Returns a conservative estimate of the largest datagram size that the