MOQT Session-level handling of draft-18 GOAWAY.

Request-level GOAWAY will be processed and silently dropped on all bidi streams until further restrictions are clarified in https://github.com/moq-wg/moq-transport/issues/1655.

PiperOrigin-RevId: 991292836
diff --git a/quiche/quic/moqt/moqt_bidi_stream.cc b/quiche/quic/moqt/moqt_bidi_stream.cc
index 6d748a0..15167c8 100644
--- a/quiche/quic/moqt/moqt_bidi_stream.cc
+++ b/quiche/quic/moqt/moqt_bidi_stream.cc
@@ -91,6 +91,25 @@
                              /*fin=*/false);
 }
 
+// TODO(martinduke): There should be an application callback when these are
+// received, but https://github.com/moq-wg/moq-transport/issues/1655 is
+// resolving which child classes of MoqtBidiStreamBase should handle them. For
+// now, we just handle them silently, as in draft-18 they are all allowed, to
+// avoid changing the application API twice. However, it will return an error if
+// the message is incorrectly formatted or there is more than one.
+absl::Status MoqtBidiStreamBase::OnControlMessage(const MoqtGoAway& message) {
+  if (received_goaway_) {
+    return absl::InvalidArgumentError(
+        "Received multiple GOAWAY on a bidi stream");
+  }
+  if (message.request_id.has_value()) {
+    return absl::InvalidArgumentError(
+        "GOAWAY with request ID on a bidi stream");
+  }
+  received_goaway_ = true;
+  return absl::OkStatus();
+}
+
 void MoqtBidiStreamBase::OnFatalError(absl::Status status) {
   QUICHE_DCHECK(!status.ok());
   if (session_error_callback_ == nullptr) {
diff --git a/quiche/quic/moqt/moqt_bidi_stream.h b/quiche/quic/moqt/moqt_bidi_stream.h
index 2c154c6..fef0eeb 100644
--- a/quiche/quic/moqt/moqt_bidi_stream.h
+++ b/quiche/quic/moqt/moqt_bidi_stream.h
@@ -146,6 +146,8 @@
     return stream() != nullptr ? stream()->GetStreamId() : 0;
   }
 
+  absl::Status OnControlMessage(const MoqtGoAway& message);
+
  protected:
   // Called when a WebTransport stream has been associated with the object.
   // Should be used to set the priority for the stream.
@@ -184,6 +186,7 @@
   MoqtRequestUpdateQueue request_update_queue_;
   SessionErrorCallback session_error_callback_;
   absl::Status stream_status_ = absl::OkStatus();
+  bool received_goaway_ = false;
 };
 
 // DispatchControlMessage is wrapped into a class so that the caller class can
diff --git a/quiche/quic/moqt/moqt_bidi_stream_test.cc b/quiche/quic/moqt/moqt_bidi_stream_test.cc
index bfc110e..bc3feae 100644
--- a/quiche/quic/moqt/moqt_bidi_stream_test.cc
+++ b/quiche/quic/moqt/moqt_bidi_stream_test.cc
@@ -11,6 +11,7 @@
 
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "quiche/quic/core/quic_time.h"
 #include "quiche/quic/core/quic_types.h"
 #include "quiche/quic/moqt/moqt_error.h"
 #include "quiche/quic/moqt/moqt_fetch_task.h"
