Allow SCONE packets to update the server connection ID.

Changes protected by QuicConnection::parse_scone_packets_, which is never true in production.

PiperOrigin-RevId: 968604914
diff --git a/quiche/quic/core/http/end_to_end_test.cc b/quiche/quic/core/http/end_to_end_test.cc
index 0b0a62a..5eb3158 100644
--- a/quiche/quic/core/http/end_to_end_test.cc
+++ b/quiche/quic/core/http/end_to_end_test.cc
@@ -32,6 +32,7 @@
 #include "openssl/ssl.h"
 #include "quiche/quic/core/congestion_control/rtt_stats.h"
 #include "quiche/quic/core/crypto/crypto_protocol.h"
+#include "quiche/quic/core/crypto/proof_source.h"
 #include "quiche/quic/core/crypto/quic_client_session_cache.h"
 #include "quiche/quic/core/crypto/quic_compressed_certs_cache.h"
 #include "quiche/quic/core/crypto/quic_crypto_client_config.h"
@@ -84,6 +85,7 @@
 #include "quiche/quic/core/quic_utils.h"
 #include "quiche/quic/core/quic_versions.h"
 #include "quiche/quic/core/scone.h"
+#include "quiche/quic/core/tls_server_handshaker.h"
 #include "quiche/quic/core/web_transport_interface.h"
 #include "quiche/quic/platform/api/quic_expect_bug.h"
 #include "quiche/quic/platform/api/quic_flags.h"
@@ -94,6 +96,7 @@
 #include "quiche/quic/platform/api/quic_test_loopback.h"
 #include "quiche/quic/test_tools/bad_packet_writer.h"
 #include "quiche/quic/test_tools/crypto_test_utils.h"
+#include "quiche/quic/test_tools/fake_proof_source_handle.h"
 #include "quiche/quic/test_tools/packet_dropping_test_writer.h"
 #include "quiche/quic/test_tools/packet_reordering_writer.h"
 #include "quiche/quic/test_tools/qpack/qpack_encoder_peer.h"
@@ -130,7 +133,6 @@
 #include "quiche/common/platform/api/quiche_logging.h"
 #include "quiche/common/platform/api/quiche_reference_counted.h"
 #include "quiche/common/platform/api/quiche_test.h"
-#include "quiche/common/quiche_endian.h"
 #include "quiche/common/quiche_mem_slice.h"
 #include "quiche/common/simple_buffer_allocator.h"
 #include "quiche/common/test_tools/quiche_test_utils.h"
@@ -282,6 +284,59 @@
   QuicDefaultClient* client_;
 };
 
+// In production, the proof source is often asynchronous, which can affect the
+// code path. These classes create this pattern by returning PENDING on
+// SelectCertificate() and scheduling completion on |server_thread|.
+class AsyncTlsServerHandshaker : public TlsServerHandshaker {
+ public:
+  // Note that |server_thread| is a pointer to a unique_ptr. server_thread_
+  // has typically not been created when this is called.
+  AsyncTlsServerHandshaker(QuicSession* session,
+                           const QuicCryptoServerConfig* crypto_config,
+                           std::unique_ptr<ServerThread>* server_thread)
+      : TlsServerHandshaker(session, crypto_config),
+        crypto_config_(crypto_config),
+        server_thread_(server_thread) {}
+
+  std::unique_ptr<ProofSourceHandle> MaybeCreateProofSourceHandle() override {
+    auto handle = std::make_unique<FakeProofSourceHandle>(
+        crypto_config_->proof_source(), this,
+        FakeProofSourceHandle::Action::DELEGATE_ASYNC,
+        FakeProofSourceHandle::Action::DELEGATE_SYNC, QuicDelayedSSLConfig(),
+        [this]() {
+          (*server_thread_)->Schedule([this]() {
+            fake_proof_source_handle_->CompletePendingOperation();
+          });
+        });
+    fake_proof_source_handle_ = handle.get();
+    return handle;
+  }
+
+ private:
+  const QuicCryptoServerConfig* crypto_config_;
+  std::unique_ptr<ServerThread>* server_thread_;
+  FakeProofSourceHandle* fake_proof_source_handle_ = nullptr;
+};
+
+class AsyncCryptoStreamFactory : public QuicTestServer::CryptoStreamFactory {
+ public:
+  // Note that |server_thread| is a pointer to a unique_ptr. server_thread_
+  // has typically not been created when this is called.
+  explicit AsyncCryptoStreamFactory(
+      std::unique_ptr<ServerThread>* server_thread)
+      : server_thread_(server_thread) {}
+
+  std::unique_ptr<QuicCryptoServerStreamBase> CreateCryptoStream(
+      const QuicCryptoServerConfig* crypto_config,
+      QuicServerSessionBase* session) override {
+    return std::make_unique<AsyncTlsServerHandshaker>(session, crypto_config,
+                                                      server_thread_);
+  }
+
+ private:
+  std::unique_ptr<ServerThread>* server_thread_;
+};
+
 class EndToEndTest : public QuicTestWithParam<TestParams> {
  protected:
   EndToEndTest()
@@ -647,8 +702,8 @@
       fd_ = kQuicInvalidSocketFd;
     }
     auto test_server = std::make_unique<QuicTestServer>(
-        crypto_test_utils::ProofSourceForTesting(), server_config_,
-        server_supported_versions_, &memory_cache_backend_,
+        crypto_test_utils::ProofSourceForTesting("", num_certs_in_chain_),
+        server_config_, server_supported_versions_, &memory_cache_backend_,
         expected_server_connection_id_length_);
     test_server->SetEventLoopFactory(GetParam().event_loop);
     const QuicSocketAddress server_listening_address =
