Key update support: Add support for initiating periodic key updates for both client and server [4/x]

Summary:
Allow the server/client transport to initiate periodic key update. It's defaulted to being disabled.

The new logic for initiating and verifying a key update was handled correctly by the peer is consolidated in QuicTransportFunctions.

Reviewed By: mjoras

Differential Revision: D53109624

fbshipit-source-id: 0c3a944978fc0e0a84252da953dc116aa7c26379

quic/QuicConstants.h

@@ -613,6 +613,13 @@ constexpr uint32_t kDefaultMaxDatagramsBuffered = 75;
 // milliseconds
 constexpr std::chrono::milliseconds kMinIntervalBetweenSessionTickets = 100ms;
 
+// Number of packets to write with the current cipher before initiating a key
+// update. This is a conservative number below the confidentiality limit (2^23)
+// derived in the spec for packets of size up to 64k bytes:
+// https://www.rfc-editor.org/rfc/rfc9001.html#name-confidentiality-limit
+constexpr uint64_t kDefaultKeyUpdatePacketCountInterval =
+    static_cast<uint64_t>(8 * 1000 * 1000);
+
 enum class ZeroRttSourceTokenMatchingPolicy : uint8_t {
   REJECT_IF_NO_EXACT_MATCH = 0,
   LIMIT_IF_NO_EXACT_MATCH = 1,

quic/api/QuicTransportFunctions.cpp

@@ -843,7 +843,20 @@ void updateConnection(
   conn.lossState.totalPacketsSent++;
   conn.lossState.totalStreamBytesSent += streamBytesSent;
   conn.lossState.totalNewStreamBytesSent += newStreamBytesSent;
+
+  // Count the number of packets sent in the current phase.
+  // This is used to initiate key updates if enabled.
+  if (packet.header.getProtectionType() == conn.oneRttWritePhase) {
+    if (conn.oneRttWritePendingVerification &&
+        conn.oneRttWritePacketsSentInCurrentPhase == 0) {
+      // This is the first packet in the new phase after we have initiated a key
+      // update. We need to keep track of it to confirm the peer acks it in the
+      // same phase.
+      conn.oneRttWritePendingVerificationPacketNumber =
+          packet.header.getPacketSequenceNum();
+    }
     conn.oneRttWritePacketsSentInCurrentPhase++;
+  }
 
   if (!retransmittable && !isPing) {
     DCHECK(!packetEvent);
@@ -1340,7 +1353,7 @@ void writeShortClose(
     const Aead& aead,
     const PacketNumberCipher& headerCipher) {
   auto header = ShortHeader(
-      ProtectionType::KeyPhaseZero,
+      connection.oneRttWritePhase,
       connId,
       getNextPacketNum(connection, PacketNumberSpace::AppData));
   writeCloseCommon(
@@ -1919,4 +1932,66 @@ void updateOneRttWriteCipher(
   conn.oneRttWritePhase = oneRttPhase;
   conn.oneRttWritePacketsSentInCurrentPhase = 0;
 }
+
+void maybeHandleIncomingKeyUpdate(QuicConnectionStateBase& conn) {
+  if (conn.readCodec->getCurrentOneRttReadPhase() != conn.oneRttWritePhase) {
+    // Peer has initiated a key update.
+    updateOneRttWriteCipher(
+        conn,
+        conn.handshakeLayer->getNextOneRttWriteCipher(),
+        conn.readCodec->getCurrentOneRttReadPhase());
+
+    conn.readCodec->setNextOneRttReadCipher(
+        conn.handshakeLayer->getNextOneRttReadCipher());
+  }
+}
+
+void maybeInitiateKeyUpdate(QuicConnectionStateBase& conn) {
+  if (conn.transportSettings.initiateKeyUpdate &&
+      conn.oneRttWritePacketsSentInCurrentPhase >
+          conn.transportSettings.keyUpdatePacketCountInterval) {
+    if (conn.readCodec->canInitiateKeyUpdate()) {
+      conn.readCodec->advanceOneRttReadPhase();
+
+      updateOneRttWriteCipher(
+          conn,
+          conn.handshakeLayer->getNextOneRttWriteCipher(),
+          conn.readCodec->getCurrentOneRttReadPhase());
+      conn.readCodec->setNextOneRttReadCipher(
+          conn.handshakeLayer->getNextOneRttReadCipher());
+      // Signal the transport that a key update has been initiated.
+      conn.oneRttWritePendingVerification = true;
+      conn.oneRttWritePendingVerificationPacketNumber.clear();
+    }
+  }
+}
+
+void maybeVerifyPendingKeyUpdate(
+    QuicConnectionStateBase& conn,
+    const OutstandingPacketWrapper& outstandingPacket,
+    const RegularQuicPacket& ackPacket) {
+  if (!(protectionTypeToEncryptionLevel(
+            outstandingPacket.packet.header.getProtectionType()) ==
+        EncryptionLevel::AppData)) {
+    // This is not an app data packet. We can't have initiated a key update yet.
+    return;
+  }
+
+  if (conn.oneRttWritePendingVerificationPacketNumber &&
+      outstandingPacket.packet.header.getPacketSequenceNum() >=
+          conn.oneRttWritePendingVerificationPacketNumber.value()) {
+    // There is a pending key update. This packet should be acked in
+    // the current phase.
+    if (ackPacket.header.getProtectionType() == conn.oneRttWritePhase) {
+      // Key update is verified.
+      conn.oneRttWritePendingVerificationPacketNumber.clear();
+      conn.oneRttWritePendingVerification = false;
+    } else {
+      throw QuicTransportException(
+          "Packet with key update was acked in the wrong phase",
+          TransportErrorCode::CRYPTO_ERROR);
+    }
+  }
+}
+
 } // namespace quic

quic/api/QuicTransportFunctions.h

@@ -333,4 +333,11 @@ void updateOneRttWriteCipher(
     QuicConnectionStateBase& conn,
     std::unique_ptr<Aead> aead,
     ProtectionType oneRttPhase);
+void maybeHandleIncomingKeyUpdate(QuicConnectionStateBase& conn);
+void maybeInitiateKeyUpdate(QuicConnectionStateBase& conn);
+void maybeVerifyPendingKeyUpdate(
+    QuicConnectionStateBase& conn,
+    const OutstandingPacketWrapper& outstandingPacket,
+    const RegularQuicPacket& ackPacket);
+
 } // namespace quic

quic/api/test/QuicTypedTransportTest.cpp

@@ -1295,6 +1295,173 @@ TYPED_TEST(QuicTypedTransportAfterStartTest, HandleIncomingKeyUpdate) {
   this->destroyTransport();
 }
 
+/**
+ * Initiate a key update - Successful attempt
+ */
+TYPED_TEST(QuicTypedTransportAfterStartTest, InitiateKeyUpdateSuccess) {
+  ASSERT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseZero);
+  auto numberOfWrittenPacketsInPhase =
+      this->getConn().oneRttWritePacketsSentInCurrentPhase;
+
+  auto streamId = this->getTransport()->createBidirectionalStream().value();
+
+  {
+    // Send and receive a packet in the current phase
+    this->getTransport()->writeChain(
+        streamId, IOBuf::copyBuffer("hello"), false);
+    this->loopForWrites();
+    EXPECT_EQ(
+        this->getConn().oneRttWritePacketsSentInCurrentPhase,
+        ++numberOfWrittenPacketsInPhase);
+
+    this->deliverPacket(this->buildPeerPacketWithStreamData(
+        streamId, IOBuf::copyBuffer("hello2"), ProtectionType::KeyPhaseZero));
+
+    // Both read and writer ciphers should be in phase zero.
+    EXPECT_EQ(
+        this->getConn().readCodec->getCurrentOneRttReadPhase(),
+        ProtectionType::KeyPhaseZero);
+    EXPECT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseZero);
+  }
+
+  {
+    // Force initiate a key update
+    QuicConnectionStateBase& conn = this->getNonConstConn();
+    conn.transportSettings.initiateKeyUpdate = true;
+    conn.transportSettings.keyUpdatePacketCountInterval =
+        numberOfWrittenPacketsInPhase;
+
+    // A key update should be triggered after this write is completed.
+    this->getTransport()->writeChain(
+        streamId, IOBuf::copyBuffer("hello3"), false);
+    this->loopForWrites();
+
+    numberOfWrittenPacketsInPhase = 0;
+    EXPECT_EQ(
+        this->getConn().oneRttWritePacketsSentInCurrentPhase,
+        numberOfWrittenPacketsInPhase);
+
+    // Both read and writer ciphers should have advanced to phase one.
+    EXPECT_EQ(
+        this->getConn().readCodec->getCurrentOneRttReadPhase(),
+        ProtectionType::KeyPhaseOne);
+    EXPECT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseOne);
+  }
+
+  {
+    // Another key update should not be allowed until this one is verified
+    QuicConnectionStateBase& conn = this->getNonConstConn();
+    EXPECT_FALSE(conn.readCodec->canInitiateKeyUpdate());
+    EXPECT_FALSE(conn.readCodec->advanceOneRttReadPhase());
+  }
+
+  {
+    // Receiving a packet in the new phase verifies the pending key update
+    this->deliverPacket(this->buildPeerPacketWithStreamData(
+        streamId, IOBuf::copyBuffer("hello4"), ProtectionType::KeyPhaseOne));
+
+    EXPECT_TRUE(this->getConn().readCodec->canInitiateKeyUpdate());
+  }
+
+  this->getNonConstConn().outstandings.reset();
+
+  this->destroyTransport();
+}
+
+/**
+ * Initiate a key update - Failed attempt
+ */
+TYPED_TEST(QuicTypedTransportAfterStartTest, InitiateKeyUpdateFailure) {
+  ASSERT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseZero);
+  auto numberOfWrittenPacketsInPhase =
+      this->getConn().oneRttWritePacketsSentInCurrentPhase;
+
+  auto streamId = this->getTransport()->createBidirectionalStream().value();
+
+  {
+    // Send and receive a packet in the current phase
+    this->getTransport()->writeChain(
+        streamId, IOBuf::copyBuffer("hello"), false);
+    this->loopForWrites();
+    EXPECT_EQ(
+        this->getConn().oneRttWritePacketsSentInCurrentPhase,
+        ++numberOfWrittenPacketsInPhase);
+
+    this->deliverPacket(this->buildPeerPacketWithStreamData(
+        streamId, IOBuf::copyBuffer("hello2"), ProtectionType::KeyPhaseZero));
+
+    // Both read and writer ciphers should be in phase zero.
+    EXPECT_EQ(
+        this->getConn().readCodec->getCurrentOneRttReadPhase(),
+        ProtectionType::KeyPhaseZero);
+    EXPECT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseZero);
+  }
+
+  {
+    // Force initiate a key update.
+    QuicConnectionStateBase& conn = this->getNonConstConn();
+    conn.transportSettings.initiateKeyUpdate = true;
+    conn.transportSettings.keyUpdatePacketCountInterval =
+        numberOfWrittenPacketsInPhase;
+
+    // A key update should be triggered after this write is completed.
+    this->getTransport()->writeChain(
+        streamId, IOBuf::copyBuffer("hello3"), false);
+    const auto maybeWriteInterval = this->loopForWrites();
+    ASSERT_EQ(1, this->getNumPacketsWritten(maybeWriteInterval));
+    // const auto packetSentTime =
+    //     CHECK_NOTNULL(this->getNewestAppDataOutstandingPacket())->metadata.time;
+
+    numberOfWrittenPacketsInPhase = 0;
+    EXPECT_EQ(
+        this->getConn().oneRttWritePacketsSentInCurrentPhase,
+        numberOfWrittenPacketsInPhase);
+
+    // Both read and writer ciphers should have advanced to phase one.
+    EXPECT_EQ(
+        this->getConn().readCodec->getCurrentOneRttReadPhase(),
+        ProtectionType::KeyPhaseOne);
+    EXPECT_EQ(this->getConn().oneRttWritePhase, ProtectionType::KeyPhaseOne);
+  }
+
+  {
+    // Send a packet in the new phase.
+    this->getTransport()->writeChain(
+        streamId, IOBuf::copyBuffer("hello3"), false);
+    const auto maybeWriteInterval = this->loopForWrites();
+    ASSERT_EQ(1, this->getNumPacketsWritten(maybeWriteInterval));
+    EXPECT_EQ(
+        this->getConn().oneRttWritePacketsSentInCurrentPhase,
+        ++numberOfWrittenPacketsInPhase);
+
+    // The peer acks the new packet in the old phase. This is a CRYPTO_ERROR.
+    auto oldestOutstandingPkt =
+        this->getOldestOutstandingPacket(PacketNumberSpace::AppData);
+    auto newestOutstandingPkt =
+        this->getNewestOutstandingPacket(PacketNumberSpace::AppData);
+    quic::AckBlocks acks = {
+        {oldestOutstandingPkt->packet.header.getPacketSequenceNum(),
+         newestOutstandingPkt->packet.header.getPacketSequenceNum()}};
+    auto buf = quic::test::packetToBuf(
+        AckPacketBuilder()
+            .setDstConn(&this->getNonConstConn())
+            .setPacketNumberSpace(PacketNumberSpace::AppData)
+            .setAckPacketNumStore(&this->peerPacketNumStore)
+            .setAckBlocks(acks)
+            .setAckDelay(0us)
+            .setShortHeaderProtectionType(ProtectionType::KeyPhaseZero)
+            .build());
+    buf->coalesce();
+
+    // Crypto error: Packet with key update was acked in the wrong phase
+    EXPECT_THROW(this->deliverPacket(std::move(buf)), std::runtime_error);
+  }
+
+  this->getNonConstConn().outstandings.reset();
+
+  this->destroyTransport();
+}
+
 template <typename T>
 struct AckEventMatcherBuilder {
   using Builder = AckEventMatcherBuilder;

quic/client/QuicClientTransport.cpp

@@ -372,18 +372,6 @@ void QuicClientTransport::processUdpPacketData(
         "Invalid connection id", TransportErrorCode::PROTOCOL_VIOLATION);
   }
 
