Defer stream available callbacks until connCallback_ is set

Summary: If the peer sends higher stream limits in the settings, these callbacks may be invoked before the app has a chance to set the connection callback.

Reviewed By: sharmafb

Differential Revision: D64214648

fbshipit-source-id: 6a8a9b8d4d9e02a2baad672d69a43ba61daba918

quic/api/QuicTransportBaseLite.cpp

@@ -646,6 +646,14 @@ void QuicTransportBaseLite::setConnectionSetupCallback(
 void QuicTransportBaseLite::setConnectionCallback(
     folly::MaybeManagedPtr<ConnectionCallback> callback) {
   connCallback_ = callback;
+  if (connCallback_) {
+    runOnEvbAsync([](auto self) { self->processCallbacksAfterNetworkData(); });
+  }
+}
+
+void QuicTransportBaseLite::setConnectionCallbackFromCtor(
+    folly::MaybeManagedPtr<ConnectionCallback> callback) {
+  connCallback_ = callback;
 }
 
 Optional<LocalErrorCode> QuicTransportBaseLite::setControlStream(StreamId id) {
@@ -1501,6 +1509,10 @@ void QuicTransportBaseLite::processCallbacksAfterNetworkData() {
   if (closeState_ != CloseState::OPEN) {
     return;
   }
+  if (!connCallback_ || !conn_->streamManager) {
+    return;
+  }
+
   // We reuse this storage for storing streams which need callbacks.
   std::vector<StreamId> tempStorage;
 

quic/api/QuicTransportBaseLite.h

@@ -471,6 +471,9 @@ class QuicTransportBaseLite : virtual public QuicSocketLite,
   virtual void createBufAccessor(size_t /* capacity */) {}
 
  protected:
+  void setConnectionCallbackFromCtor(
+      folly::MaybeManagedPtr<ConnectionCallback> callback);
+
   /**
    * A wrapper around writeSocketData
    *

quic/api/test/QuicTransportBaseTest.cpp

@@ -235,8 +235,6 @@ class TestQuicTransport
       ConnectionCallback* connCb)
       : QuicTransportBase(std::move(evb), std::move(socket)),
         observerContainer_(std::make_shared<SocketObserverContainer>(this)) {
-    setConnectionSetupCallback(connSetupCb);
-    setConnectionCallback(connCb);
     auto conn = std::make_unique<QuicServerConnectionState>(
         FizzServerQuicHandshakeContext::Builder().build());
     conn->clientConnectionId = ConnectionId({10, 9, 8, 7});
@@ -247,6 +245,8 @@ class TestQuicTransport
     aead = test::createNoOpAead();
     headerCipher = test::createNoOpHeaderCipher();
     connIdAlgo_ = std::make_unique<DefaultConnectionIdAlgo>();
+    setConnectionSetupCallback(connSetupCb);
+    setConnectionCallbackFromCtor(connCb);
   }
 
   ~TestQuicTransport() override {
@@ -711,6 +711,23 @@ TEST_P(QuicTransportImplTestBase, IdleTimeoutExpiredDestroysTransport) {
   transport->invokeIdleTimeout();
 }
 
+TEST_P(QuicTransportImplTestBase, DelayConnCallback) {
+  transport->transportConn->streamManager->setMaxLocalBidirectionalStreams(
+      0, /*force=*/true);
+  transport->setConnectionCallback(nullptr);
+
+  transport->addMaxStreamsFrame(
+      MaxStreamsFrame(10, /*isBidirectionalIn=*/true));
+
+  transport->setConnectionCallback(&connCallback);
+  EXPECT_CALL(connCallback, onBidirectionalStreamsAvailable(_))
+      .WillOnce(Invoke([](uint64_t numAvailableStreams) {
+        EXPECT_EQ(numAvailableStreams, 10);
+      }));
+  transport->getEventBase()->loopOnce();
+  transport.reset();
+}
+
 TEST_P(QuicTransportImplTestBase, IdleTimeoutStreamMaessage) {
   auto stream1 = transport->createBidirectionalStream().value();
   auto stream2 = transport->createBidirectionalStream().value();

quic/api/test/TestQuicTransport.h

@@ -26,8 +26,6 @@ class TestQuicTransport
       ConnectionCallback* connCb)
       : QuicTransportBase(std::move(evb), std::move(socket)),
         observerContainer_(std::make_shared<SocketObserverContainer>(this)) {
-    setConnectionSetupCallback(connSetupCb);
-    setConnectionCallback(connCb);
     conn_.reset(new QuicServerConnectionState(
         FizzServerQuicHandshakeContext::Builder().build()));
     conn_->clientConnectionId = ConnectionId({9, 8, 7, 6});
@@ -36,6 +34,8 @@ class TestQuicTransport
     conn_->observerContainer = observerContainer_;
     aead = test::createNoOpAead();
     headerCipher = test::createNoOpHeaderCipher();
+    setConnectionSetupCallback(connSetupCb);
+    setConnectionCallbackFromCtor(connCb);
   }
 
   ~TestQuicTransport() override {

quic/server/QuicServerTransport.cpp

@@ -67,7 +67,7 @@ QuicServerTransport::QuicServerTransport(
   conn_.reset(tempConn.release());
   conn_->observerContainer = wrappedObserverContainer_.getWeakPtr();
   setConnectionSetupCallback(connSetupCb);
-  setConnectionCallback(connStreamsCb);
+  setConnectionCallbackFromCtor(connStreamsCb);
   registerAllTransportKnobParamHandlers();
 }