@@ -677,6 +732,10 @@
       absl::down_cast<QuicTestServer*>(server_thread_->server())
           ->SetSpdyStreamFactory(stream_factory_);
     }
+    if (async_crypto_stream_factory_ != nullptr) {
+      absl::down_cast<QuicTestServer*>(server_thread_->server())
+          ->SetCryptoStreamFactory(async_crypto_stream_factory_);
+    }
 
     server_thread_->Start();
   }
@@ -1119,6 +1178,8 @@
   bool use_preferred_address_ = false;
   QuicSocketAddress server_preferred_address_;
   QuicPacketWriterParams packet_writer_params_;
+  AsyncCryptoStreamFactory* async_crypto_stream_factory_ = nullptr;
+  int num_certs_in_chain_ = 1;
 };
 
 // Run all end to end tests with all supported versions.
@@ -4763,6 +4824,8 @@
 // SconePacketWriter will look for SCONE packets from the client and write
 // bandwidth values into them as a SCONE network element would. It only does
 // so for flows where it has observed the SCONE indicator.
+// TODO(martinduke): Refactor this class. It only works well when the SCONE
+// flow is client->server.
 static constexpr size_t kMaxSconeReports = 5;
 class SconePacketWriter : public PacketDroppingTestWriter {
  public:
@@ -4803,6 +4866,12 @@
     return kReportedValues[*last_report_index_].second;
   }
 
+  // For server->client SCONE flows, the writer won't see the SCONE indicator.
+  void ForceFlowtoScone(QuicSocketAddress address) {
+    observed_scone_endpoints_.insert(address.host());
+  }
+  bool SawSconeIndicator() const { return !observed_scone_endpoints_.empty(); }
+
  private:
   bool FlowIsScone(const QuicIpAddress& self_address,
                    const QuicSocketAddress& peer_address, const char* buffer) {
@@ -4970,6 +5039,48 @@
   }
 }
 
