Miscellaneous changes to quic/load_balancer: - Add LoadBalancerDecoder::GetConfig to return the LoadBalancerConfig for a config_id. - Add LoadBalancerDecoder::GetConfigId(uint8_t connection_id_first_byte) to get the config id from the first byte of the connection ID. - Add LoadBalancerServerId::Create(absl::string_view data) to create a LoadBalancerServerId from a string_view. PiperOrigin-RevId: 505150431
diff --git a/quiche/quic/load_balancer/load_balancer_decoder.cc b/quiche/quic/load_balancer/load_balancer_decoder.cc index a946b36..cdff314 100644 --- a/quiche/quic/load_balancer/load_balancer_decoder.cc +++ b/quiche/quic/load_balancer/load_balancer_decoder.cc
@@ -75,8 +75,12 @@ if (connection_id.IsEmpty()) { return absl::optional<uint8_t>(); } - const uint8_t first_byte = connection_id.data()[0]; - uint8_t codepoint = (first_byte >> 6); + return GetConfigId(*reinterpret_cast<const uint8_t*>(connection_id.data())); +} + +absl::optional<uint8_t> LoadBalancerDecoder::GetConfigId( + const uint8_t connection_id_first_byte) { + uint8_t codepoint = (connection_id_first_byte >> 6); if (codepoint < kNumLoadBalancerConfigs) { return codepoint; }
diff --git a/quiche/quic/load_balancer/load_balancer_decoder.h b/quiche/quic/load_balancer/load_balancer_decoder.h index 2c326cb..f5ba6f3 100644 --- a/quiche/quic/load_balancer/load_balancer_decoder.h +++ b/quiche/quic/load_balancer/load_balancer_decoder.h
@@ -20,7 +20,17 @@ // Remove support for a config. Does nothing if there is no config for // |config_id|. Does nothing and creates a bug if |config_id| is greater than // 2. - void DeleteConfig(const uint8_t config_id); + void DeleteConfig(uint8_t config_id); + + // Return the config for |config_id|, or nullptr if not found. + const LoadBalancerConfig* GetConfig(const uint8_t config_id) const { + if (config_id >= kNumLoadBalancerConfigs || + !config_[config_id].has_value()) { + return nullptr; + } + + return &config_[config_id].value(); + } // Extract a server ID from |connection_id|. If there is no config for the // codepoint, |connection_id| is too short, or there's a decrypt error, @@ -30,10 +40,15 @@ const QuicConnectionId& connection_id) const; // Returns the config ID stored in the first two bits of |connection_id|, or - // empty if |connection_id| is empty. + // empty if |connection_id| is empty, or the first two bits of the first byte + // of |connection_id| are 0b11. static absl::optional<uint8_t> GetConfigId( const QuicConnectionId& connection_id); + // Returns the config ID stored in the first two bits of + // |connection_id_first_byte|, or empty if the first two bits are 0b11. + static absl::optional<uint8_t> GetConfigId(uint8_t connection_id_first_byte); + private: // Decoders can support up to 3 configs at once. absl::optional<LoadBalancerConfig> config_[kNumLoadBalancerConfigs];
diff --git a/quiche/quic/load_balancer/load_balancer_decoder_test.cc b/quiche/quic/load_balancer/load_balancer_decoder_test.cc index 728467f..b2bdfd3 100644 --- a/quiche/quic/load_balancer/load_balancer_decoder_test.cc +++ b/quiche/quic/load_balancer/load_balancer_decoder_test.cc
@@ -208,8 +208,10 @@ EXPECT_FALSE( LoadBalancerDecoder::GetConfigId(QuicConnectionId()).has_value()); for (uint8_t i = 0; i < 3; i++) { - auto config_id = LoadBalancerDecoder::GetConfigId( - QuicConnectionId({static_cast<unsigned char>(i << 6)})); + const QuicConnectionId connection_id({static_cast<unsigned char>(i << 6)}); + auto config_id = LoadBalancerDecoder::GetConfigId(connection_id); + EXPECT_EQ(config_id, + LoadBalancerDecoder::GetConfigId(connection_id.data()[0])); EXPECT_TRUE(config_id.has_value()); EXPECT_EQ(*config_id, i); } @@ -217,6 +219,22 @@ LoadBalancerDecoder::GetConfigId(QuicConnectionId({0xc0})).has_value()); } +TEST_F(LoadBalancerDecoderTest, GetConfig) { + LoadBalancerDecoder decoder; + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(2, 3, 4)); + + EXPECT_EQ(decoder.GetConfig(0), nullptr); + EXPECT_EQ(decoder.GetConfig(1), nullptr); + EXPECT_EQ(decoder.GetConfig(3), nullptr); + EXPECT_EQ(decoder.GetConfig(4), nullptr); + + const LoadBalancerConfig* config = decoder.GetConfig(2); + ASSERT_NE(config, nullptr); + EXPECT_EQ(config->server_id_len(), 3); + EXPECT_EQ(config->nonce_len(), 4); + EXPECT_FALSE(config->IsEncrypted()); +} + } // namespace } // namespace test
diff --git a/quiche/quic/load_balancer/load_balancer_server_id.h b/quiche/quic/load_balancer/load_balancer_server_id.h index 18fc40e..cdd3c95 100644 --- a/quiche/quic/load_balancer/load_balancer_server_id.h +++ b/quiche/quic/load_balancer/load_balancer_server_id.h
@@ -26,7 +26,13 @@ public: // Copies all the bytes from |data| into a new LoadBalancerServerId. static absl::optional<LoadBalancerServerId> Create( - const absl::Span<const uint8_t> data); + absl::Span<const uint8_t> data); + + // For callers with a string_view at hand. + static absl::optional<LoadBalancerServerId> Create(absl::string_view data) { + return Create(absl::MakeSpan(reinterpret_cast<const uint8_t*>(data.data()), + data.length())); + } // Server IDs are opaque bytes, but defining these operators allows us to sort // them into a tree and define ranges.