-  if (conn_->readCodec->getCurrentOneRttReadPhase() !=
-      conn_->oneRttWritePhase) {
-    // Peer has initiated a key update.
-    updateOneRttWriteCipher(
-        *conn_,
-        clientConn_->clientHandshakeLayer->getNextOneRttWriteCipher(),
-        conn_->readCodec->getCurrentOneRttReadPhase());
-
-    conn_->readCodec->setNextOneRttReadCipher(
-        clientConn_->clientHandshakeLayer->getNextOneRttReadCipher());
-  }
-
   // Add the packet to the AckState associated with the packet number space.
   auto& ackState = getAckState(*conn_, pnSpace);
   uint64_t distanceFromExpectedPacketNum =
@@ -419,6 +407,8 @@ void QuicClientTransport::processUdpPacketData(
                 // processing loop.
                 conn_->handshakeLayer->handshakeConfirmed();
               }
+              maybeVerifyPendingKeyUpdate(
+                  *conn_, outstandingPacket, regularPacket);
               switch (packetFrame.type()) {
                 case QuicWriteFrame::Type::WriteAckFrame: {
                   const WriteAckFrame& frame = *packetFrame.asWriteAckFrame();
@@ -664,6 +654,8 @@ void QuicClientTransport::processUdpPacketData(
     handshakeConfirmed(*conn_);
   }
 
+  maybeHandleIncomingKeyUpdate(*conn_);
+
   // Try reading bytes off of crypto, and performing a handshake.
   auto cryptoData = readDataFromCryptoStream(
       *getCryptoStream(*conn_->cryptoState, encryptionLevel));
@@ -923,6 +915,7 @@ void QuicClientTransport::writeData() {
   // use.
   SCOPE_EXIT {
     conn_->pendingEvents.numProbePackets = {};
+    maybeInitiateKeyUpdate(*conn_);
   };
   if (conn_->initialWriteCipher) {
     auto& initialCryptoStream =

quic/client/handshake/ClientHandshake.h

@@ -62,14 +62,14 @@ class ClientHandshake : public Handshake {
    * one rtt write cipher using the current traffic secret and advance the
    * traffic secret.
    */
-  std::unique_ptr<Aead> getNextOneRttWriteCipher();
+  std::unique_ptr<Aead> getNextOneRttWriteCipher() override;
 
   /**
    * An API to get oneRttReadCiphers on key rotation. Each call will return a
    * one rtt read cipher using the current traffic secret and advance the
    * traffic secret.
    */
-  std::unique_ptr<Aead> getNextOneRttReadCipher();
+  std::unique_ptr<Aead> getNextOneRttReadCipher() override;
 
   /**
    * Triggered when we have received a handshake done frame from the server.

quic/common/test/TestPacketBuilders.cpp

@@ -51,6 +51,12 @@ AckPacketBuilder&& AckPacketBuilder::setAead(const Aead* aeadIn) {
   return std::move(*this);
 }
 
+AckPacketBuilder&& AckPacketBuilder::setShortHeaderProtectionType(
+    ProtectionType protectionTypeIn) {
+  shortHeaderProtectionType = protectionTypeIn;
+  return std::move(*this);
+}
+
 RegularQuicPacketBuilder::Packet AckPacketBuilder::build() && {
   // This function sends ACK to dstConn
   auto srcConnId =
@@ -119,7 +125,7 @@ RegularQuicPacketBuilder::Packet AckPacketBuilder::build() && {
         ackPacketNum,
         QuicVersion::MVFST);
   } else {
-    header = ShortHeader(ProtectionType::KeyPhaseZero, dstConnId, ackPacketNum);
+    header = ShortHeader(shortHeaderProtectionType, dstConnId, ackPacketNum);
   }
   RegularQuicPacketBuilder builder(
       CHECK_NOTNULL(dstConn)->udpSendPacketLen,

quic/common/test/TestPacketBuilders.h

@@ -30,6 +30,7 @@ struct AckPacketBuilderFields {
   folly::Optional<quic::AckBlocks> maybeAckBlocks;
   folly::Optional<std::chrono::microseconds> maybeAckDelay;
   folly::Optional<const Aead*> maybeAead; // not required
+  ProtectionType shortHeaderProtectionType{ProtectionType::KeyPhaseZero};
   explicit AckPacketBuilderFields() = default;
 };
 
@@ -42,6 +43,7 @@ struct AckPacketBuilder : public AckPacketBuilderFields {
   Builder&& setAckBlocks(const quic::AckBlocks& ackBlocksIn);
   Builder&& setAckDelay(std::chrono::microseconds ackDelayIn);
   Builder&& setAead(const Aead* aeadIn);
+  Builder&& setShortHeaderProtectionType(ProtectionType protectionTypeIn);
   RegularQuicPacketBuilder::Packet build() &&;
   quic::Buf buildBuf() &&;
   explicit AckPacketBuilder() = default;

quic/handshake/HandshakeLayer.h

@@ -10,6 +10,7 @@
 #include <quic/QuicConstants.h>
 #include <quic/codec/PacketNumberCipher.h>
 #include <quic/codec/Types.h>
+#include <quic/handshake/Aead.h>
 
 namespace quic {
 
@@ -25,6 +26,20 @@ class Handshake {
   virtual const folly::Optional<std::string>& getApplicationProtocol()
       const = 0;
 
+  /**
+   * An API to get oneRttReadCiphers on key rotation. Each call will return a
+   * one rtt read cipher using the current traffic secret and advance the
+   * traffic secret.
+   */
+  virtual std::unique_ptr<Aead> getNextOneRttReadCipher() = 0;
+
+  /**
+   * An API to get oneRttWriteCiphers on key rotation. Each call will return a
+   * one rtt write cipher using the current traffic secret and advance the
+   * traffic secret.
+   */
+  virtual std::unique_ptr<Aead> getNextOneRttWriteCipher() = 0;
+
   virtual void handshakeConfirmed() {
     LOG(FATAL) << "Not implemented";
   }

quic/server/QuicServerTransport.cpp

@@ -267,6 +267,7 @@ void QuicServerTransport::writeData() {
   // use.
   SCOPE_EXIT {
     conn_->pendingEvents.numProbePackets = {};
+    maybeInitiateKeyUpdate(*conn_);
   };
   if (conn_->initialWriteCipher) {
     auto& initialCryptoStream =

quic/server/handshake/ServerHandshake.h

@@ -123,7 +123,7 @@ class ServerHandshake : public Handshake {
    * one rtt write cipher using the current traffic secret and advance the
    * traffic secret.
    */
-  std::unique_ptr<Aead> getNextOneRttWriteCipher();
+  std::unique_ptr<Aead> getNextOneRttWriteCipher() override;
 
   /**
    * An API to get oneRttReadCiphers. Each call will generate a one rtt
@@ -137,7 +137,7 @@ class ServerHandshake : public Handshake {
    * one rtt read cipher using the current traffic secret and advance the
    * traffic secret.
    */
-  std::unique_ptr<Aead> getNextOneRttReadCipher();
+  std::unique_ptr<Aead> getNextOneRttReadCipher() override;
 
   /**
    * An edge triggered API to get the zeroRttReadCipher. Once you receive the

quic/server/state/ServerStateMachine.cpp

@@ -1018,17 +1018,6 @@ void onServerReadDataFromOpen(
       }
     }
 
-    if (conn.readCodec->getCurrentOneRttReadPhase() != conn.oneRttWritePhase) {
-      // Peer has initiated a key update.
-      updateOneRttWriteCipher(
-          conn,
-          conn.serverHandshakeLayer->getNextOneRttWriteCipher(),
-          conn.readCodec->getCurrentOneRttReadPhase());
-
-      conn.readCodec->setNextOneRttReadCipher(
-          conn.serverHandshakeLayer->getNextOneRttReadCipher());
-    }
-
     auto& ackState = getAckState(conn, packetNumberSpace);
     uint64_t distanceFromExpectedPacketNum = addPacketToAckState(
         conn, ackState, packetNum, readData.udpPacket.timings);
@@ -1053,9 +1042,12 @@ void onServerReadDataFromOpen(
               conn,
               packetNumberSpace,
               ackFrame,
-              [&](const OutstandingPacketWrapper&,
+              [&](const OutstandingPacketWrapper& outstandingPacket,
                   const QuicWriteFrame& packetFrame,
                   const ReadAckFrame&) {
+                maybeVerifyPendingKeyUpdate(
+                    conn, outstandingPacket, regularPacket);
+
                 switch (packetFrame.type()) {
                   case QuicWriteFrame::Type::WriteStreamFrame: {
                     const WriteStreamFrame& frame =
@@ -1310,6 +1302,8 @@ void onServerReadDataFromOpen(
       handshakeConfirmed(conn);
     }
 
+    maybeHandleIncomingKeyUpdate(conn);
+
     // Update writable limit before processing the handshake data. This is so
     // that if we haven't decided whether or not to validate the peer, we won't
     // increase the limit.

quic/state/StateData.h

@@ -350,6 +350,15 @@ struct QuicConnectionStateBase : public folly::DelayedDestruction {
   std::unique_ptr<Aead> oneRttWriteCipher;
   ProtectionType oneRttWritePhase{ProtectionType::KeyPhaseZero};
   uint64_t oneRttWritePacketsSentInCurrentPhase = 0;
+  // Whether there is a locally-initiated key update that is pending
+  // verification. This is used to decide whether the next variable should be
+  // updated or not.
+  bool oneRttWritePendingVerification{false};
+  // In a key update is pending, this holds the first packet number sent in the
+  // current (updated) phase. The peer must acknowledge this packet in the same
+  // phase, responding in a different phase is a protocol violation. Once the
+  // packet is acked, this value will be cleared.
+  folly::Optional<PacketNum> oneRttWritePendingVerificationPacketNumber;
 
   // Write cipher for packets with initial keys.
   std::unique_ptr<Aead> initialWriteCipher;

quic/state/TransportSettings.h

@@ -332,6 +332,11 @@ struct TransportSettings {
   // allowing cwnd growth. 0 disables. The amount given to callbacks has the
   // current amount of stream bytes buffered subtracted from it.
   uint8_t backpressureHeadroomFactor{0};
+
+  // Whether to initiate key updates
+  bool initiateKeyUpdate{false};
+  // How many sent packets to send before initiating a key update
+  uint64_t keyUpdatePacketCountInterval{kDefaultKeyUpdatePacketCountInterval};
 };
 
 } // namespace quic