Enforce strict amplification limit:
1) throttle sending if next created packet + bytes_sent >= factor * bytes_received,
2) buffer coalesced packet if size + bytes_sent >= factor * bytes_received.

Protected by FLAGS_quic_strict_amplification_factor.

PiperOrigin-RevId: 458577060
diff --git a/quiche/quic/core/http/end_to_end_test.cc b/quiche/quic/core/http/end_to_end_test.cc
index c957aba..f595893 100644
--- a/quiche/quic/core/http/end_to_end_test.cc
+++ b/quiche/quic/core/http/end_to_end_test.cc
@@ -2850,6 +2850,9 @@
 
 TEST_P(EndToEndTest, ConnectionMigrationClientIPChanged) {
   ASSERT_TRUE(Initialize());
+  if (GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
+    return;
+  }
   SendSynchronousFooRequestAndCheckResponse();
 
   // Store the client IP address which was used to send the first request.
@@ -2890,7 +2893,8 @@
 
 TEST_P(EndToEndTest, IetfConnectionMigrationClientIPChangedMultipleTimes) {
   ASSERT_TRUE(Initialize());
-  if (!GetClientConnection()->connection_migration_use_new_cid()) {
+  if (!GetClientConnection()->connection_migration_use_new_cid() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   SendSynchronousFooRequestAndCheckResponse();
@@ -3001,7 +3005,8 @@
 
 TEST_P(EndToEndTest,
        ConnectionMigrationWithNonZeroConnectionIDClientIPChangedMultipleTimes) {
-  if (!version_.SupportsClientConnectionIds()) {
+  if (!version_.SupportsClientConnectionIds() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     ASSERT_TRUE(Initialize());
     return;
   }
@@ -3138,7 +3143,8 @@
 TEST_P(EndToEndTest, ConnectionMigrationNewTokenForNewIp) {
   ASSERT_TRUE(Initialize());
   if (!version_.HasIetfQuicFrames() ||
-      !client_->client()->session()->connection()->validate_client_address()) {
+      !client_->client()->session()->connection()->validate_client_address() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   SendSynchronousFooRequestAndCheckResponse();
diff --git a/quiche/quic/core/quic_connection.cc b/quiche/quic/core/quic_connection.cc
index 2c8d341..aeca5a1 100644
--- a/quiche/quic/core/quic_connection.cc
+++ b/quiche/quic/core/quic_connection.cc
@@ -3238,7 +3238,7 @@
     return true;
   }
 
-  if (LimitedByAmplificationFactor()) {
+  if (LimitedByAmplificationFactor(packet_creator_.max_packet_length())) {
     // Server is constrained by the amplification restriction.
     QUIC_CODE_COUNT(quic_throttled_by_amplification_limit);
     QUIC_DVLOG(1) << ENDPOINT
@@ -4739,7 +4739,7 @@
     pending_retransmission_alarm_ = true;
     return;
   }
-  if (LimitedByAmplificationFactor()) {
+  if (LimitedByAmplificationFactor(packet_creator_.max_packet_length())) {
     // Do not set retransmission timer if connection is anti-amplification limit
     // throttled. Otherwise, nothing can be sent when timer fires.
     retransmission_alarm_->Cancel();
@@ -5854,7 +5854,9 @@
     if (!flushed) {
       // Connection is write blocked.
       QUIC_BUG_IF(quic_bug_12714_33,
-                  !writer_->IsWriteBlocked() && !LimitedByAmplificationFactor())
+                  !writer_->IsWriteBlocked() &&
+                      !LimitedByAmplificationFactor(
+                          packet_creator_.max_packet_length()))
           << "Writer not blocked and not throttled by amplification factor, "
              "but ACK not flushed for packet space:"
           << i;
@@ -5976,14 +5978,20 @@
     QUIC_RELOADABLE_FLAG_COUNT(
         quic_fix_bytes_accounting_for_buffered_coalesced_packets);
   }
-  if (!buffered_packets_.empty() || HandleWriteBlocked()) {
+  const size_t padding_size =
+      length - std::min<size_t>(length, coalesced_packet_.length());
+  // Buffer coalesced packet if padding + bytes_sent exceeds amplifcation limit.
+  if (!buffered_packets_.empty() || HandleWriteBlocked() ||
+      (enforce_strict_amplification_factor_ &&
+       LimitedByAmplificationFactor(padding_size))) {
     QUIC_DVLOG(1) << ENDPOINT
                   << "Buffering coalesced packet of len: " << length;
     buffered_packets_.emplace_back(
         buffer, static_cast<QuicPacketLength>(length),
         coalesced_packet_.self_address(), coalesced_packet_.peer_address());
     if (!GetQuicReloadableFlag(
-            quic_fix_bytes_accounting_for_buffered_coalesced_packets)) {
+            quic_fix_bytes_accounting_for_buffered_coalesced_packets) &&
+        !enforce_strict_amplification_factor_) {
       return true;
     }
   } else {
@@ -6007,7 +6015,6 @@
   }
   // Account for added padding.
   if (length > coalesced_packet_.length()) {
-    size_t padding_size = length - coalesced_packet_.length();
     if (IsDefaultPath(coalesced_packet_.self_address(),
                       coalesced_packet_.peer_address())) {
       if (EnforceAntiAmplificationLimit()) {
@@ -6108,9 +6115,10 @@
 
 // TODO(danzh) Pass in path object or its reference of some sort to use this
 // method to check anti-amplification limit on non-default path.
-bool QuicConnection::LimitedByAmplificationFactor() const {
+bool QuicConnection::LimitedByAmplificationFactor(QuicByteCount bytes) const {
   return EnforceAntiAmplificationLimit() &&
-         default_path_.bytes_sent_before_address_validation >=
+         (default_path_.bytes_sent_before_address_validation +
+          (enforce_strict_amplification_factor_ ? bytes : 0)) >=
              anti_amplification_factor_ *
                  default_path_.bytes_received_before_address_validation;
 }
diff --git a/quiche/quic/core/quic_connection.h b/quiche/quic/core/quic_connection.h
index db2d320..e2c8f49 100644
--- a/quiche/quic/core/quic_connection.h
+++ b/quiche/quic/core/quic_connection.h
@@ -1709,7 +1709,9 @@
   QuicPacketNumber GetLargestAckedPacket() const;
 
   // Whether connection is limited by amplification factor.
-  bool LimitedByAmplificationFactor() const;
+  // If enforce_strict_amplification_factor_ is true, this will return true if
+  // connection is amplification limited after sending |bytes|.
+  bool LimitedByAmplificationFactor(QuicByteCount bytes) const;
 
   // Called before sending a packet to get packet send time and to set the
   // release time delay in |per_packet_options_|. Return the time when the
@@ -2240,6 +2242,11 @@
 
   bool only_send_probing_frames_on_alternative_path_ =
       GetQuicReloadableFlag(quic_not_bundle_ack_on_alternative_path);
+
+  // If true, throttle sending if next created packet will exceed amplification
+  // limit.
+  const bool enforce_strict_amplification_factor_ =
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor);
 };
 
 }  // namespace quic
diff --git a/quiche/quic/core/quic_connection_test.cc b/quiche/quic/core/quic_connection_test.cc
index 0fea5c6..13cf40f 100644
--- a/quiche/quic/core/quic_connection_test.cc
+++ b/quiche/quic/core/quic_connection_test.cc
@@ -1709,7 +1709,8 @@
 
 TEST_P(QuicConnectionTest, PeerIpAddressChangeAtServer) {
   set_perspective(Perspective::IS_SERVER);
-  if (!connection_.validate_client_address()) {
+  if (!connection_.validate_client_address() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
@@ -1905,6 +1906,9 @@
 
   QuicFrames frames2;
   frames2.push_back(QuicFrame(frame2_));
+  if (GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
+    frames2.push_back(QuicFrame(QuicPaddingFrame(-1)));
+  }
   ProcessFramesPacketWithAddresses(frames2, kSelfAddress, kNewPeerAddress,
                                    ENCRYPTION_FORWARD_SECURE);
   EXPECT_EQ(kNewPeerAddress, connection_.peer_address());
@@ -1926,6 +1930,9 @@
 }
 
 TEST_P(QuicConnectionTest, EffectivePeerAddressChangeAtServer) {
+  if (GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
+    return;
+  }
   set_perspective(Perspective::IS_SERVER);
   QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
   EXPECT_EQ(Perspective::IS_SERVER, connection_.perspective());
@@ -2102,7 +2109,8 @@
 TEST_P(QuicConnectionTest,
        ReversePathValidationResponseReceivedFromUnexpectedPeerAddress) {
   set_perspective(Perspective::IS_SERVER);
-  if (!connection_.connection_migration_use_new_cid()) {
+  if (!connection_.connection_migration_use_new_cid() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
@@ -7612,6 +7620,7 @@
 TEST_P(QuicConnectionTest, ServerReceivesChloOnNonCryptoStream) {
   set_perspective(Perspective::IS_SERVER);
   QuicPacketCreatorPeer::SetSendVersionInPacket(creator_, false);
+  QuicConnectionPeer::SetAddressValidated(&connection_);
 
   CryptoHandshakeMessage message;
   CryptoFramer framer;
@@ -7838,6 +7847,9 @@
     return;
   }
   set_perspective(Perspective::IS_SERVER);
+  // Receives packet 1000 in initial data.
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(1);
+  ProcessCryptoPacketAtLevel(1000, ENCRYPTION_INITIAL);
 
   if (version().handshake_protocol == PROTOCOL_TLS1_3) {
     EXPECT_CALL(visitor_, BeforeConnectionCloseSent());
@@ -9544,7 +9556,8 @@
 }
 
 TEST_P(QuicConnectionTest, AntiAmplificationLimit) {
-  if (!connection_.version().SupportsAntiAmplificationLimit()) {
+  if (!connection_.version().SupportsAntiAmplificationLimit() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
@@ -9600,7 +9613,8 @@
 }
 
 TEST_P(QuicConnectionTest, 3AntiAmplificationLimit) {
-  if (!connection_.version().SupportsAntiAmplificationLimit()) {
+  if (!connection_.version().SupportsAntiAmplificationLimit() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
@@ -9668,7 +9682,8 @@
 }
 
 TEST_P(QuicConnectionTest, 10AntiAmplificationLimit) {
-  if (!connection_.version().SupportsAntiAmplificationLimit()) {
+  if (!connection_.version().SupportsAntiAmplificationLimit() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
@@ -9873,7 +9888,8 @@
 TEST_P(QuicConnectionTest, FailToCoalescePacket) {
   // EXPECT_QUIC_BUG tests are expensive so only run one instance of them.
   if (!IsDefaultTestConfiguration() ||
-      !connection_.version().CanSendCoalescedPackets()) {
+      !connection_.version().CanSendCoalescedPackets() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
 
@@ -11919,6 +11935,7 @@
   frames.push_back(QuicFrame(frame1_));
   QuicPathFrameBuffer path_frame_buffer{0, 1, 2, 3, 4, 5, 6, 7};
   frames.push_back(QuicFrame(QuicPathChallengeFrame(0, path_frame_buffer)));
+  frames.push_back(QuicFrame(QuicPaddingFrame(-1)));
   const QuicSocketAddress kNewPeerAddress(QuicIpAddress::Loopback4(),
                                           /*port=*/23456);
 
@@ -14500,7 +14517,8 @@
 // Regression test for b/182571515
 TEST_P(QuicConnectionTest, LostDataThenGetAcknowledged) {
   set_perspective(Perspective::IS_SERVER);
-  if (!connection_.validate_client_address()) {
+  if (!connection_.validate_client_address() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
 
@@ -14923,7 +14941,8 @@
 
 // Regression test for b/201643321.
 TEST_P(QuicConnectionTest, FailedToRetransmitShlo) {
-  if (!version().HasIetfQuicFrames()) {
+  if (!version().HasIetfQuicFrames() ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     return;
   }
   set_perspective(Perspective::IS_SERVER);
@@ -15485,7 +15504,8 @@
   QuicConnectionPeer::SendPing(&connection_);
   const QuicConnectionStats& stats = connection_.GetStats();
   if (GetQuicReloadableFlag(
-          quic_fix_bytes_accounting_for_buffered_coalesced_packets)) {
+          quic_fix_bytes_accounting_for_buffered_coalesced_packets) ||
+      GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
     // Verify padding is accounted.
     EXPECT_EQ(stats.bytes_sent, connection_.max_packet_length());
   } else {
@@ -15493,6 +15513,80 @@
   }
 }
 
+TEST_P(QuicConnectionTest, StrictAntiAmplificationLimit) {
+  if (!connection_.version().SupportsAntiAmplificationLimit()) {
+    return;
+  }
+  EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(AnyNumber());
+  set_perspective(Perspective::IS_SERVER);
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  // Verify no data can be sent at the beginning because bytes received is 0.
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(0);
+  connection_.SendCryptoDataWithString("foo", 0);
+  EXPECT_FALSE(connection_.CanWrite(HAS_RETRANSMITTABLE_DATA));
+  EXPECT_FALSE(connection_.CanWrite(NO_RETRANSMITTABLE_DATA));
+  EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
+
+  const size_t anti_amplification_factor =
+      GetQuicFlag(FLAGS_quic_anti_amplification_factor);
+  // Receives packet 1.
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(1);
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _))
+      .Times(anti_amplification_factor);
+  ProcessCryptoPacketAtLevel(1, ENCRYPTION_INITIAL);
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetEncrypter(ENCRYPTION_FORWARD_SECURE,
+                           std::make_unique<TaggingEncrypter>(0x03));
+
+  for (size_t i = 1; i < anti_amplification_factor - 1; ++i) {
+    connection_.SendCryptoDataWithString("foo", i * 3);
+  }
+  // Send an addtion packet with max_packet_size - 1.
+  connection_.SetMaxPacketLength(connection_.max_packet_length() - 1);
+  connection_.SendCryptoDataWithString("bar",
+                                       (anti_amplification_factor - 1) * 3);
+  EXPECT_LT(writer_->total_bytes_written(),
+            anti_amplification_factor *
+                QuicConnectionPeer::BytesReceivedOnDefaultPath(&connection_));
+  if (GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
+    // 3 connection closes which will be buffered.
+    EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(3);
+    // Verify retransmission alarm is not set.
+    EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
+  } else {
+    // Crypto + 3 connection closes.
+    EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(4);
+    EXPECT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  }
+  // Try to send another packet with max_packet_size.
+  connection_.SetMaxPacketLength(connection_.max_packet_length() + 1);
+  connection_.SendCryptoDataWithString("bar", anti_amplification_factor * 3);
+  EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
+  // Close connection.
+  EXPECT_CALL(visitor_, BeforeConnectionCloseSent());
+  EXPECT_CALL(visitor_, OnConnectionClosed(_, _));
+  connection_.CloseConnection(
+      QUIC_INTERNAL_ERROR, "error",
+      ConnectionCloseBehavior::SEND_CONNECTION_CLOSE_PACKET);
+  EXPECT_EQ(0u, connection_.NumQueuedPackets());
+  if (GetQuicFlag(FLAGS_quic_enforce_strict_amplification_factor)) {
+    EXPECT_LT(writer_->total_bytes_written(),
+              anti_amplification_factor *
+                  QuicConnectionPeer::BytesReceivedOnDefaultPath(&connection_));
+  } else {
+    EXPECT_LT(writer_->total_bytes_written(),
+              (anti_amplification_factor + 2) *
+                  QuicConnectionPeer::BytesReceivedOnDefaultPath(&connection_));
+    EXPECT_GT(writer_->total_bytes_written(),
+              (anti_amplification_factor + 1) *
+                  QuicConnectionPeer::BytesReceivedOnDefaultPath(&connection_));
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quiche/quic/core/quic_protocol_flags_list.h b/quiche/quic/core/quic_protocol_flags_list.h
index 20d9a88..361bcfc 100644
--- a/quiche/quic/core/quic_protocol_flags_list.h
+++ b/quiche/quic/core/quic_protocol_flags_list.h
@@ -215,4 +215,7 @@
 
 QUIC_PROTOCOL_FLAG(bool, quic_use_lower_server_response_mtu_for_test, false,
                    "If true, cap server response packet size at 1250.")
+
+QUIC_PROTOCOL_FLAG(bool, quic_enforce_strict_amplification_factor, false,
+                   "If true, enforce strict amplification factor")
 #endif
diff --git a/quiche/quic/test_tools/quic_connection_peer.cc b/quiche/quic/test_tools/quic_connection_peer.cc
index e38bfc3..db827aa 100644
--- a/quiche/quic/test_tools/quic_connection_peer.cc
+++ b/quiche/quic/test_tools/quic_connection_peer.cc
@@ -438,6 +438,12 @@
   return &connection->path_validator_;
 }
 
+// static
+QuicByteCount QuicConnectionPeer::BytesReceivedOnDefaultPath(
+    QuicConnection* connection) {
+  return connection->default_path_.bytes_received_before_address_validation;
+}
+
 //  static
 QuicByteCount QuicConnectionPeer::BytesSentOnAlternativePath(
     QuicConnection* connection) {
diff --git a/quiche/quic/test_tools/quic_connection_peer.h b/quiche/quic/test_tools/quic_connection_peer.h
index d1a41f9..7ad551f 100644
--- a/quiche/quic/test_tools/quic_connection_peer.h
+++ b/quiche/quic/test_tools/quic_connection_peer.h
@@ -175,6 +175,8 @@
 
   static QuicPathValidator* path_validator(QuicConnection* connection);
 
+  static QuicByteCount BytesReceivedOnDefaultPath(QuicConnection* connection);
+
   static QuicByteCount BytesSentOnAlternativePath(QuicConnection* connection);
 
   static QuicByteCount BytesReceivedOnAlternativePath(
diff --git a/quiche/quic/test_tools/quic_test_utils.cc b/quiche/quic/test_tools/quic_test_utils.cc
index 6582ee7..e52130e 100644
--- a/quiche/quic/test_tools/quic_test_utils.cc
+++ b/quiche/quic/test_tools/quic_test_utils.cc
@@ -559,6 +559,11 @@
 
 PacketSavingConnection::~PacketSavingConnection() {}
 
+SerializedPacketFate PacketSavingConnection::GetSerializedPacketFate(
+    bool /*is_mtu_discovery*/, EncryptionLevel /*encryption_level*/) {
+  return SEND_TO_WRITER;
+}
+
 void PacketSavingConnection::SendOrQueuePacket(SerializedPacket packet) {
   encrypted_packets_.push_back(std::make_unique<QuicEncryptedPacket>(
       CopyBuffer(packet), packet.encrypted_length, true));
@@ -1363,6 +1368,7 @@
   }
 
   last_packet_size_ = packet.length();
+  total_bytes_written_ += packet.length();
   last_packet_header_ = framer_.header();
   if (!framer_.connection_close_frames().empty()) {
     ++connection_close_packets_;
diff --git a/quiche/quic/test_tools/quic_test_utils.h b/quiche/quic/test_tools/quic_test_utils.h
index 5e4b52e..2dfda1a 100644
--- a/quiche/quic/test_tools/quic_test_utils.h
+++ b/quiche/quic/test_tools/quic_test_utils.h
@@ -732,6 +732,9 @@
 
   ~PacketSavingConnection() override;
 
+  SerializedPacketFate GetSerializedPacketFate(
+      bool is_mtu_discovery, EncryptionLevel encryption_level) override;
+
   void SendOrQueuePacket(SerializedPacket packet) override;
 
   MOCK_METHOD(void, OnPacketSent, (EncryptionLevel, TransmissionType));
@@ -1859,7 +1862,9 @@
     return framer_.coalesced_packet();
   }
 
-  size_t last_packet_size() { return last_packet_size_; }
+  size_t last_packet_size() const { return last_packet_size_; }
+
+  size_t total_bytes_written() const { return total_bytes_written_; }
 
   const QuicPacketHeader& last_packet_header() const {
     return last_packet_header_;
@@ -1934,6 +1939,7 @@
   ParsedQuicVersion version_;
   SimpleQuicFramer framer_;
   size_t last_packet_size_ = 0;
+  size_t total_bytes_written_ = 0;
   QuicPacketHeader last_packet_header_;
   bool write_blocked_ = false;
   bool write_should_fail_ = false;