From c34e6eb1e0e9915c765ad78eb37ad6ea8a873c23 Mon Sep 17 00:00:00 2001
From: Konstantin Tsoy <ktsoy@meta.com>
Date: Thu, 23 May 2024 11:09:34 -0700
Subject: [PATCH] Check that transport is still open in
 maybeStopWriteLooperAndArmSocketWritableEvent();

Summary: We can throw in `writeSocketDataAndCatch()` and that would result in transport closure with drain (with drain keeping the socket alive), but the congestion controller would be reset in close: https://fburl.com/code/zxwfye5u and `maybeStopWriteLooperAndArmSocketWritableEvent()` trips over it later when executed in scope exit.

Reviewed By: mjoras

Differential Revision: D57728369

fbshipit-source-id: 51a4719ae97fab0e90e3ae093a3f56be5a096aff
---
 quic/api/QuicTransportBase.cpp          | 11 +++++--
 quic/api/test/QuicTransportBaseTest.cpp | 42 +++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp
index ad28dc698..1b0473aa0 100644
--- a/quic/api/QuicTransportBase.cpp
+++ b/quic/api/QuicTransportBase.cpp
@@ -3584,7 +3584,7 @@ void QuicTransportBase::onSocketWritable() noexcept {
 }
 
 void QuicTransportBase::maybeStopWriteLooperAndArmSocketWritableEvent() {
-  if (!socket_) {
+  if (!socket_ || (closeState_ == CloseState::CLOSED)) {
     return;
   }
   if (conn_->transportSettings.useSockWritableEvents &&
@@ -3595,9 +3595,14 @@ void QuicTransportBase::maybeStopWriteLooperAndArmSocketWritableEvent() {
     bool haveBufferToRetry = writeReason == WriteDataReason::BUFFERED_WRITE;
     bool haveNewDataToWrite =
         (writeReason != WriteDataReason::NO_WRITE) && !haveBufferToRetry;
+    bool haveCongestionControlWindow = true;
+    if (conn_->congestionController) {
+      haveCongestionControlWindow =
+          conn_->congestionController->getWritableBytes() > 0;
+    }
+    bool haveFlowControlWindow = getSendConnFlowControlBytesAPI(*conn_) > 0;
     bool connHasWriteWindow =
-        (conn_->congestionController->getWritableBytes() > 0) &&
-        (getSendConnFlowControlBytesAPI(*conn_) > 0);
+        haveCongestionControlWindow && haveFlowControlWindow;
     if (haveBufferToRetry || (haveNewDataToWrite && connHasWriteWindow)) {
       // Re-arm the write event and stop the write
       // looper.
diff --git a/quic/api/test/QuicTransportBaseTest.cpp b/quic/api/test/QuicTransportBaseTest.cpp
index 12e4db795..bb88263a3 100644
--- a/quic/api/test/QuicTransportBaseTest.cpp
+++ b/quic/api/test/QuicTransportBaseTest.cpp
@@ -574,6 +574,14 @@ class TestQuicTransport
     QuicTransportBase::maybeStopWriteLooperAndArmSocketWritableEvent();
   }
 
+  void closeImpl(
+      folly::Optional<QuicError> error,
+      bool drainConnection = true,
+      bool sendCloseImmediately = true) {
+    QuicTransportBase::closeImpl(
+        std::move(error), drainConnection, sendCloseImmediately);
+  }
+
   void onSocketWritable() noexcept override {
     QuicTransportBase::onSocketWritable();
   }
@@ -4906,5 +4914,39 @@ TEST_P(
   transport.reset();
 }
 
+TEST_P(
+    QuicTransportImplTestBase,
+    TestMaybeStopWriteLooperAndArmSocketWritableEventOnClosedSocket) {
+  auto transportSettings = transport->getTransportSettings();
+  transportSettings.useSockWritableEvents = true;
+  transport->setTransportSettings(transportSettings);
+  transport->getConnectionState().streamManager->refreshTransportSettings(
+      transportSettings);
+
+  transport->transportConn->oneRttWriteCipher = test::createNoOpAead();
+
+  // Create a stream with outgoing data.
+  auto streamId = transport->createBidirectionalStream().value();
+  const auto& conn = transport->transportConn;
+  auto stream = transport->getStream(streamId);
+  stream->writeBuffer.append(IOBuf::copyBuffer("hello"));
+
+  // Insert streamId into the list.
+  conn->streamManager->addWritable(*stream);
+  conn->streamManager->updateWritableStreams(*stream);
+
+  // Write looper is running.
+  transport->writeLooper()->run(true /* thisIteration */);
+  EXPECT_TRUE(transport->writeLooper()->isRunning());
+
+  // Close the socket.
+  transport->closeImpl((QuicError(
+      QuicErrorCode(TransportErrorCode::INTERNAL_ERROR),
+      std::string("writeSocketDataAndCatch()  error"))));
+  transport->maybeStopWriteLooperAndArmSocketWritableEvent();
+
+  transport.reset();
+}
+
 } // namespace test
 } // namespace quic