Enforce `QuicConnectionMigrationManager` and `QuicSpdyClientSessionWithMigration` to use QuicForceBlockablePacketWriter during construction.

And add a new subclass `QuicClientPathValidationContext` to enforce that client-side path validation contexts provide a `QuicForceBlockablePacketWriter`.

With both enforcement, remove unsafe static_cast from QuicConnection::writer() to `QuicForceBlockablePacketWriter` in `QuicConnectionMigrationManager` by caching the `QuicForceBlockablePacketWriter` object in QuicSpdyClientSessionWithMigration whenever there are path changes.

PiperOrigin-RevId: 804421523
diff --git a/quiche/quic/core/http/quic_connection_migration_manager.cc b/quiche/quic/core/http/quic_connection_migration_manager.cc
index 1988d84..a88dc5f 100644
--- a/quiche/quic/core/http/quic_connection_migration_manager.cc
+++ b/quiche/quic/core/http/quic_connection_migration_manager.cc
@@ -91,13 +91,13 @@
       QuicConnectionMigrationManager* absl_nonnull migration_manager)
       : migration_manager_(migration_manager) {}
   void OnPathValidationSuccess(
-      std::unique_ptr<quic::QuicPathValidationContext> context,
+      std::unique_ptr<QuicPathValidationContext> context,
       quic::QuicTime start_time) override {
     migration_manager_->OnConnectionMigrationProbeSucceeded(std::move(context),
                                                             start_time);
   }
   void OnPathValidationFailure(
-      std::unique_ptr<quic::QuicPathValidationContext> context) override {
+      std::unique_ptr<QuicPathValidationContext> context) override {
     migration_manager_->OnProbeFailed(std::move(context));
   }
 
@@ -364,7 +364,7 @@
 
 void QuicConnectionMigrationManager::
     PathContextCreationResultDelegateForImmediateMigration::OnCreationSucceeded(
-        std::unique_ptr<QuicPathValidationContext> context) {
+        std::unique_ptr<QuicClientPathValidationContext> context) {
   migration_manager_->FinishMigrate(std::move(context), close_session_on_error_,
                                     std::move(migration_callback_));
 }
