From baf0544fc105e13d6bc96a2b3451eed7952d9145 Mon Sep 17 00:00:00 2001
From: Matt Joras <mjoras@meta.com>
Date: Wed, 25 Sep 2024 15:01:16 -0700
Subject: [PATCH] Add option to allow PTOing pure PING packets.

Summary: As in title. Without this the keepalive option is less effective since a single PING packet loss can cause issues.

Reviewed By: kvtsoy

Differential Revision: D63397494

fbshipit-source-id: 7ef6b6f54189609e3a96409ac9c035c475dc0569
---
 quic/api/QuicTransportFunctions.cpp          |  3 +++
 quic/api/test/QuicTransportFunctionsTest.cpp | 12 ++++++++++
 quic/codec/QuicPacketRebuilder.cpp           |  3 +++
 quic/codec/test/QuicPacketRebuilderTest.cpp  | 25 ++++++++++++++++++++
 quic/state/TransportSettings.h               |  2 ++
 5 files changed, 45 insertions(+)

diff --git a/quic/api/QuicTransportFunctions.cpp b/quic/api/QuicTransportFunctions.cpp
index 221608422..5bc07c165 100644
--- a/quic/api/QuicTransportFunctions.cpp
+++ b/quic/api/QuicTransportFunctions.cpp
@@ -784,6 +784,9 @@ void updateConnection(
       case QuicWriteFrame::Type::PingFrame:
         conn.pendingEvents.sendPing = false;
         isPing = true;
+        if (conn.transportSettings.ptoPingFrames) {
+          retransmittable = true;
+        }
         break;
       case QuicWriteFrame::Type::QuicSimpleFrame: {
         const QuicSimpleFrame& simpleFrame = *frame.asQuicSimpleFrame();
diff --git a/quic/api/test/QuicTransportFunctionsTest.cpp b/quic/api/test/QuicTransportFunctionsTest.cpp
index 39c988fd0..845fd04a9 100644
--- a/quic/api/test/QuicTransportFunctionsTest.cpp
+++ b/quic/api/test/QuicTransportFunctionsTest.cpp
@@ -218,6 +218,18 @@ TEST_F(QuicTransportFunctionsTest, PingPacketGoesToOPList) {
   EXPECT_FALSE(conn->pendingEvents.setLossDetectionAlarm);
 }
 
+TEST_F(QuicTransportFunctionsTest, PingPacketGoesToOPListAndLossAlarm) {
+  auto conn = createConn();
+  conn->transportSettings.ptoPingFrames = true;
+  auto packet = buildEmptyPacket(*conn, PacketNumberSpace::AppData);
+  packet.packet.frames.push_back(PingFrame());
+  EXPECT_EQ(0, conn->outstandings.packets.size());
+  updateConnection(
+      *conn, none, packet.packet, Clock::now(), 50, 0, false /* isDSRPacket */);
+  EXPECT_EQ(1, conn->outstandings.packets.size());
+  EXPECT_TRUE(conn->pendingEvents.setLossDetectionAlarm);
+}
+
 TEST_F(QuicTransportFunctionsTest, TestUpdateConnection) {
   auto conn = createConn();
   auto mockCongestionController =
diff --git a/quic/codec/QuicPacketRebuilder.cpp b/quic/codec/QuicPacketRebuilder.cpp
index c18b5426d..c7f149dea 100644
--- a/quic/codec/QuicPacketRebuilder.cpp
+++ b/quic/codec/QuicPacketRebuilder.cpp
@@ -176,6 +176,9 @@ Optional<ClonedPacketIdentifier> PacketRebuilder::rebuildFromPacket(
       case QuicWriteFrame::Type::PingFrame: {
         const PingFrame& pingFrame = *frame.asPingFrame();
         writeSuccess = writeFrame(pingFrame, builder_) != 0;
+        if (conn_.transportSettings.ptoPingFrames) {
+          notPureAck |= writeSuccess;
+        }
         break;
       }
       case QuicWriteFrame::Type::QuicSimpleFrame: {
diff --git a/quic/codec/test/QuicPacketRebuilderTest.cpp b/quic/codec/test/QuicPacketRebuilderTest.cpp
index ebcba1101..abc257544 100644
--- a/quic/codec/test/QuicPacketRebuilderTest.cpp
+++ b/quic/codec/test/QuicPacketRebuilderTest.cpp
@@ -546,6 +546,31 @@ TEST_F(QuicPacketRebuilderTest, PurePingWontRebuild) {
   EXPECT_EQ(0, conn.outstandings.numClonedPackets());
 }
 
+TEST_F(QuicPacketRebuilderTest, PurePingWillRebuild) {
+  ShortHeader shortHeader1(
+      ProtectionType::KeyPhaseZero, getTestConnectionId(), 0);
+  RegularQuicPacketBuilder regularBuilder(
+      kDefaultUDPSendPacketLen, std::move(shortHeader1), 0);
+  regularBuilder.encodePacketHeader();
+  PingFrame pingFrame;
+  writeFrame(pingFrame, regularBuilder);
+  auto packet = std::move(regularBuilder).buildPacket();
+  auto outstandingPacket = makeDummyOutstandingPacket(packet.packet, 50);
+  EXPECT_EQ(1, outstandingPacket.packet.frames.size());
+  QuicServerConnectionState conn(
+      FizzServerQuicHandshakeContext::Builder().build());
+  conn.transportSettings.ptoPingFrames = true;
+  ShortHeader shortHeader2(
+      ProtectionType::KeyPhaseZero, getTestConnectionId(), 0);
+  RegularQuicPacketBuilder regularBuilder2(
+      kDefaultUDPSendPacketLen, std::move(shortHeader2), 0);
+  regularBuilder2.encodePacketHeader();
+  PacketRebuilder rebuilder(regularBuilder2, conn);
+  EXPECT_TRUE(rebuilder.rebuildFromPacket(outstandingPacket).has_value());
+  EXPECT_TRUE(outstandingPacket.maybeClonedPacketIdentifier.has_value());
+  EXPECT_EQ(1, conn.outstandings.numClonedPackets());
+}
+
 TEST_F(QuicPacketRebuilderTest, LastStreamFrameSkipLen) {
   QuicServerConnectionState conn(
       FizzServerQuicHandshakeContext::Builder().build());
diff --git a/quic/state/TransportSettings.h b/quic/state/TransportSettings.h
index 6140feec4..e11d8ff5b 100644
--- a/quic/state/TransportSettings.h
+++ b/quic/state/TransportSettings.h
@@ -410,6 +410,8 @@ struct TransportSettings {
   // receiving and then processing.
   bool networkDataPerSocketRead{false};
   bool cloneAllPacketsWithCryptoFrame{false};
+  // Whether we allow PTOing of packets with just PING frames.
+  bool ptoPingFrames{false};
 };
 
 } // namespace quic