In TlsServerHandshaker, do not call ProofSourceHandle::SelectCertificate if QUIC connection has disconnected.

Protected by FLAGS_quic_reloadable_flag_quic_tls_no_select_cert_if_disconnected.

PiperOrigin-RevId: 410354467
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index af9c02c..176f87c 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -65,6 +65,8 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_discard_initial_packet_with_key_dropped, true)
 // If true, do not bundle 2nd ACK with connection close if there is an ACK queued.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_single_ack_in_packet2, false)
+// If true, do not call ProofSourceHandle::SelectCertificate if QUIC connection has disconnected.
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_tls_no_select_cert_if_disconnected, true)
 // If true, do not count bytes sent/received on the alternative path into the bytes sent/received on the default path.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_count_bytes_on_alternative_path_seperately, true)
 // If true, do not re-arm PTO while sending application data during handshake.
diff --git a/quic/core/tls_server_handshaker.cc b/quic/core/tls_server_handshaker.cc
index d35b923..15eed38 100644
--- a/quic/core/tls_server_handshaker.cc
+++ b/quic/core/tls_server_handshaker.cc
@@ -210,6 +210,10 @@
   if (session->connection()->context()->tracer) {
     tls_connection_.EnableInfoCallback();
   }
+
+  if (no_select_cert_if_disconnected_) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_tls_no_select_cert_if_disconnected, 1, 2);
+  }
 }
 
 TlsServerHandshaker::~TlsServerHandshaker() { CancelOutstandingCallbacks(); }
@@ -965,6 +969,13 @@
           ? std::string(alps_result.alps_buffer.get(), alps_result.alps_length)
           : std::string();
 
+  if (no_select_cert_if_disconnected_ &&
+      !session()->connection()->connected()) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_tls_no_select_cert_if_disconnected, 2, 2);
+    select_cert_status_ = QUIC_FAILURE;
+    return ssl_select_cert_error;
+  }
+
   const QuicAsyncStatus status = proof_source_handle_->SelectCertificate(
       session()->connection()->self_address().Normalized(),
       session()->connection()->peer_address().Normalized(),
diff --git a/quic/core/tls_server_handshaker.h b/quic/core/tls_server_handshaker.h
index 57aa652..01cf19b 100644
--- a/quic/core/tls_server_handshaker.h
+++ b/quic/core/tls_server_handshaker.h
@@ -377,6 +377,8 @@
       last_received_cached_network_params_;
 
   bool cert_matched_sni_ = false;
+  const bool no_select_cert_if_disconnected_ =
+      GetQuicReloadableFlag(quic_tls_no_select_cert_if_disconnected);
 };
 
 }  // namespace quic
diff --git a/quic/core/tls_server_handshaker_test.cc b/quic/core/tls_server_handshaker_test.cc
index 3e707b2..214866c 100644
--- a/quic/core/tls_server_handshaker_test.cc
+++ b/quic/core/tls_server_handshaker_test.cc
@@ -26,6 +26,7 @@
 #include "quic/test_tools/failing_proof_source.h"
 #include "quic/test_tools/fake_proof_source.h"
 #include "quic/test_tools/fake_proof_source_handle.h"
+#include "quic/test_tools/quic_config_peer.h"
 #include "quic/test_tools/quic_test_utils.h"
 #include "quic/test_tools/simple_session_cache.h"
 #include "quic/test_tools/test_certificates.h"
@@ -95,6 +96,10 @@
     ON_CALL(*this, MaybeCreateProofSourceHandle())
         .WillByDefault(testing::Invoke(
             this, &TestTlsServerHandshaker::RealMaybeCreateProofSourceHandle));
+
+    ON_CALL(*this, OverrideQuicConfigDefaults(_))
+        .WillByDefault(testing::Invoke(
+            this, &TestTlsServerHandshaker::RealOverrideQuicConfigDefaults));
   }
 
   MOCK_METHOD(std::unique_ptr<ProofSourceHandle>,
@@ -102,6 +107,9 @@
               (),
               (override));
 
+  MOCK_METHOD(void, OverrideQuicConfigDefaults, (QuicConfig * config),
+              (override));
+
   void SetupProofSourceHandle(
       FakeProofSourceHandle::Action select_cert_action,
       FakeProofSourceHandle::Action compute_signature_action,
@@ -142,6 +150,10 @@
     return TlsServerHandshaker::MaybeCreateProofSourceHandle();
   }
 
+  void RealOverrideQuicConfigDefaults(QuicConfig* config) {
+    return TlsServerHandshaker::OverrideQuicConfigDefaults(config);
+  }
+
   // Owned by TlsServerHandshaker.
   FakeProofSourceHandle* fake_proof_source_handle_ = nullptr;
   ProofSource* proof_source_ = nullptr;
@@ -1023,6 +1035,44 @@
   EXPECT_FALSE(server_handshaker_->received_client_cert());
 }
 