@@ -372,9 +372,7 @@
 void QuicConnectionMigrationManager::
     PathContextCreationResultDelegateForImmediateMigration::OnCreationFailed(
         QuicNetworkHandle network, absl::string_view error) {
-  static_cast<QuicForceBlockablePacketWriter*>(
-      migration_manager_->connection_->writer())
-      ->ForceWriteBlocked(false);
+  migration_manager_->session_->writer()->ForceWriteBlocked(false);
   std::move(migration_callback_)(network, MigrationResult::FAILURE);
   if (close_session_on_error_) {
     migration_manager_->session_->OnConnectionToBeClosedDueToMigrationError(
@@ -398,7 +396,7 @@
 
 void QuicConnectionMigrationManager::
     PathContextCreationResultDelegateForProbing::OnCreationSucceeded(
-        std::unique_ptr<QuicPathValidationContext> context) {
+        std::unique_ptr<QuicClientPathValidationContext> context) {
   migration_manager_->FinishStartProbing(std::move(probing_callback_),
                                          std::move(context));
 }
@@ -447,8 +445,7 @@
         "Connection is migrating with an invalid network handle.");
   }
   QUIC_DVLOG(1) << "Force blocking the current packet writer";
-  static_cast<QuicForceBlockablePacketWriter*>(connection_->writer())
-      ->ForceWriteBlocked(true);
+  session_->writer()->ForceWriteBlocked(true);
   if (config_.disable_blackhole_detection_on_immediate_migrate) {
     // Turn off the black hole detector since the writer is blocked.
     // Blackhole will be re-enabled once a packet is sent again.
@@ -478,14 +475,13 @@
 }
 
 void QuicConnectionMigrationManager::FinishMigrate(
-    std::unique_ptr<QuicPathValidationContext> path_context,
+    std::unique_ptr<QuicClientPathValidationContext> path_context,
     bool close_session_on_error, MigrationCallback callback) {
   // Migrate to the new socket.
   MigrationCause current_migration_cause = current_migration_cause_;
   QuicNetworkHandle network = path_context->network();
   if (!session_->MigrateToNewPath(std::move(path_context))) {
-    static_cast<QuicForceBlockablePacketWriter*>(connection_->writer())
-        ->ForceWriteBlocked(false);
+    session_->writer()->ForceWriteBlocked(false);
     std::move(callback)(network, MigrationResult::FAILURE);
     if (close_session_on_error) {
       session_->OnConnectionToBeClosedDueToMigrationError(
@@ -514,8 +510,7 @@
                 << MigrationCauseToString(current_migration_cause_);
   // Force blocking the packet writer to avoid any writes since there is no
   // alternate network available.
-  static_cast<QuicForceBlockablePacketWriter*>(connection_->writer())
-      ->ForceWriteBlocked(true);
+  session_->writer()->ForceWriteBlocked(true);
   if (config_.disable_blackhole_detection_on_immediate_migrate) {
     // Turn off the black hole detector since the writer is blocked.
     // Blackhole will be re-enabled once a packet is sent again.
@@ -768,7 +763,7 @@
 
 void QuicConnectionMigrationManager::FinishStartProbing(
     StartProbingCallback probing_callback,
-    std::unique_ptr<QuicPathValidationContext> path_context) {
+    std::unique_ptr<QuicClientPathValidationContext> path_context) {
   session_->PrepareForProbingOnPath(*path_context);
   switch (current_migration_cause_) {
     case MigrationCause::CHANGE_PORT_ON_PATH_DEGRADING:
@@ -836,7 +831,10 @@
     return;
   }
   // Migrate to the probed socket immediately.
-  if (!session_->MigrateToNewPath(std::move(path_context))) {
+  if (!session_->MigrateToNewPath(
+          std::unique_ptr<QuicClientPathValidationContext>(
+              static_cast<QuicClientPathValidationContext*>(
+                  path_context.release())))) {
     if (debug_visitor_) {
       debug_visitor_->OnConnectionMigrationFailedAfterProbe();
     }
@@ -881,7 +879,10 @@
     return;
   }
   // Migrate to the probed socket immediately.
-  if (!session_->MigrateToNewPath(std::move(path_context))) {
+  if (!session_->MigrateToNewPath(
+          std::unique_ptr<QuicClientPathValidationContext>(
+              static_cast<QuicClientPathValidationContext*>(
+                  path_context.release())))) {
     if (debug_visitor_) {
       debug_visitor_->OnConnectionMigrationFailedAfterProbe();
     }
diff --git a/quiche/quic/core/http/quic_connection_migration_manager.h b/quiche/quic/core/http/quic_connection_migration_manager.h
index 243d63d..25ebf12 100644
--- a/quiche/quic/core/http/quic_connection_migration_manager.h
+++ b/quiche/quic/core/http/quic_connection_migration_manager.h
@@ -239,7 +239,7 @@
         bool close_session_on_error, MigrationCallback migration_callback);
 
     void OnCreationSucceeded(
-        std::unique_ptr<QuicPathValidationContext> context) override;
+        std::unique_ptr<QuicClientPathValidationContext> context) override;
 
     void OnCreationFailed(QuicNetworkHandle network,
                           absl::string_view error) override;
@@ -261,7 +261,7 @@
         StartProbingCallback probing_callback);
 
     void OnCreationSucceeded(
-        std::unique_ptr<QuicPathValidationContext> context) override;
+        std::unique_ptr<QuicClientPathValidationContext> context) override;
 
     void OnCreationFailed(QuicNetworkHandle network,
                           absl::string_view error) override;
@@ -294,8 +294,9 @@
                bool close_session_on_error,
                MigrationCallback migration_callback);
   // Helper to finish session migration once the |path_context| is provided.
-  void FinishMigrate(std::unique_ptr<QuicPathValidationContext> path_context,
-                     bool close_session_on_error, MigrationCallback callback);
+  void FinishMigrate(
+      std::unique_ptr<QuicClientPathValidationContext> path_context,
+      bool close_session_on_error, MigrationCallback callback);
 
   void StartMigrateBackToDefaultNetworkTimer(QuicTimeDelta delay);
   void CancelMigrateBackToDefaultNetworkTimer();
@@ -312,7 +313,7 @@
                     const QuicSocketAddress& peer_address);
   void FinishStartProbing(
       StartProbingCallback probing_callback,
-      std::unique_ptr<QuicPathValidationContext> path_context);
+      std::unique_ptr<QuicClientPathValidationContext> path_context);
 
   bool MaybeCloseIdleSession(bool has_write_error,
                              ConnectionCloseBehavior close_behavior);
diff --git a/quiche/quic/core/http/quic_connection_migration_manager_test.cc b/quiche/quic/core/http/quic_connection_migration_manager_test.cc
index aff6e28..f933d5b 100644
--- a/quiche/quic/core/http/quic_connection_migration_manager_test.cc
+++ b/quiche/quic/core/http/quic_connection_migration_manager_test.cc
@@ -112,15 +112,18 @@
   bool force_write_blocked_ = false;
 };
 
-class TestQuicClientPathValidationContext : public QuicPathValidationContext {
+class TestQuicClientPathValidationContext
+    : public QuicClientPathValidationContext {
  public:
   TestQuicClientPathValidationContext(
       const quic::QuicSocketAddress& self_address,
       const quic::QuicSocketAddress& peer_address, QuicNetworkHandle network)
-      : QuicPathValidationContext(self_address, peer_address, network),
+      : QuicClientPathValidationContext(self_address, peer_address, network),
         writer_(std::make_unique<TestQuicForceBlockablePacketWriter>()) {}
 
-  quic::QuicPacketWriter* WriterToUse() override { return writer_.get(); }
+  QuicForceBlockablePacketWriter* ForceBlockableWriterToUse() override {
+    return writer_.get();
+  }
 
   bool ShouldConnectionOwnWriter() const override { return true; }
 
@@ -359,15 +362,16 @@
     : public QuicSpdyClientSessionWithMigration {
  public:
   TestQuicSpdyClientSessionWithMigration(
-      QuicConnection* connection, QuicSession::Visitor* visitor,
-      const QuicConfig& config,
+      QuicConnection* connection, QuicForceBlockablePacketWriter* writer,
+      QuicSession::Visitor* visitor, const QuicConfig& config,
       const ParsedQuicVersionVector& supported_versions,
       QuicNetworkHandle default_network, QuicNetworkHandle current_network,
       std::unique_ptr<QuicPathContextFactory> path_context_factory,
       QuicConnectionMigrationConfig migration_config)
       : QuicSpdyClientSessionWithMigration(
-            connection, visitor, config, supported_versions, default_network,
-            current_network, std::move(path_context_factory), migration_config),
+            connection, writer, visitor, config, supported_versions,
+            default_network, current_network, std::move(path_context_factory),
+            migration_config),
         crypto_stream_(this) {
     ON_CALL(*this, IsSessionProxied()).WillByDefault(Return(false));
     ON_CALL(*this, OnMigrationToPathDone(_, _))
@@ -387,9 +391,10 @@
               (QuicPathValidationContext & context), (override));
   MOCK_METHOD(bool, IsSessionProxied, (), (const));
   MOCK_METHOD(QuicTimeDelta, TimeSinceLastStreamClose, (), (override));
-  MOCK_METHOD(bool, PrepareForMigrationToPath, (QuicPathValidationContext&));
+  MOCK_METHOD(bool, PrepareForMigrationToPath,
+              (QuicClientPathValidationContext&));
   MOCK_METHOD(void, OnMigrationToPathDone,
-              (std::unique_ptr<QuicPathValidationContext>, bool));
+              (std::unique_ptr<QuicClientPathValidationContext>, bool));
   MOCK_METHOD(void, OnConnectionToBeClosedDueToMigrationError,
               (MigrationCause migration_cause, QuicErrorCode quic_error),
               (override));
@@ -497,8 +502,8 @@
     migration_config_.migrate_idle_session = migrate_idle_session_;
 
     session_ = std::make_unique<TestQuicSpdyClientSessionWithMigration>(
-        connection_, &session_visitor_, config_, versions_, default_network_,
-        initial_network_,
+        connection_, default_writer_, &session_visitor_, config_, versions_,
+        default_network_, initial_network_,
         std::unique_ptr<QuicPathContextFactory>(path_context_factory_),
         migration_config_);
     session_->Initialize();
diff --git a/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc b/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
index 7b55176..6e80690 100644
--- a/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
+++ b/quiche/quic/core/http/quic_spdy_client_session_with_migration.cc
@@ -4,11 +4,14 @@
 
 #include "quiche/quic/core/http/quic_spdy_client_session_with_migration.h"
 
+#include "quiche/quic/core/quic_force_blockable_packet_writer.h"
+
 namespace quic {
 
 QuicSpdyClientSessionWithMigration::QuicSpdyClientSessionWithMigration(
-    QuicConnection* connection, QuicSession::Visitor* visitor,
-    const QuicConfig& config, const ParsedQuicVersionVector& supported_versions,
+    QuicConnection* connection, QuicForceBlockablePacketWriter* writer,
+    QuicSession::Visitor* visitor, const QuicConfig& config,
+    const ParsedQuicVersionVector& supported_versions,
     QuicNetworkHandle default_network, QuicNetworkHandle current_network,
     std::unique_ptr<QuicPathContextFactory> path_context_factory,
     QuicConnectionMigrationConfig migration_config)
@@ -17,7 +20,11 @@
       path_context_factory_(std::move(path_context_factory)),
       migration_manager_(this, connection->clock(), default_network,
                          current_network, path_context_factory_.get(),
-                         migration_config) {}
+                         migration_config),
+      writer_(writer) {
+  QUICHE_DCHECK_EQ(writer_, connection->writer())
+      << "Writer is not the connection writer";
+}
 
 QuicSpdyClientSessionWithMigration::~QuicSpdyClientSessionWithMigration() =
     default;
@@ -41,7 +48,7 @@
 }
 
 bool QuicSpdyClientSessionWithMigration::MigrateToNewPath(
-    std::unique_ptr<QuicPathValidationContext> path_context) {
+    std::unique_ptr<QuicClientPathValidationContext> path_context) {
   if (!PrepareForMigrationToPath(*path_context)) {
     QUIC_CLIENT_HISTOGRAM_BOOL("QuicSession.PrepareForMigrationToPath", false,
                                "");
@@ -57,6 +64,8 @@
         "No unused server connection ID");
     QUIC_DVLOG(1) << "MigratePath fails as there is no CID available";
   }
+  writer_ = path_context->ForceBlockableWriterToUse();
+  QUICHE_DCHECK_EQ(writer_, connection()->writer());
   OnMigrationToPathDone(std::move(path_context), success);
   return success;
 }
diff --git a/quiche/quic/core/http/quic_spdy_client_session_with_migration.h b/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
index f4aeab2..6d3346d 100644
--- a/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
+++ b/quiche/quic/core/http/quic_spdy_client_session_with_migration.h
@@ -12,6 +12,7 @@
 #include "quiche/quic/core/quic_config.h"
 #include "quiche/quic/core/quic_connection.h"
 #include "quiche/quic/core/quic_error_codes.h"
+#include "quiche/quic/core/quic_force_blockable_packet_writer.h"
 #include "quiche/quic/core/quic_path_context_factory.h"
 #include "quiche/quic/core/quic_session.h"
 #include "quiche/quic/core/quic_time.h"
@@ -33,8 +34,8 @@
     : public QuicSpdyClientSessionBase {
  public:
   QuicSpdyClientSessionWithMigration(
-      QuicConnection* connection, QuicSession::Visitor* visitor,
-      const QuicConfig& config,
+      QuicConnection* connection, QuicForceBlockablePacketWriter* writer,
+      QuicSession::Visitor* visitor, const QuicConfig& config,
       const ParsedQuicVersionVector& supported_versions,
       QuicNetworkHandle default_network, QuicNetworkHandle current_network,
       std::unique_ptr<QuicPathContextFactory> path_context_factory,
@@ -78,7 +79,7 @@
   // network.
   // Returns true on successful migration.
   bool MigrateToNewPath(
-      std::unique_ptr<QuicPathValidationContext> path_context);
+      std::unique_ptr<QuicClientPathValidationContext> path_context);
 
   void SetMigrationDebugVisitor(
       quic::QuicConnectionMigrationDebugVisitor* visitor);
@@ -93,19 +94,23 @@
     return migration_manager_;
   }
 
+  QuicForceBlockablePacketWriter* writer() { return writer_; }
+
  protected:
   // Called in MigrateToNewPath() prior to calling MigratePath().
   // Return false if MigratePath() should be skipped.
   virtual bool PrepareForMigrationToPath(
-      QuicPathValidationContext& context) = 0;
+      QuicClientPathValidationContext& context) = 0;
 
   // Called in MigrateToNewPath() after MigratePath() for clean up.
   virtual void OnMigrationToPathDone(
-      std::unique_ptr<QuicPathValidationContext> context, bool success) = 0;
+      std::unique_ptr<QuicClientPathValidationContext> context,
+      bool success) = 0;
 
  private:
   std::unique_ptr<QuicPathContextFactory> path_context_factory_;
   QuicConnectionMigrationManager migration_manager_;
