New QuicSocket::ConnectionCallback for first peer packet processed

Summary: as title

Reviewed By: JunqiWang

Differential Revision: D19434037

fbshipit-source-id: 7f2b6491fcce68840dc6ff57f99657dd5d535e9f

quic/api/QuicSocket.h

@@ -81,6 +81,11 @@ class QuicSocket {
      * Called when the transport is ready to send/receive data.
      */
     virtual void onTransportReady() noexcept {}
+
+    /**
+     * Called after the transport successfully processes the received packet.
+     */
+    virtual void onFirstPeerPacketProcessed() noexcept {}
   };
 
   /**

quic/api/test/Mocks.h

@@ -118,6 +118,7 @@ class MockConnectionCallback : public QuicSocket::ConnectionCallback {
       void(std::pair<QuicErrorCode, std::string>));
   GMOCK_METHOD0_(, noexcept, , onReplaySafe, void());
   GMOCK_METHOD0_(, noexcept, , onTransportReady, void());
+  GMOCK_METHOD0_(, noexcept, , onFirstPeerPacketProcessed, void());
 };
 
 class MockDeliveryCallback : public QuicSocket::DeliveryCallback {

quic/client/QuicClientTransport.cpp

@@ -665,7 +665,11 @@ void QuicClientTransport::onReadData(
     QUIC_TRACE(packet_drop, *conn_, "already_closed");
     return;
   }
+  bool waitingForFirstPacket = !hasReceivedPackets(*conn_);
   processUDPData(peer, std::move(networkData));
+  if (connCallback_ && waitingForFirstPacket && hasReceivedPackets(*conn_)) {
+    connCallback_->onFirstPeerPacketProcessed();
+  }
   if (!transportReadyNotified_ && hasWriteCipher()) {
     transportReadyNotified_ = true;
     CHECK_NOTNULL(connCallback_)->onTransportReady();

quic/client/test/QuicClientTransportTest.cpp

@@ -1700,6 +1700,53 @@ class QuicClientTransportTest : public Test {
   QuicVersion version{QuicVersion::QUIC_DRAFT};
 };
 
+TEST_F(QuicClientTransportTest, FirstPacketProcessedCallback) {
+  client->addNewPeerAddress(serverAddr);
+  client->start(&clientConnCallback);
+
+  originalConnId = client->getConn().clientConnectionId;
+  ServerConnectionIdParams params(0, 0, 0);
+  client->getNonConstConn().serverConnectionId =
+      connIdAlgo_->encodeConnectionId(params);
+
+  AckBlocks acks;
+  acks.insert(0);
+  auto& aead = getInitialCipher();
+  auto& headerCipher = getInitialHeaderCipher();
+  auto ackPacket = packetToBufCleartext(
+      createAckPacket(
+          client->getNonConstConn(),
+          initialPacketNum,
+          acks,
+          PacketNumberSpace::Initial,
+          &aead),
+      aead,
+      headerCipher,
+      initialPacketNum);
+  initialPacketNum++;
+  EXPECT_CALL(clientConnCallback, onFirstPeerPacketProcessed()).Times(1);
+  deliverData(serverAddr, ackPacket->coalesce());
+  EXPECT_FALSE(client->hasWriteCipher());
+
+  // Another ack won't trigger it again:
+  auto oneMoreAckPacket = packetToBufCleartext(
+      createAckPacket(
+          client->getNonConstConn(),
+          initialPacketNum,
+          acks,
+          PacketNumberSpace::Initial,
+          &aead),
+      aead,
+      headerCipher,
+      initialPacketNum);
+  initialPacketNum++;
+  EXPECT_CALL(clientConnCallback, onFirstPeerPacketProcessed()).Times(0);
+  deliverData(serverAddr, oneMoreAckPacket->coalesce());
+  EXPECT_FALSE(client->hasWriteCipher());
+
+  client->closeNow(folly::none);
+}
+
 TEST_F(QuicClientTransportTest, CustomTransportParam) {
   EXPECT_TRUE(client->setCustomTransportParameter(
       std::make_unique<CustomIntegralTransportParameter>(

quic/server/QuicServerTransport.cpp

@@ -101,6 +101,7 @@ void QuicServerTransport::onReadData(
   ServerEvents::ReadData readData;
   readData.peer = peer;
   readData.networkData = std::move(networkData);
+  bool waitingForFirstPacket = !hasReceivedPackets(*conn_);
   onServerReadData(*serverConn_, readData);
   processPendingData(true);
 
@@ -114,6 +115,9 @@ void QuicServerTransport::onReadData(
           shared_from_this(), *conn_->serverConnectionId);
     }
   }
+  if (connCallback_ && waitingForFirstPacket && hasReceivedPackets(*conn_)) {
+    connCallback_->onFirstPeerPacketProcessed();
+  }
   maybeWriteNewSessionTicket();
   maybeNotifyConnectionIdBound();
   maybeIssueConnectionIds();

quic/server/test/QuicServerTransportTest.cpp

@@ -3233,6 +3233,28 @@ class QuicUnencryptedServerTransportTest : public QuicServerTransportTest {
   void setupConnection() override {}
 };
 
+TEST_F(QuicUnencryptedServerTransportTest, FirstPacketProcessedCallback) {
+  getFakeHandshakeLayer()->allowZeroRttKeys();
+  EXPECT_CALL(connCallback, onFirstPeerPacketProcessed()).Times(1);
+  recvClientHello();
+  loopForWrites();
+  AckBlocks acks;
+  acks.insert(0);
+  auto aead = getInitialCipher();
+  auto headerCipher = getInitialHeaderCipher();
+  EXPECT_CALL(connCallback, onFirstPeerPacketProcessed()).Times(0);
+  deliverData(packetToBufCleartext(
+      createAckPacket(
+          server->getNonConstConn(),
+          clientNextInitialPacketNum,
+          acks,
+          PacketNumberSpace::Initial,
+          aead.get()),
+      *aead,
+      *headerCipher,
+      clientNextInitialPacketNum));
+}
+
 TEST_F(QuicUnencryptedServerTransportTest, TestUnencryptedStream) {
   auto data = IOBuf::copyBuffer("bad data");
   PacketNum nextPacket = clientNextInitialPacketNum++;