+TEST_P(TlsServerHandshakerTest, CloseConnectionBeforeSelectCert) {
+  InitializeServerWithFakeProofSourceHandle();
+  server_handshaker_->SetupProofSourceHandle(
+      /*select_cert_action=*/FakeProofSourceHandle::Action::
+          FAIL_SYNC_DO_NOT_CHECK_CLOSED,
+      /*compute_signature_action=*/FakeProofSourceHandle::Action::
+          FAIL_SYNC_DO_NOT_CHECK_CLOSED);
+
+  EXPECT_CALL(*server_handshaker_, OverrideQuicConfigDefaults(_))
+      .WillOnce(testing::Invoke([](QuicConfig* config) {
+        QuicConfigPeer::SetReceivedMaxUnidirectionalStreams(config,
+                                                            /*max_streams=*/0);
+      }));
+
+  EXPECT_CALL(*server_connection_,
+              CloseConnection(QUIC_ZERO_RTT_RESUMPTION_LIMIT_REDUCED, _, _))
+      .WillOnce(testing::Invoke(
+          [this](QuicErrorCode error, const std::string& details,
+                 ConnectionCloseBehavior connection_close_behavior) {
+            server_connection_->ReallyCloseConnection(
+                error, details, connection_close_behavior);
+            ASSERT_FALSE(server_connection_->connected());
+          }));
+
+  AdvanceHandshakeWithFakeClient();
+  if (!GetQuicReloadableFlag(quic_tls_no_select_cert_if_disconnected)) {
+    // SelectCertificate is called when flag is false.
+    EXPECT_FALSE(server_handshaker_->fake_proof_source_handle()
+                     ->all_select_cert_args()
+                     .empty());
+    return;
+  }
+
+  EXPECT_TRUE(server_handshaker_->fake_proof_source_handle()
+                  ->all_select_cert_args()
+                  .empty());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/test_tools/fake_proof_source_handle.cc b/quic/test_tools/fake_proof_source_handle.cc
index 46a1d09..0b2e06f 100644
--- a/quic/test_tools/fake_proof_source_handle.cc
+++ b/quic/test_tools/fake_proof_source_handle.cc
@@ -80,7 +80,9 @@
     const std::vector<uint8_t>& quic_transport_params,
     const absl::optional<std::vector<uint8_t>>& early_data_context,
     const QuicSSLConfig& ssl_config) {
-  QUICHE_CHECK(!closed_);
+  if (select_cert_action_ != Action::FAIL_SYNC_DO_NOT_CHECK_CLOSED) {
+    QUICHE_CHECK(!closed_);
+  }
   all_select_cert_args_.push_back(SelectCertArgs(
       server_address, client_address, ssl_capabilities, hostname, client_hello,
       alpn, alps, quic_transport_params, early_data_context, ssl_config));
@@ -90,7 +92,8 @@
     select_cert_op_.emplace(delegate_, callback_, select_cert_action_,
                             all_select_cert_args_.back(), dealyed_ssl_config_);
     return QUIC_PENDING;
-  } else if (select_cert_action_ == Action::FAIL_SYNC) {
+  } else if (select_cert_action_ == Action::FAIL_SYNC ||
+             select_cert_action_ == Action::FAIL_SYNC_DO_NOT_CHECK_CLOSED) {
     callback()->OnSelectCertificateDone(
         /*ok=*/false,
         /*is_sync=*/true, nullptr, /*handshake_hints=*/absl::string_view(),
@@ -121,7 +124,9 @@
     uint16_t signature_algorithm,
     absl::string_view in,
     size_t max_signature_size) {
-  QUICHE_CHECK(!closed_);
+  if (compute_signature_action_ != Action::FAIL_SYNC_DO_NOT_CHECK_CLOSED) {
+    QUICHE_CHECK(!closed_);
+  }
   all_compute_signature_args_.push_back(
       ComputeSignatureArgs(server_address, client_address, hostname,
                            signature_algorithm, in, max_signature_size));
@@ -132,7 +137,9 @@
                                   compute_signature_action_,
                                   all_compute_signature_args_.back());
     return QUIC_PENDING;
-  } else if (compute_signature_action_ == Action::FAIL_SYNC) {
+  } else if (compute_signature_action_ == Action::FAIL_SYNC ||
+             compute_signature_action_ ==
+                 Action::FAIL_SYNC_DO_NOT_CHECK_CLOSED) {
     callback()->OnComputeSignatureDone(/*ok=*/false, /*is_sync=*/true,
                                        /*signature=*/"", /*details=*/nullptr);
     return QUIC_FAILURE;
diff --git a/quic/test_tools/fake_proof_source_handle.h b/quic/test_tools/fake_proof_source_handle.h
index b4203e8..bb2c60e 100644
--- a/quic/test_tools/fake_proof_source_handle.h
+++ b/quic/test_tools/fake_proof_source_handle.h
@@ -25,6 +25,8 @@
     // Handle the operation asynchronously. Fail the operation when the caller
     // calls CompletePendingOperation().
     FAIL_ASYNC,
+    // Similar to FAIL_SYNC, but do not QUICHE_CHECK(!closed_) when invoked.
+    FAIL_SYNC_DO_NOT_CHECK_CLOSED,
   };
 
   // |delegate| must do cert selection and signature synchronously.