+  QuicForceBlockablePacketWriter* absl_nonnull writer_;
 };
 
 }  // namespace quic
diff --git a/quiche/quic/core/quic_path_context_factory.h b/quiche/quic/core/quic_path_context_factory.h
index 0134587..5fa5c2d 100644
--- a/quiche/quic/core/quic_path_context_factory.h
+++ b/quiche/quic/core/quic_path_context_factory.h
@@ -8,14 +8,40 @@
 #include <memory>
 
 #include "absl/strings/string_view.h"
+#include "quiche/quic/core/quic_force_blockable_packet_writer.h"
 #include "quiche/quic/core/quic_path_validator.h"
 #include "quiche/quic/platform/api/quic_socket_address.h"
 #include "quiche/common/platform/api/quiche_export.h"
 
 namespace quic {
 
-// An interface for creating QuicPathValidationContext objects used for probing
-// and migrating path.
+// A client side interface that enforces the writer to be force blockable.
+class QUICHE_EXPORT QuicClientPathValidationContext
+    : public QuicPathValidationContext {
+ public:
+  QuicClientPathValidationContext(const QuicSocketAddress& self_address,
+                                  const QuicSocketAddress& peer_address,
+                                  QuicNetworkHandle network)
+      : QuicClientPathValidationContext(self_address, peer_address,
+                                        peer_address, network) {}
+
+  QuicClientPathValidationContext(
+      const QuicSocketAddress& self_address,
+      const QuicSocketAddress& peer_address,
+      const QuicSocketAddress& effective_peer_address,
+      QuicNetworkHandle network)
+      : QuicPathValidationContext(self_address, peer_address,
+                                  effective_peer_address, network) {}
+
+  QuicPacketWriter* WriterToUse() override {
+    return ForceBlockableWriterToUse();
+  }
+
+  virtual QuicForceBlockablePacketWriter* ForceBlockableWriterToUse() = 0;
+};
+
+// An interface for creating QuicClientPathValidationContext objects used for
+// probing and migrating path.
 class QUICHE_EXPORT QuicPathContextFactory {
  public:
   // An interface to handle creation success and failure given the creation
@@ -27,7 +53,7 @@
     // Called when the manager successfully created a path context.
     // |context| must not be nullptr.
     virtual void OnCreationSucceeded(
-        std::unique_ptr<QuicPathValidationContext> context) = 0;
+        std::unique_ptr<QuicClientPathValidationContext> context) = 0;
 
     // Called when the manager fails to create a path context on |network|.
     // |error| will be populated with details.