+// Repro for part of b/548012868. It will crash if the client does not accept
+// a new connection ID in a SCONE header bundled with INITIAL.
+TEST_P(EndToEndTest, SconeProtocolServerToClientAsynchronous) {
+  if (!version_.IsIetfQuic() || override_server_connection_id_length_ != 16) {
+    // Because the server in this test suite uses
+    // DeterministicConnectionIdGenerator, an 8-byte connection ID is unchanged,
+    // meaning it will not be rejected by the client.
+    ASSERT_TRUE(Initialize());
+    return;
+  }
+  client_config_.set_parse_scone_packets(true);
+  server_config_.set_scone_packet_interval(QuicTime::Delta::FromSeconds(20));
+
+  // Build a cert chain with 8 certs, so that the HANDSHAKE messages fill
+  // several packets.
+  num_certs_in_chain_ = 8;
+  AsyncCryptoStreamFactory async_crypto_stream_factory(&server_thread_);
+  async_crypto_stream_factory_ = &async_crypto_stream_factory;
+
+  // server_writer_ is allocated with 'new' in SetUp(), but will be replaced
+  // here with a SconePacketWriter.
+  delete server_writer_;
+  server_writer_ = new SconePacketWriter();
+  absl::down_cast<SconePacketWriter*>(server_writer_)
+      ->ForceFlowtoScone(server_address_);
+  delete client_writer_;
+  client_writer_ = new SconePacketWriter();
+
+  ASSERT_TRUE(Initialize());
+  EXPECT_TRUE(client_->client()->WaitForOneRttKeysAvailable());
+  QuicTestClientSession* client_session =
+      absl::down_cast<QuicTestClientSession*>(client_->client()->session());
+  ASSERT_NE(client_session, nullptr);
+
+  EXPECT_TRUE(
+      absl::down_cast<SconePacketWriter*>(client_writer_)->SawSconeIndicator());
+  // The very first packet in the handshake came with SCONE. If it was rejected
+  // by the client because it was a new connection ID, there will be no
+  // bandwidth report.
+  EXPECT_NE(client_session->received_bandwidth(), QuicBandwidth::Zero());
+}
+
 TEST_P(EndToEndTest, VersionNegotiationDowngradeAttackIsDetected) {
   ResetClientWriterForVersionNegotiationTest();
   ParsedQuicVersion target_version = server_supported_versions_.back();
diff --git a/quiche/quic/core/quic_connection.cc b/quiche/quic/core/quic_connection.cc
index eedbef6..6f2240c 100644
--- a/quiche/quic/core/quic_connection.cc
+++ b/quiche/quic/core/quic_connection.cc
@@ -1005,6 +1005,19 @@
     return true;
   }
 
+  // header.is_scone_header cannot be true unless parse_scone_packets is true.
+  if (header.is_scone_header && perspective_ == Perspective::IS_CLIENT &&
+      version().IsIetfQuic() &&
+
+      (!server_connection_id_replaced_by_initial_ ||
+       // The SCONE sender can omit server connection ID if followed by short
+       // header packets.
+       server_connection_id.IsEmpty())) {
+    QUIC_BUG_IF(quic_bug_scone_parsed_unexpectedly, !parse_scone_packets_)
+        << "header.is_scone_header true without setting parse_scone_packets";
+    return true;
+  }
+
   if (PacketCanReplaceServerConnectionId(header, perspective_)) {
     QUIC_DLOG(INFO) << ENDPOINT << "Accepting packet with new connection ID "
                     << server_connection_id << " instead of "
diff --git a/quiche/quic/core/quic_connection_test.cc b/quiche/quic/core/quic_connection_test.cc
index acb6de8..bfea0bc 100644
--- a/quiche/quic/core/quic_connection_test.cc
+++ b/quiche/quic/core/quic_connection_test.cc
@@ -13,7 +13,6 @@
 #include <utility>
 #include <vector>
 
-#include "absl/base/macros.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
@@ -19047,6 +19046,49 @@
   }
 }
 
+TEST_P(QuicConnectionTest, SconeCanChangeServerConnectionId) {
+  if (!version().IsIetfQuic()) {
+    return;
+  }
+  // Set up connection to accept SCONE packets to avoid triggering QUIC_BUG.
+  QuicConfig config;
+  config.set_parse_scone_packets(true);
+  EXPECT_CALL(*send_algorithm_, SetFromConfig);
+  EXPECT_CALL(*send_algorithm_, EnableECT1()).WillRepeatedly(Return(false));
+  EXPECT_CALL(*send_algorithm_, EnableECT0()).WillRepeatedly(Return(false));
+  connection_.SetFromConfig(config);
+
+  // Create a valid SCONE header that meets all the conditions.
+  QuicPacketHeader header;
+  header.destination_connection_id = EmptyQuicConnectionId();
+  header.source_connection_id = EmptyQuicConnectionId();
+  header.version = UnsupportedQuicVersion();
+  header.form = IETF_QUIC_LONG_HEADER_PACKET;
+  header.long_packet_type = HANDSHAKE;
+  header.version_flag = true;
+  header.is_scone_header = true;
+  EXPECT_TRUE(connection_.OnUnauthenticatedPublicHeader(header));
+
+  // Break the conditions one at a time.
+  header.is_scone_header = false;
+  EXPECT_FALSE(connection_.OnUnauthenticatedPublicHeader(header));
+  header.is_scone_header = true;
+
+  header.source_connection_id = TestConnectionId(0x1234);
+  EXPECT_TRUE(connection_.OnUnauthenticatedPublicHeader(header));
+  // Replace initial server connection ID. Non-empty connection ID is no longer
+  // valid.
+  EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
+  peer_creator_.SetServerConnectionId(TestConnectionId(0x5678));
+  QuicFrame frame = MakeCryptoFrame();
+  ForceProcessFramePacket(frame);
+  DeleteFrame(&frame);
+  EXPECT_EQ(
+      QuicConnectionPeer::GetDefaultPath(&connection_)->server_connection_id,
+      TestConnectionId(0x5678));
+  EXPECT_FALSE(connection_.OnUnauthenticatedPublicHeader(header));
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quiche/quic/core/quic_framer.cc b/quiche/quic/core/quic_framer.cc
index 785b29b..5884abe 100644
--- a/quiche/quic/core/quic_framer.cc
+++ b/quiche/quic/core/quic_framer.cc
@@ -2464,6 +2464,7 @@
       if (version_label == kSconeVersionHigh) {
         *scone_value |= 0x01;
       }
+      header->is_scone_header = true;
     }
     header->destination_connection_id =
         QuicConnectionId(destination_connection_id);
