Move StatelessRejector to a heap-allocated model

net/tools/quic/quic_dispatcher.cc

Index: net/tools/quic/quic_dispatcher.cc
diff --git a/net/tools/quic/quic_dispatcher.cc b/net/tools/quic/quic_dispatcher.cc
index c65929e6d57b8cc00e53a6102cdd5bb58fec00a8..ba15c19aab08e11b8bff67ff55fd677c211227ca 100644
--- a/net/tools/quic/quic_dispatcher.cc
+++ b/net/tools/quic/quic_dispatcher.cc
@@ -328,11 +328,15 @@ bool QuicDispatcher::OnUnauthenticatedHeader(const QuicPacketHeader& header) {
   // Packet's connection ID is unknown.  Apply the validity checks.
   QuicPacketFate fate = ValidityChecks(header);
   if (fate == kFateProcess) {
-    fate = MaybeRejectStatelessly(connection_id, header);
+    // Execute stateless rejection logic to determine the packet fate, then
+    // invoke ProcessUnauthenticatedHeaderFate.
+    MaybeRejectStatelessly(connection_id, header);
+  } else {
+    // If the fate is already known, process it without executing stateless
+    // rejection logic.
+    ProcessUnauthenticatedHeaderFate(fate, connection_id, header.packet_number);
   }
 
-  // Perform whatever actions are needed on the current packet.
-  ProcessUnauthenticatedHeaderFate(fate, connection_id, header.packet_number);
   return false;
 }
 
@@ -761,9 +765,29 @@ bool QuicDispatcher::OnUnauthenticatedUnknownPublicHeader(
   return true;
 }
 
