Move StatelessRejector to a heap-allocated model

net/tools/quic/quic_dispatcher.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "net/tools/quic/quic_dispatcher.h"
 
 #include <utility>
 
 #include "base/debug/stack_trace.h"
 #include "base/logging.h"
@@ skipping 310 matching lines @@
     time_wait_list_manager_->ProcessPacket(
         current_server_address_, current_client_address_,
         header.public_header.connection_id, header.packet_number,
         *current_packet_);
     return false;
   }
 
   // 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;
 }
 
 void QuicDispatcher::ProcessUnauthenticatedHeaderFate(
     QuicPacketFate fate,
     QuicConnectionId connection_id,
     QuicPacketNumber packet_number) {
   switch (fate) {
     case kFateProcess: {
       ProcessChlo();
@@ skipping 408 matching lines @@
 
 void QuicDispatcher::SetLastError(QuicErrorCode error) {
   last_error_ = error;
 }
 
 bool QuicDispatcher::OnUnauthenticatedUnknownPublicHeader(
     const QuicPacketPublicHeader& header) {
   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 ||
       header.public_header.versions.front() <= QUIC_VERSION_32 ||
       !ShouldAttemptCheapStatelessRejection()) {
     // Not use cheap stateless reject.
     if (FLAGS_quic_buffer_packet_till_chlo &&
         !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) {
       QUIC_BUG
           << "Have to drop packet because buffering non-chlo packet is "
              "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()) {
     // This CHLO is prohibited by policy.
     StatelessConnectionTerminator terminator(connection_id, &framer_, helper(),
                                              time_wait_list_manager_.get());
     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;
     }
+
+    default:
+      QUIC_BUG << "Rejector has invalid state " << rejector->state();
+      fate = kFateDrop;
+      break;
   }
-
-  QUIC_BUG << "Rejector has unknown invalid state.";
-  return kFateDrop;
+  ProcessUnauthenticatedHeaderFate(fate, rejector->connection_id(),
+                                   packet_number);
 }
 
 const QuicVersionVector& QuicDispatcher::GetSupportedVersions() {
   return version_manager_->GetSupportedVersions();
 }
 
 void QuicDispatcher::DeliverPacketsToSession(
     const std::list<BufferedPacket>& packets,
     QuicServerSessionBase* session) {
   for (const BufferedPacket& packet : packets) {
     session->ProcessUdpPacket(packet.server_address, packet.client_address,
                               *(packet.packet));
   }
 }
 
 }  // namespace net