diff --git a/quiche/quic/core/quic_packets.h b/quiche/quic/core/quic_packets.h
index ee1c27d..29a717d 100644
--- a/quiche/quic/core/quic_packets.h
+++ b/quiche/quic/core/quic_packets.h
@@ -145,6 +145,8 @@
   bool has_possible_stateless_reset_token : 1;
   // Latency spin bit on the short packet header (RFC 9000 Section 17.4)
   bool spin_bit : 1;
+  // Packet is SCONE (version will be ParsedQuicVersion::Unsupported)
+  bool is_scone_header : 1 = false;
   // -- end bitfield-able bools in the first cacheline --
 
   // There are 8 bytes still available in the first cacheline.  Start with long
diff --git a/quiche/quic/test_tools/crypto_test_utils.cc b/quiche/quic/test_tools/crypto_test_utils.cc
index b7e7e85..abf3419 100644
--- a/quiche/quic/test_tools/crypto_test_utils.cc
+++ b/quiche/quic/test_tools/crypto_test_utils.cc
@@ -901,11 +901,13 @@
 
 class TestProofSource : public ProofSourceX509 {
  public:
-  explicit TestProofSource(const std::string& trust_anchor_id)
+  explicit TestProofSource(const std::string& trust_anchor_id,
+                           int num_certs_in_chain = 1)
       : ProofSourceX509(
             quiche::QuicheReferenceCountedPointer<ProofSource::Chain>(
                 new ProofSource::Chain(
-                    std::vector<std::string>{std::string(kTestCertificate)},
+                    std::vector<std::string>(num_certs_in_chain,
+                                             std::string(kTestCertificate)),
                     trust_anchor_id)),
             std::move(*CertificatePrivateKey::LoadFromDer(
                 kTestCertificatePrivateKey))) {
@@ -991,8 +993,8 @@
 }  // namespace
 
 std::unique_ptr<ProofSource> ProofSourceForTesting(
-    const std::string& trust_anchor_id) {
-  return std::make_unique<TestProofSource>(trust_anchor_id);
+    const std::string& trust_anchor_id, int num_certs_in_chain) {
+  return std::make_unique<TestProofSource>(trust_anchor_id, num_certs_in_chain);
 }
 
 std::unique_ptr<ProofVerifier> ProofVerifierForTesting() {
diff --git a/quiche/quic/test_tools/crypto_test_utils.h b/quiche/quic/test_tools/crypto_test_utils.h
index 79fc473..1626251 100644
--- a/quiche/quic/test_tools/crypto_test_utils.h
+++ b/quiche/quic/test_tools/crypto_test_utils.h
@@ -9,6 +9,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -176,7 +177,7 @@
 // with it as described at
 // https://tlswg.org/tls-trust-anchor-ids/draft-ietf-tls-trust-anchor-ids.html#section-4.1.
 std::unique_ptr<ProofSource> ProofSourceForTesting(
-    const std::string& trust_anchor_id = "");
+    const std::string& trust_anchor_id = "", int num_certs_in_chain = 1);
 
 // Returns a new |ProofVerifier| that uses the QUIC testing root CA.
 std::unique_ptr<ProofVerifier> ProofVerifierForTesting();
diff --git a/quiche/quic/test_tools/fake_proof_source_handle.cc b/quiche/quic/test_tools/fake_proof_source_handle.cc
index 56cee5c..352eacb 100644
--- a/quiche/quic/test_tools/fake_proof_source_handle.cc
+++ b/quiche/quic/test_tools/fake_proof_source_handle.cc
@@ -20,10 +20,9 @@
 #include "quiche/quic/core/quic_connection_id.h"
 #include "quiche/quic/core/quic_types.h"
 #include "quiche/quic/platform/api/quic_bug_tracker.h"
-#include "quiche/quic/platform/api/quic_flags.h"
 #include "quiche/quic/platform/api/quic_socket_address.h"
 #include "quiche/common/platform/api/quiche_logging.h"
-#include "quiche/common/platform/api/quiche_reference_counted.h"
+#include "quiche/common/quiche_callbacks.h"
 
 namespace quic {
 namespace test {
@@ -70,12 +69,14 @@
 FakeProofSourceHandle::FakeProofSourceHandle(
     ProofSource* absl_nonnull delegate,
     ProofSourceHandleCallback* absl_nonnull callback, Action select_cert_action,
-    Action compute_signature_action, QuicDelayedSSLConfig delayed_ssl_config)
+    Action compute_signature_action, QuicDelayedSSLConfig delayed_ssl_config,
+    quiche::SingleUseCallback<void()> pending_callback)
     : delegate_(delegate),
       callback_(callback),
       select_cert_action_(select_cert_action),
       compute_signature_action_(compute_signature_action),
-      delayed_ssl_config_(delayed_ssl_config) {
+      delayed_ssl_config_(delayed_ssl_config),
+      pending_callback_(std::move(pending_callback)) {
   QUICHE_CHECK(delegate);
   QUICHE_CHECK(callback);
 }
@@ -108,6 +109,10 @@
       select_cert_action_ == Action::FAIL_ASYNC) {
     select_cert_op_.emplace(delegate_, callback_, select_cert_action_,
                             all_select_cert_args_.back(), delayed_ssl_config_);
+    if (pending_callback_ != nullptr) {
+      std::move(pending_callback_)();
+      pending_callback_ = nullptr;
+    }
     return QUIC_PENDING;
   } else if (select_cert_action_ == Action::FAIL_SYNC ||
              select_cert_action_ == Action::FAIL_SYNC_DO_NOT_CHECK_CLOSED) {
diff --git a/quiche/quic/test_tools/fake_proof_source_handle.h b/quiche/quic/test_tools/fake_proof_source_handle.h
index 3347482..fe64d6e 100644
--- a/quiche/quic/test_tools/fake_proof_source_handle.h
+++ b/quiche/quic/test_tools/fake_proof_source_handle.h
@@ -18,6 +18,7 @@
 #include "quiche/quic/core/quic_connection_id.h"
 #include "quiche/quic/core/quic_types.h"
 #include "quiche/quic/platform/api/quic_socket_address.h"
+#include "quiche/common/quiche_callbacks.h"
 
 namespace quic {
 namespace test {
@@ -43,11 +44,14 @@
 
   // |delegate| must do cert selection and signature synchronously.
   // |delayed_ssl_config| is the config passed to OnSelectCertificateDone.
+  // |pending_callback| is called when a pending operation is started, so that
+  // CompletePendingOperation() can be scheduled on a thread.
   FakeProofSourceHandle(
       ProofSource* absl_nonnull delegate,
       ProofSourceHandleCallback* absl_nonnull callback,
       Action select_cert_action, Action compute_signature_action,
-      QuicDelayedSSLConfig delayed_ssl_config = QuicDelayedSSLConfig());
+      QuicDelayedSSLConfig delayed_ssl_config = QuicDelayedSSLConfig(),
+      quiche::SingleUseCallback<void()> pending_callback = nullptr);
 
   ~FakeProofSourceHandle() override = default;
 
@@ -203,6 +207,9 @@
   // Save all the select cert and compute signature args for tests to inspect.
   std::vector<SelectCertArgs> all_select_cert_args_;
   std::vector<ComputeSignatureArgs> all_compute_signature_args_;
+
+  // Called when SelectCertificate() returns PENDING.
+  quiche::SingleUseCallback<void()> pending_callback_;
 };
 
 }  // namespace test