Add a ConnectionIdLength() interface to ConnectionIdGeneratorInterface. As the connection ID length can change dynamically, QuicDispatcher will leverage this interface to know how many bytes to extract from short headers. PiperOrigin-RevId: 480435323
diff --git a/quiche/quic/core/connection_id_generator.h b/quiche/quic/core/connection_id_generator.h index 65284d8..d27b44f 100644 --- a/quiche/quic/core/connection_id_generator.h +++ b/quiche/quic/core/connection_id_generator.h
@@ -23,6 +23,9 @@ // and consider replacing it. Returns empty if not replaced. virtual absl::optional<QuicConnectionId> MaybeReplaceConnectionId( const QuicConnectionId& original, const ParsedQuicVersion& version) = 0; + // Returns the length of a connection ID generated by this generator with the + // specified first byte. + virtual uint8_t ConnectionIdLength(uint8_t first_byte) const = 0; virtual ~ConnectionIdGeneratorInterface() = default; };
diff --git a/quiche/quic/core/deterministic_connection_id_generator.h b/quiche/quic/core/deterministic_connection_id_generator.h index cd27811..74d42d9 100644 --- a/quiche/quic/core/deterministic_connection_id_generator.h +++ b/quiche/quic/core/deterministic_connection_id_generator.h
@@ -27,6 +27,9 @@ absl::optional<QuicConnectionId> MaybeReplaceConnectionId( const QuicConnectionId& original, const ParsedQuicVersion& version) override; + uint8_t ConnectionIdLength(uint8_t /*first_byte*/) const override { + return expected_connection_id_length_; + } private: const uint8_t expected_connection_id_length_;
diff --git a/quiche/quic/core/deterministic_connection_id_generator_test.cc b/quiche/quic/core/deterministic_connection_id_generator_test.cc index 67016b5..6c3ee21 100644 --- a/quiche/quic/core/deterministic_connection_id_generator_test.cc +++ b/quiche/quic/core/deterministic_connection_id_generator_test.cc
@@ -118,5 +118,9 @@ } } +TEST_P(DeterministicConnectionIdGeneratorTest, ReturnLength) { + EXPECT_EQ(generator_.ConnectionIdLength(0x01), connection_id_length_); +} + } // namespace test } // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.cc b/quiche/quic/load_balancer/load_balancer_encoder.cc index 914b890..9e71d6c 100644 --- a/quiche/quic/load_balancer/load_balancer_encoder.cc +++ b/quiche/quic/load_balancer/load_balancer_encoder.cc
@@ -44,9 +44,6 @@ } // namespace -constexpr uint8_t kLoadBalancerLengthMask = 0x3f; -constexpr uint8_t kLoadBalancerUnroutableConfigId = 0xc0; - absl::optional<LoadBalancerEncoder> LoadBalancerEncoder::Create( QuicRandom &random, LoadBalancerEncoderVisitorInterface *const visitor, const bool len_self_encoded, const uint8_t unroutable_connection_id_len) { @@ -89,6 +86,7 @@ seed_ = absl::MakeUint128(random_.RandUint64(), random_.RandUint64()) % NumberOfNonces(config.nonce_len()); num_nonces_left_ = NumberOfNonces(config.nonce_len()); + connection_id_lengths_[config.config_id()] = config.total_len(); return true; } @@ -102,10 +100,10 @@ } QuicConnectionId LoadBalancerEncoder::GenerateConnectionId() { - uint8_t length = (config_.has_value()) ? config_->total_len() - : unroutable_connection_id_len_; - uint8_t config_id = config_.has_value() ? (config_->config_id() << 6) + uint8_t config_id = config_.has_value() ? config_->config_id() : kLoadBalancerUnroutableConfigId; + uint8_t shifted_config_id = config_id << 6; + uint8_t length = connection_id_lengths_[config_id]; if (config_.has_value() != server_id_.has_value()) { QUIC_BUG(quic_bug_435375038_04) << "Existence of config and server_id are out of sync"; @@ -114,12 +112,12 @@ uint8_t first_byte; // first byte if (len_self_encoded_) { - first_byte = config_id | (length - 1); + first_byte = shifted_config_id | (length - 1); } else { random_.RandBytes(static_cast<void *>(&first_byte), 1); - first_byte = config_id | (first_byte & kLoadBalancerLengthMask); + first_byte = shifted_config_id | (first_byte & kLoadBalancerLengthMask); } - if (config_id == kLoadBalancerUnroutableConfigId) { + if (!config_.has_value()) { return MakeUnroutableConnectionId(first_byte); } QuicConnectionId id; @@ -178,19 +176,27 @@ const QuicConnectionId &original, const ParsedQuicVersion &version) { // Pre-IETF versions of QUIC can respond poorly to new connection IDs issued // during the handshake. - uint8_t needed_length = config_.has_value() ? config_->total_len() - : unroutable_connection_id_len_; + uint8_t needed_length = config_.has_value() + ? config_->total_len() + : connection_id_lengths_[kNumLoadBalancerConfigs]; return (!version.HasIetfQuicFrames() && original.length() == needed_length) ? absl::optional<QuicConnectionId>() : GenerateConnectionId(); } +uint8_t LoadBalancerEncoder::ConnectionIdLength(uint8_t first_byte) const { + if (len_self_encoded()) { + return (first_byte &= kLoadBalancerLengthMask) + 1; + } + return connection_id_lengths_[first_byte >> 6]; +} + QuicConnectionId LoadBalancerEncoder::MakeUnroutableConnectionId( uint8_t first_byte) { QuicConnectionId id; - id.set_length(unroutable_connection_id_len_); + id.set_length(connection_id_lengths_[kLoadBalancerUnroutableConfigId]); id.mutable_data()[0] = first_byte; - random_.RandBytes(&id.mutable_data()[1], unroutable_connection_id_len_ - 1); + random_.RandBytes(&id.mutable_data()[1], connection_id_lengths_[3] - 1); return id; }
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.h b/quiche/quic/load_balancer/load_balancer_encoder.h index 266ebc1..1099d7a 100644 --- a/quiche/quic/load_balancer/load_balancer_encoder.h +++ b/quiche/quic/load_balancer/load_balancer_encoder.h
@@ -18,6 +18,18 @@ // Default length of a 4-tuple connection ID. inline constexpr uint8_t kLoadBalancerUnroutableLen = 8; +// When the encoder is self-encoding the connection ID length, these are the +// bits of the first byte that do so. +constexpr uint8_t kLoadBalancerLengthMask = 0x3f; +// The bits of the connection ID first byte that encode the config ID. +constexpr uint8_t kLoadBalancerConfigIdMask = 0xc0; +// The config ID that means the connection ID does not contain routing +// information. +constexpr uint8_t kLoadBalancerUnroutableConfigId = kNumLoadBalancerConfigs; +// The bits of the connection ID first byte that correspond to a connection ID +// that does not contain routing information. +constexpr uint8_t kLoadBalancerUnroutablePrefix = + kLoadBalancerUnroutableConfigId << 6; // Interface which receives notifications when the current config is updated. class QUIC_EXPORT_PRIVATE LoadBalancerEncoderVisitorInterface { @@ -111,6 +123,7 @@ absl::optional<QuicConnectionId> MaybeReplaceConnectionId( const QuicConnectionId& original, const ParsedQuicVersion& version) override; + uint8_t ConnectionIdLength(uint8_t first_byte) const override; protected: LoadBalancerEncoder(QuicRandom& random, @@ -119,8 +132,9 @@ const uint8_t unroutable_connection_id_len) : random_(random), len_self_encoded_(len_self_encoded), - visitor_(visitor), - unroutable_connection_id_len_(unroutable_connection_id_len) {} + visitor_(visitor) { + std::fill_n(connection_id_lengths_, 4, unroutable_connection_id_len); + } private: friend class test::LoadBalancerEncoderPeer; @@ -130,11 +144,11 @@ QuicRandom& random_; const bool len_self_encoded_; LoadBalancerEncoderVisitorInterface* const visitor_; - const uint8_t unroutable_connection_id_len_; absl::optional<LoadBalancerConfig> config_; absl::uint128 seed_, num_nonces_left_ = 0; absl::optional<LoadBalancerServerId> server_id_; + uint8_t connection_id_lengths_[kNumLoadBalancerConfigs + 1]; }; } // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_encoder_test.cc b/quiche/quic/load_balancer/load_balancer_encoder_test.cc index a331713..12d8fc7 100644 --- a/quiche/quic/load_balancer/load_balancer_encoder_test.cc +++ b/quiche/quic/load_balancer/load_balancer_encoder_test.cc
@@ -397,6 +397,53 @@ EXPECT_EQ(*encoder->GenerateNextConnectionId(TestConnectionId(1)), expected); } +TEST_F(LoadBalancerEncoderTest, ConnectionIdLengthsEncoded) { + // The first byte literally encodes the length. + auto len_encoder = LoadBalancerEncoder::Create(random_, nullptr, true); + ASSERT_TRUE(len_encoder.has_value()); + EXPECT_EQ(len_encoder->ConnectionIdLength(0xc8), 9); + EXPECT_EQ(len_encoder->ConnectionIdLength(0x4a), 11); + EXPECT_EQ(len_encoder->ConnectionIdLength(0x09), 10); + // The length is not self-encoded anymore. + auto encoder = LoadBalancerEncoder::Create(random_, nullptr, false); + ASSERT_TRUE(encoder.has_value()); + EXPECT_EQ(encoder->ConnectionIdLength(0xc8), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x4a), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x09), kQuicDefaultConnectionIdLength); + // Add config ID 0, so that ID now returns a different length. + uint8_t config_id = 0; + uint8_t server_id_len = 3; + uint8_t nonce_len = 6; + uint8_t config_0_len = server_id_len + nonce_len + 1; + auto config0 = LoadBalancerConfig::CreateUnencrypted(config_id, server_id_len, + nonce_len); + ASSERT_TRUE(config0.has_value()); + EXPECT_TRUE( + encoder->UpdateConfig(*config0, MakeServerId(kServerId, server_id_len))); + EXPECT_EQ(encoder->ConnectionIdLength(0xc8), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x4a), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x09), config_0_len); + // Replace config ID 0 with 1. There are probably still packets with config + // ID 0 arriving, so keep that length in memory. + config_id = 1; + nonce_len++; + uint8_t config_1_len = server_id_len + nonce_len + 1; + auto config1 = LoadBalancerConfig::CreateUnencrypted(config_id, server_id_len, + nonce_len); + ASSERT_TRUE(config1.has_value()); + // Old config length still there after replacement + EXPECT_TRUE( + encoder->UpdateConfig(*config1, MakeServerId(kServerId, server_id_len))); + EXPECT_EQ(encoder->ConnectionIdLength(0xc8), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x4a), config_1_len); + EXPECT_EQ(encoder->ConnectionIdLength(0x09), config_0_len); + // Old config length still there after delete + encoder->DeleteConfig(); + EXPECT_EQ(encoder->ConnectionIdLength(0xc8), kQuicDefaultConnectionIdLength); + EXPECT_EQ(encoder->ConnectionIdLength(0x4a), config_1_len); + EXPECT_EQ(encoder->ConnectionIdLength(0x09), config_0_len); +} + } // namespace } // namespace test
diff --git a/quiche/quic/test_tools/mock_connection_id_generator.h b/quiche/quic/test_tools/mock_connection_id_generator.h index 66176f9..42209d6 100644 --- a/quiche/quic/test_tools/mock_connection_id_generator.h +++ b/quiche/quic/test_tools/mock_connection_id_generator.h
@@ -20,6 +20,9 @@ (const quic::QuicConnectionId& original, const quic::ParsedQuicVersion& version), (override)); + + MOCK_METHOD(uint8_t, ConnectionIdLength, (uint8_t first_byte), + (const, override)); }; } // namespace test