@@ -31,6 +32,7 @@
 class TestMoqtBidiStream : public MoqtBidiStreamBase {
  public:
   using MoqtBidiStreamBase::MoqtBidiStreamBase;
+  using MoqtBidiStreamBase::OnControlMessage;
 
   void OnStreamBound() override {}
   absl::Status OnRawControlMessage(
@@ -121,17 +123,29 @@
   stream.Receive(framer.SerializeObjectAck(MoqtObjectAck()).AsStringView());
   stream_->OnCanRead();
 
-  stream.Receive(framer.SerializeGoAway(MoqtGoAway()).AsStringView());
+  stream.Receive(framer.SerializePublishDone(MoqtPublishDone()).AsStringView());
   EXPECT_CALL(error_callback_, Call)
       .WillOnce([](MoqtError error, absl::string_view message) {
         EXPECT_EQ(error, MoqtError::kProtocolViolation);
-        EXPECT_EQ(
-            message,
-            "Received an unexpected message of type GOAWAY on a test stream");
+        EXPECT_EQ(message,
+                  "Received an unexpected message of type PUBLISH_DONE on a "
+                  "test stream");
       });
   stream_->OnCanRead();
 }
 
+TEST_F(MoqtBidiStreamTest, ReceiveGoAway) {
+  QUICHE_EXPECT_OK(stream_->OnControlMessage(
+      MoqtGoAway("", quic::QuicTimeDelta::Zero(), std::nullopt)));
+  EXPECT_TRUE(IsInvalidArgument(stream_->OnControlMessage(
+      MoqtGoAway("", quic::QuicTimeDelta::Zero(), std::nullopt))));
+}
+
+TEST_F(MoqtBidiStreamTest, ReceiveGoAwayWithRequestId) {
+  EXPECT_TRUE(IsInvalidArgument(stream_->OnControlMessage(
+      MoqtGoAway("", quic::QuicTimeDelta::Zero(), /*request_id=*/0))));
+}
+
 TEST_F(MoqtBidiStreamTest, SendRequestOk) {
   stream_->BindStream(&mock_stream_);
   EXPECT_CALL(mock_stream_, CanWrite).WillRepeatedly(testing::Return(true));
diff --git a/quiche/quic/moqt/moqt_fetch_stream.h b/quiche/quic/moqt/moqt_fetch_stream.h
index 6f54cea..57d4604 100644
--- a/quiche/quic/moqt/moqt_fetch_stream.h
+++ b/quiche/quic/moqt/moqt_fetch_stream.h
@@ -76,6 +76,7 @@
                       std::optional<DataStreamIndex> index) override;
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -118,6 +119,7 @@
   }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override { stream_parser()->set_allow_fin(true); }
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
diff --git a/quiche/quic/moqt/moqt_namespace_stream.h b/quiche/quic/moqt/moqt_namespace_stream.h
index 5935220..49c9178 100644
--- a/quiche/quic/moqt/moqt_namespace_stream.h
+++ b/quiche/quic/moqt/moqt_namespace_stream.h
@@ -51,6 +51,7 @@
   ~MoqtSubscribeNamespaceRequestStream();
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -153,6 +154,7 @@
       MoqtIncomingSubscribeNamespaceCallback& application);
   ~MoqtSubscribeNamespaceResponseStream() { Detach(); }
 
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override {
     // TODO(martinduke): Set the priority for this stream.
   }
diff --git a/quiche/quic/moqt/moqt_publish_namespace_stream.h b/quiche/quic/moqt/moqt_publish_namespace_stream.h
index 47b6a15..5b497f1 100644
--- a/quiche/quic/moqt/moqt_publish_namespace_stream.h
+++ b/quiche/quic/moqt/moqt_publish_namespace_stream.h
@@ -48,6 +48,7 @@
   ~MoqtPublishNamespaceRequestStream() { Detach(); }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -83,6 +84,7 @@
         weak_ptr_factory_(this) {}
   ~MoqtPublishNamespaceResponseStream() { Detach(); }
 
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override {
     // TODO(martinduke): Set the priority for this stream.
   }