-QuicDispatcher::QuicPacketFate QuicDispatcher::MaybeRejectStatelessly(
-    QuicConnectionId connection_id,
-    const QuicPacketHeader& header) {
+class StatelessRejectorProcessDoneCallback
+    : public StatelessRejector::ProcessDoneCallback {
+ public:
+  StatelessRejectorProcessDoneCallback(QuicDispatcher* dispatcher,
+                                       QuicPacketNumber packet_number,
+                                       QuicVersion first_version)
+      : dispatcher_(dispatcher),
+        packet_number_(packet_number),
+        first_version_(first_version) {}
+
+  void Run(std::unique_ptr<StatelessRejector> rejector) override {
+    dispatcher_->OnStatelessRejectorProcessDone(std::move(rejector),
+                                                packet_number_, first_version_);
+  }
+
+ private:
+  QuicDispatcher* dispatcher_;
+  QuicPacketNumber packet_number_;
+  QuicVersion first_version_;
+};
+
+void QuicDispatcher::MaybeRejectStatelessly(QuicConnectionId connection_id,
+                                            const QuicPacketHeader& header) {
   // TODO(rch): This logic should probably live completely inside the rejector.
   if (!FLAGS_quic_use_cheap_stateless_rejects ||
       !FLAGS_enable_quic_stateless_reject_support ||
@@ -774,18 +798,22 @@ QuicDispatcher::QuicPacketFate QuicDispatcher::MaybeRejectStatelessly(
         !ChloExtractor::Extract(*current_packet_, GetSupportedVersions(),
                                 nullptr)) {
       // Buffer non-CHLO packets.
-      return kFateBuffer;
+      ProcessUnauthenticatedHeaderFate(kFateBuffer, connection_id,
+                                       header.packet_number);
+      return;
     }
-    return kFateProcess;
+    ProcessUnauthenticatedHeaderFate(kFateProcess, connection_id,
+                                     header.packet_number);
+    return;
   }
 
-  StatelessRejector rejector(
+  std::unique_ptr<StatelessRejector> rejector(new StatelessRejector(
       header.public_header.versions.front(), GetSupportedVersions(),
       crypto_config_, &compressed_certs_cache_, helper()->GetClock(),
       helper()->GetRandomGenerator(), current_packet_->length(),
-      current_client_address_, current_server_address_);
+      current_client_address_, current_server_address_));
   ChloValidator validator(session_helper_.get(), current_server_address_,
-                          &rejector);
+                          rejector.get());
   if (!ChloExtractor::Extract(*current_packet_, GetSupportedVersions(),
                               &validator)) {
     if (!FLAGS_quic_buffer_packet_till_chlo) {
@@ -794,9 +822,13 @@ QuicDispatcher::QuicPacketFate QuicDispatcher::MaybeRejectStatelessly(
              "not supported while trying to do stateless reject. "
              "--gfe2_reloadable_flag_quic_buffer_packet_till_chlo false"
              " --gfe2_reloadable_flag_quic_use_cheap_stateless_rejects true";
-      return kFateDrop;
+      ProcessUnauthenticatedHeaderFate(kFateDrop, connection_id,
+                                       header.packet_number);
+      return;
     }
-    return kFateBuffer;
+    ProcessUnauthenticatedHeaderFate(kFateBuffer, connection_id,
+                                     header.packet_number);
+    return;
   }
 
   if (!validator.can_accept()) {
@@ -806,40 +838,63 @@ QuicDispatcher::QuicPacketFate QuicDispatcher::MaybeRejectStatelessly(
     terminator.CloseConnection(QUIC_HANDSHAKE_FAILED,
                                validator.error_details());
     OnConnectionClosedStatelessly(QUIC_HANDSHAKE_FAILED);
-    return kFateTimeWait;
+    ProcessUnauthenticatedHeaderFate(kFateTimeWait, connection_id,
+                                     header.packet_number);
+    return;
   }
 
-  // This packet included a CHLO. See if it can be rejected statelessly.
-  switch (rejector.state()) {
+  // Continue stateless rejector processing
+  std::unique_ptr<StatelessRejectorProcessDoneCallback> cb(
+      new StatelessRejectorProcessDoneCallback(
+          this, header.packet_number, header.public_header.versions.front()));
+  StatelessRejector::Process(std::move(rejector), std::move(cb));
+}
+
+void QuicDispatcher::OnStatelessRejectorProcessDone(
+    std::unique_ptr<StatelessRejector> rejector,
+    QuicPacketNumber packet_number,
+    QuicVersion first_version) {
+  QuicPacketFate fate;
+  switch (rejector->state()) {
     case StatelessRejector::FAILED: {
       // There was an error processing the client hello.
-      StatelessConnectionTerminator terminator(
-          connection_id, &framer_, helper(), time_wait_list_manager_.get());
-      terminator.CloseConnection(rejector.error(), rejector.error_details());
-      return kFateTimeWait;
+      StatelessConnectionTerminator terminator(rejector->connection_id(),
+                                               &framer_, helper(),
+                                               time_wait_list_manager_.get());
+      terminator.CloseConnection(rejector->error(), rejector->error_details());
+      fate = kFateTimeWait;
+      break;
     }
 
     case StatelessRejector::UNSUPPORTED:
       // Cheap stateless rejects are not supported so process the packet.
-      return kFateProcess;
+      fate = kFateProcess;
+      break;
 
     case StatelessRejector::ACCEPTED:
       // Contains a valid CHLO, so process the packet and create a connection.
-      return kFateProcess;
+      fate = kFateProcess;
+      break;
 
     case StatelessRejector::REJECTED: {
-      DCHECK_EQ(framer_.version(), header.public_header.versions.front());
-      StatelessConnectionTerminator terminator(
-          connection_id, &framer_, helper(), time_wait_list_manager_.get());
+      DCHECK_EQ(framer_.version(), first_version);
+      StatelessConnectionTerminator terminator(rejector->connection_id(),
+                                               &framer_, helper(),
+                                               time_wait_list_manager_.get());
       terminator.RejectConnection(
-          rejector.reply().GetSerialized().AsStringPiece());
+          rejector->reply().GetSerialized().AsStringPiece());
       OnConnectionRejectedStatelessly();
-      return kFateTimeWait;
+      fate = kFateTimeWait;
+      break;
     }
-  }
 
-  QUIC_BUG << "Rejector has unknown invalid state.";
-  return kFateDrop;
+    default:
+      QUIC_BUG << "Rejector has invalid state " << rejector->state();
+      fate = kFateDrop;
+      break;
+  }
+  ProcessUnauthenticatedHeaderFate(fate, rejector->connection_id(),
+                                   packet_number);
 }
 
 const QuicVersionVector& QuicDispatcher::GetSupportedVersions() {