diff --git a/quiche/quic/moqt/moqt_publish_stream.h b/quiche/quic/moqt/moqt_publish_stream.h
index 4beb0c7..c2c1205 100644
--- a/quiche/quic/moqt/moqt_publish_stream.h
+++ b/quiche/quic/moqt/moqt_publish_stream.h
@@ -43,6 +43,7 @@
   ~MoqtPublishRequestStream();
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -93,6 +94,7 @@
   ~MoqtPublishResponseStream() { Detach(); }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override {
     stream_parser()->set_allow_fin(true);
     // TODO(martinduke): Set the priority for this stream.
diff --git a/quiche/quic/moqt/moqt_session.cc b/quiche/quic/moqt/moqt_session.cc
index 4fa267c..30006fe 100644
--- a/quiche/quic/moqt/moqt_session.cc
+++ b/quiche/quic/moqt/moqt_session.cc
@@ -99,10 +99,13 @@
   if (parameters_.using_webtrans) {
     session_->SetOnDraining([this]() {
       QUICHE_DLOG(INFO) << "WebTransport session is draining";
-      received_goaway_ = true;
-      if (callbacks_.goaway_received_callback != nullptr) {
-        std::move(callbacks_.goaway_received_callback)(absl::string_view());
+      if (callbacks_.goaway_received_callback == nullptr) {
+        return;
       }
+      MoqtSessionGoAwayCallback callback =
+          std::move(callbacks_.goaway_received_callback);
+      callbacks_.goaway_received_callback = nullptr;
+      std::move(callback)(absl::string_view(), quic::QuicTimeDelta::Zero());
     });
   }
   if (parameters_.perspective == Perspective::IS_SERVER) {
@@ -780,13 +783,16 @@
     QUIC_DLOG(INFO) << ENDPOINT << "Tried to send multiple GOAWAY";
     return;
   }
-  if (!new_session_uri.empty() && !new_session_uri.empty()) {
+  if (!new_session_uri.empty() &&
+      parameters_.perspective == Perspective::IS_CLIENT) {
     QUIC_DLOG(INFO) << ENDPOINT
                     << "Client tried to send GOAWAY with new session URI";
     return;
   }
   MoqtGoAway message;
   message.new_session_uri = std::string(new_session_uri);
+  message.timeout = kDefaultGoAwayTimeout;
+  message.request_id = next_incoming_request_id_;
   SendControlMessage(framer_.SerializeGoAway(message));
   sent_goaway_ = true;
   goaway_timeout_alarm_ = absl::WrapUnique(
@@ -1443,22 +1449,30 @@
 }
 
 absl::Status MoqtSession::OnControlMessage(const MoqtGoAway& message) {
-  if (!message.new_session_uri.empty() &&
-      perspective() == quic::Perspective::IS_SERVER) {
-    return absl::InvalidArgumentError(
-        "Received GOAWAY with new_session_uri on the server");
-  }
   if (received_goaway_) {
-    return absl::InvalidArgumentError("Received multiple GOAWAY messages");
+    return absl::InvalidArgumentError(
+        "Received multiple GOAWAY on control stream");
+  }
+  if (!message.request_id.has_value()) {
+    return absl::InvalidArgumentError("GOAWAY missing request ID");
+  }
+  if ((*message.request_id & 0x1) != (next_request_id_ & 0x1)) {
+    Error(MoqtError::kInvalidRequestId,
+          "GOAWAY request ID has incorrect parity");
+    return absl::OkStatus();
   }
   received_goaway_ = true;
-  if (callbacks_.goaway_received_callback != nullptr) {
-    std::move(callbacks_.goaway_received_callback)(message.new_session_uri);
+  if (callbacks_.goaway_received_callback == nullptr) {
+    // An HTTP/3 or WebTransport GOAWAY already invoked the callback.
+    return absl::OkStatus();
   }
+  MoqtSessionGoAwayCallback callback =
+      std::move(callbacks_.goaway_received_callback);
+  callbacks_.goaway_received_callback = nullptr;
+  std::move(callback)(message.new_session_uri, message.timeout);
   return absl::OkStatus();
 }
 
-
 void MoqtSession::OnMalformedTrack(ObjectSubscriber* track) {
   if (track->is_fetch()) {
     QUICHE_BUG(quiche_bug_malformed_fetch_track)
diff --git a/quiche/quic/moqt/moqt_session.h b/quiche/quic/moqt/moqt_session.h
index 7aba415..330b6f7 100644
--- a/quiche/quic/moqt/moqt_session.h
+++ b/quiche/quic/moqt/moqt_session.h
@@ -250,8 +250,8 @@
     ~UnknownBidiStream() {}
 
     // webtransport::StreamVisitor overrides.
-    void OnResetStreamReceived(webtransport::StreamErrorCode error) override {}
-    void OnStopSendingReceived(webtransport::StreamErrorCode error) override {}
+    void OnResetStreamReceived(webtransport::StreamErrorCode) override {}
+    void OnStopSendingReceived(webtransport::StreamErrorCode) override {}
     void OnWriteSideInDataRecvdState() override {}
     void OnCanRead() override;
     void OnCanWrite() override {}
@@ -272,8 +272,8 @@
         : session_(session->GetWeakPtr()), stream_(stream), parser_(stream) {}
 
     // webtransport::StreamVisitor overrides.
-    void OnResetStreamReceived(webtransport::StreamErrorCode error) override {}
-    void OnStopSendingReceived(webtransport::StreamErrorCode error) override {}
+    void OnResetStreamReceived(webtransport::StreamErrorCode) override {}
+    void OnStopSendingReceived(webtransport::StreamErrorCode) override {}
     void OnWriteSideInDataRecvdState() override {}
     void OnCanRead() override;
     void OnCanWrite() override {}
@@ -293,7 +293,7 @@
     void OnCanRead() override;
     void OnCanWrite() override {}
     void OnResetStreamReceived(webtransport::StreamErrorCode error) override;
-    void OnStopSendingReceived(webtransport::StreamErrorCode error) override {
+    void OnStopSendingReceived(webtransport::StreamErrorCode) override {
       // Impossible for QUIC incoming unidirectional streams.
     }
     void OnWriteSideInDataRecvdState() override {}
@@ -320,7 +320,7 @@
 
     void OnCanRead() override {}
     void OnCanWrite() override;
-    void OnResetStreamReceived(webtransport::StreamErrorCode error) override {
+    void OnResetStreamReceived(webtransport::StreamErrorCode) override {
       // Impossible for QUIC incoming unidirectional streams.
     }
     void OnStopSendingReceived(webtransport::StreamErrorCode error) override;
@@ -409,10 +409,7 @@
 
   // Handlers for the control messages on the main control stream.
   absl::Status OnControlMessage(const MoqtSetup& message);
-
-  // TODO(martinduke): All of these should be moved to bidi streams or
-  // deleted.
-  absl::Status OnControlMessage(const MoqtGoAway& /*message*/);
+  absl::Status OnControlMessage(const MoqtGoAway& message);
 
   uint64_t NextRequestId() {
     uint64_t id = next_request_id_;
diff --git a/quiche/quic/moqt/moqt_session_callbacks.h b/quiche/quic/moqt/moqt_session_callbacks.h
index e222136..71217c0 100644
--- a/quiche/quic/moqt/moqt_session_callbacks.h
+++ b/quiche/quic/moqt/moqt_session_callbacks.h
@@ -89,8 +89,8 @@
 using MoqtSessionEstablishedCallback = quiche::SingleUseCallback<void()>;
 
 // Called when a GOAWAY message is received from the server.
-using MoqtSessionGoAwayCallback =
-    quiche::SingleUseCallback<void(absl::string_view new_session_uri)>;
+using MoqtSessionGoAwayCallback = quiche::SingleUseCallback<void(
+    absl::string_view new_session_uri, quic::QuicTimeDelta timeout)>;
 
 // Called when the session is terminated.
 using MoqtSessionTerminatedCallback =
@@ -171,7 +171,7 @@
 struct MoqtSessionCallbacks {
   MoqtSessionEstablishedCallback session_established_callback = +[] {};
   MoqtSessionGoAwayCallback goaway_received_callback =
-      +[](absl::string_view) {};
+      +[](absl::string_view, quic::QuicTimeDelta) {};
   MoqtSessionTerminatedCallback session_terminated_callback =
       +[](absl::string_view) {};
   MoqtSessionDeletedCallback session_deleted_callback = +[] {};
diff --git a/quiche/quic/moqt/moqt_session_test.cc b/quiche/quic/moqt/moqt_session_test.cc
index 73a9b5d..bbed4a3 100644
--- a/quiche/quic/moqt/moqt_session_test.cc
+++ b/quiche/quic/moqt/moqt_session_test.cc
@@ -2309,9 +2309,10 @@
 TEST_F(MoqtSessionTest, ReceiveGoAwayEnforcement) {
   bidi_wrapper_ =
       MoqtSessionPeer::CreateControlStream(&session_, &mock_bidi_stream_);
-  EXPECT_CALL(session_callbacks_.goaway_received_callback, Call(""));
-  bidi_wrapper_->ReceiveMessage(MoqtGoAway("", quic::QuicTimeDelta::Zero(),
-                                           /*request_id=*/0));
+  EXPECT_CALL(session_callbacks_.goaway_received_callback,
+              Call("", quic::QuicTimeDelta::Zero()));
+  bidi_wrapper_->ReceiveMessage(
+      MoqtGoAway("", quic::QuicTimeDelta::Zero(), /*request_id=*/0));
   // New requests not allowed.
   EXPECT_CALL(mock_bidi_stream_, Writev).Times(0);
   MessageParameters parameters = SubscribeForTest();
@@ -2334,15 +2335,51 @@
   // Error on additional GOAWAY.
   EXPECT_CALL(mock_session_,
               CloseSession(static_cast<uint64_t>(MoqtError::kProtocolViolation),
-                           "Received multiple GOAWAY messages"))
+                           "Received multiple GOAWAY on control stream"))
       .Times(1);
   bool reported_error = false;
   EXPECT_CALL(session_callbacks_.session_terminated_callback, Call)
       .WillOnce([&](absl::string_view error_message) {
         reported_error = true;
-        EXPECT_EQ(error_message, "Received multiple GOAWAY messages");
+        EXPECT_EQ(error_message, "Received multiple GOAWAY on control stream");
       });
-  bidi_wrapper_->ReceiveMessage(MoqtGoAway(""));
+  bidi_wrapper_->ReceiveMessage(
+      MoqtGoAway{"", quic::QuicTimeDelta::Zero(), /*request_id=*/0});
+}
+
+TEST_F(MoqtSessionTest, GoAwayMissingRequestId) {
+  bidi_wrapper_ =
+      MoqtSessionPeer::CreateControlStream(&session_, &mock_bidi_stream_);
+  EXPECT_CALL(mock_session_,
+              CloseSession(static_cast<uint64_t>(MoqtError::kProtocolViolation),
+                           "GOAWAY missing request ID"));
+  bidi_wrapper_->ReceiveMessage(
+      MoqtGoAway{"", quic::QuicTimeDelta::Zero(), std::nullopt});
+}
+
+TEST_F(MoqtSessionTest, GoAwayRequestIdWrongParity) {
+  // Client session has even next_request_id_, so odd request_id fails.
+  bidi_wrapper_ =
+      MoqtSessionPeer::CreateControlStream(&session_, &mock_bidi_stream_);
+  EXPECT_CALL(mock_session_,
+              CloseSession(static_cast<uint64_t>(MoqtError::kInvalidRequestId),
+                           "GOAWAY request ID has incorrect parity"));
+  bidi_wrapper_->ReceiveMessage(
+      MoqtGoAway{"", quic::QuicTimeDelta::Zero(), /*request_id=*/1});
+
+  // Server session has odd next_request_id_, so even request_id fails.
+  webtransport::test::MockSession mock_server_session;
+  MoqtSession server_session(
+      &mock_server_session, MoqtSessionParameters(quic::Perspective::IS_SERVER),
+      std::make_unique<quic::test::TestAlarmFactory>(),
+      session_callbacks_.AsSessionCallbacks());
+  auto server_control_stream =
+      MoqtSessionPeer::CreateControlStream(&server_session, &mock_bidi_stream_);
+  EXPECT_CALL(mock_server_session,
+              CloseSession(static_cast<uint64_t>(MoqtError::kInvalidRequestId),
+                           "GOAWAY request ID has incorrect parity"));
+  server_control_stream->ReceiveMessage(
+      MoqtGoAway{"", quic::QuicTimeDelta::Zero(), /*request_id=*/0});
 }
 
 TEST_F(MoqtSessionTest, SendGoAwayEnforcement) {
@@ -2816,6 +2853,7 @@
       MoqtSessionPeer::CreateControlStream(&session_, &mock_bidi_stream_);
   MoqtGoAway goaway;
   goaway.new_session_uri = "";
+  goaway.request_id = 0;
   bidi_wrapper_->ReceiveMessage(goaway);
   CreateTrackPublisher();
   std::shared_ptr<MoqtTrackPublisher> track_publisher =
diff --git a/quiche/quic/moqt/moqt_subscribe_stream.h b/quiche/quic/moqt/moqt_subscribe_stream.h
index a2406ae..cd96872 100644
--- a/quiche/quic/moqt/moqt_subscribe_stream.h
+++ b/quiche/quic/moqt/moqt_subscribe_stream.h
@@ -40,6 +40,7 @@
   ~MoqtSubscribeRequestStream() { Detach(); }
 
   // StreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -77,6 +78,7 @@
   }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override { stream_parser()->set_allow_fin(true); }
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
diff --git a/quiche/quic/moqt/moqt_track_status_stream.h b/quiche/quic/moqt/moqt_track_status_stream.h
index 4cc8e67..d9c5c6e 100644
--- a/quiche/quic/moqt/moqt_track_status_stream.h
+++ b/quiche/quic/moqt/moqt_track_status_stream.h
@@ -41,6 +41,7 @@
   ~MoqtTrackStatusRequestStream() { Detach(); }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override;
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
@@ -69,6 +70,7 @@
   ~MoqtTrackStatusResponseStream() { Detach(); }
 
   // MoqtBidiStreamBase overrides.
+  using MoqtBidiStreamBase::OnControlMessage;
   void OnStreamBound() override { stream_parser()->set_allow_fin(true); }
   absl::Status OnRawControlMessage(
       const MoqtRawControlMessage& message) override;
diff --git a/quiche/quic/moqt/test_tools/moqt_framer_utils.cc b/quiche/quic/moqt/test_tools/moqt_framer_utils.cc
index 287553f..1b11bfd 100644
--- a/quiche/quic/moqt/test_tools/moqt_framer_utils.cc
+++ b/quiche/quic/moqt/test_tools/moqt_framer_utils.cc
@@ -91,6 +91,8 @@
     if (!use_webtrans && !setup.options.path.has_value()) {
       perspective = quic::Perspective::IS_SERVER;
     }
+  } else if (std::holds_alternative<MoqtGoAway>(frame)) {
+    perspective = quic::Perspective::IS_SERVER;
   }
   MoqtFramer framer(use_webtrans, perspective);
   return std::string(std::visit(FramingVisitor{framer}, frame).AsStringView());
diff --git a/quiche/quic/moqt/test_tools/moqt_mock_visitor.h b/quiche/quic/moqt/test_tools/moqt_mock_visitor.h
index 7d5154f..a24483f 100644
--- a/quiche/quic/moqt/test_tools/moqt_mock_visitor.h
+++ b/quiche/quic/moqt/test_tools/moqt_mock_visitor.h
@@ -40,7 +40,8 @@
 
 struct MockSessionCallbacks {
   testing::MockFunction<void()> session_established_callback;
-  testing::MockFunction<void(absl::string_view)> goaway_received_callback;
+  testing::MockFunction<void(absl::string_view, quic::QuicTimeDelta)>
+      goaway_received_callback;
   testing::MockFunction<void(absl::string_view)> session_terminated_callback;
   testing::MockFunction<void()> session_deleted_callback;
   testing::MockFunction<void(const TrackNamespace&,
diff --git a/quiche/quic/moqt/tools/chat_client.cc b/quiche/quic/moqt/tools/chat_client.cc
index 969f1a6..a208c24 100644
--- a/quiche/quic/moqt/tools/chat_client.cc
+++ b/quiche/quic/moqt/tools/chat_client.cc
@@ -26,6 +26,7 @@
 #include "quiche/quic/core/io/quic_event_loop.h"
 #include "quiche/quic/core/quic_default_clock.h"
 #include "quiche/quic/core/quic_server_id.h"
+#include "quiche/quic/core/quic_time.h"
 #include "quiche/quic/moqt/moqt_error.h"
 #include "quiche/quic/moqt/moqt_fetch_task.h"
 #include "quiche/quic/moqt/moqt_key_value_pair.h"
@@ -35,8 +36,6 @@
 #include "quiche/quic/moqt/moqt_outgoing_queue.h"
 #include "quiche/quic/moqt/moqt_session.h"
 #include "quiche/quic/moqt/moqt_session_callbacks.h"
-#include "quiche/quic/moqt/moqt_session_interface.h"
-#include "quiche/quic/moqt/moqt_types.h"
 #include "quiche/quic/moqt/tools/moq_chat.h"
 #include "quiche/quic/moqt/tools/moqt_client.h"
 #include "quiche/quic/platform/api/quic_default_proof_providers.h"
@@ -146,9 +145,9 @@
     session_is_open_ = true;
   };
   session_callbacks_.goaway_received_callback =
-      [](absl::string_view new_session_uri) {
+      [](absl::string_view new_session_uri, quic::QuicTimeDelta timeout) {
         std::cout << "GoAway received, new session uri = " << new_session_uri
-                  << "\n";
+                  << ", timeout = " << timeout.ToMilliseconds() << "ms\n";
         // TODO (martinduke): Connect to the new session uri.
       };
   session_callbacks_.session_terminated_callback =
diff --git a/quiche/quic/moqt/tools/moqt_relay.cc b/quiche/quic/moqt/tools/moqt_relay.cc
index ebfa5d0..222241e 100644
--- a/quiche/quic/moqt/tools/moqt_relay.cc
+++ b/quiche/quic/moqt/tools/moqt_relay.cc
@@ -17,6 +17,7 @@
 #include "quiche/quic/core/crypto/proof_verifier.h"
 #include "quiche/quic/core/io/quic_event_loop.h"
 #include "quiche/quic/core/quic_server_id.h"
+#include "quiche/quic/core/quic_time.h"
 #include "quiche/quic/moqt/moqt_error.h"
 #include "quiche/quic/moqt/moqt_fetch_task.h"
 #include "quiche/quic/moqt/moqt_key_value_pair.h"
@@ -103,9 +104,10 @@
     publisher_.SetDefaultUpstreamSession(session);
     SetNamespaceCallbacks(session);
   };
-  callbacks.goaway_received_callback = [](absl::string_view new_session_uri) {
-    QUICHE_LOG(INFO) << "GoAway received, new session uri = "
-                     << new_session_uri;
+  callbacks.goaway_received_callback = [](absl::string_view new_session_uri,
+                                          quic::QuicTimeDelta timeout) {
+    QUICHE_LOG(INFO) << "GoAway received, new session uri = " << new_session_uri
+                     << ", timeout = " << timeout;
     // There's no asynchronous means today to connect to a new URL.
     // Therefore, just ignore GOAWAY.
   };