QUIC-LB Decoder Class. Used by load balancers to manage configs and extract server IDs from connection IDs. PiperOrigin-RevId: 439865147
diff --git a/quic/load_balancer/load_balancer_decoder.cc b/quic/load_balancer/load_balancer_decoder.cc new file mode 100644 index 0000000..6d7264f --- /dev/null +++ b/quic/load_balancer/load_balancer_decoder.cc
@@ -0,0 +1,86 @@ +// Copyright (c) 2022 The Chromium Authors. All rights reserved. +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. + +#include "quic/load_balancer/load_balancer_decoder.h" + +#include "quic/platform/api/quic_bug_tracker.h" + +namespace quic { + +bool LoadBalancerDecoder::AddConfig(const LoadBalancerConfig& config) { + if (config_[config.config_id()].has_value()) { + return false; + } + config_[config.config_id()] = config; + return true; +} + +void LoadBalancerDecoder::DeleteConfig(uint8_t config_id) { + if (config_id > 2) { + QUIC_BUG(quic_bug_438896865_01) + << "Decoder deleting config with invalid config_id " + << static_cast<int>(config_id); + return; + } + config_[config_id].reset(); +} + +// This is the core logic to extract a server ID given a valid config and +// connection ID of sufficient length. +absl::optional<LoadBalancerServerId> LoadBalancerDecoder::GetServerId( + const QuicConnectionId& connection_id) const { + absl::optional<uint8_t> config_id = GetConfigId(connection_id); + if (!config_id.has_value()) { + return absl::optional<LoadBalancerServerId>(); + } + absl::optional<LoadBalancerConfig> config = config_[*config_id]; + if (!config.has_value()) { + return absl::optional<LoadBalancerServerId>(); + } + if (connection_id.length() < (1 + config->total_len())) { + // Connection ID wasn't long enough + return absl::optional<LoadBalancerServerId>(); + } + // The first byte is complete. Finish the rest. + const uint8_t* data = + reinterpret_cast<const uint8_t*>(connection_id.data()) + 1; + if (!config->IsEncrypted()) { // It's a Plaintext CID. + return LoadBalancerServerId::Create( + absl::Span<const uint8_t>(data, config->server_id_len())); + } + uint8_t result[kQuicMaxConnectionIdWithLengthPrefixLength]; + if (config->total_len() == kLoadBalancerKeyLen) { // single pass + if (!config->BlockDecrypt(data, result)) { + return absl::optional<LoadBalancerServerId>(); + } + } else { + // Do 3 or 4 passes. Only 3 are necessary if the server_id is short enough + // to fit in the first half of the connection ID (the decoder doesn't need + // to extract the nonce). + memcpy(result, data, config->total_len()); + uint8_t end = (config->server_id_len() > config->nonce_len()) ? 1 : 2; + for (uint8_t i = 4; i >= end; i--) { + if (!config->EncryptionPass(result, i)) { + return absl::optional<LoadBalancerServerId>(); + } + } + } + return LoadBalancerServerId::Create( + absl::Span<const uint8_t>(result, config->server_id_len())); +} + +absl::optional<uint8_t> LoadBalancerDecoder::GetConfigId( + const QuicConnectionId& connection_id) { + if (connection_id.IsEmpty()) { + return absl::optional<uint8_t>(); + } + const uint8_t first_byte = connection_id.data()[0]; + uint8_t codepoint = (first_byte >> 6); + if (codepoint <= 2) { + return codepoint; + } + return absl::optional<uint8_t>(); +} + +} // namespace quic
diff --git a/quic/load_balancer/load_balancer_decoder.h b/quic/load_balancer/load_balancer_decoder.h new file mode 100644 index 0000000..a2708fe --- /dev/null +++ b/quic/load_balancer/load_balancer_decoder.h
@@ -0,0 +1,45 @@ +// Copyright (c) 2022 The Chromium Authors. All rights reserved. +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. + +#ifndef QUICHE_QUIC_LOAD_BALANCER_LOAD_BALANCER_DECODER_H_ +#define QUICHE_QUIC_LOAD_BALANCER_LOAD_BALANCER_DECODER_H_ + +#include "quic/load_balancer/load_balancer_config.h" +#include "quic/load_balancer/load_balancer_server_id.h" + +namespace quic { + +// Manages QUIC-LB configurations to extract a server ID from a properly +// encoded connection ID, usually on behalf of a load balancer. +class QUIC_EXPORT_PRIVATE LoadBalancerDecoder { + public: + // Returns false if the config_id codepoint is already occupied. + bool AddConfig(const LoadBalancerConfig& config); + + // Remove support for a config + void DeleteConfig(const uint8_t config_id); + + // For these "Get" functions, the calling code might not know the length of + // the connection ID. That's OK; if not, just send at least + // kQuicMaxConnectionIdWithLengthPrefixLength bytes in a QuicConnectionId. + + // Extract a server ID from a connection ID. If there is no config for the + // codepoint, the connection ID is too short, or there's a decrypt error, + // returns empty. + absl::optional<LoadBalancerServerId> GetServerId( + 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. + static absl::optional<uint8_t> GetConfigId( + const QuicConnectionId& connection_id); + + private: + // Decoders can support up to 3 configs at once. + absl::optional<LoadBalancerConfig> config_[3]; +}; + +} // namespace quic + +#endif // QUICHE_QUIC_LOAD_BALANCER_LOAD_BALANCER_DECODER_H_
diff --git a/quic/load_balancer/load_balancer_decoder_test.cc b/quic/load_balancer/load_balancer_decoder_test.cc new file mode 100644 index 0000000..eaeef42 --- /dev/null +++ b/quic/load_balancer/load_balancer_decoder_test.cc
@@ -0,0 +1,228 @@ +// Copyright (c) 2022 The Chromium Authors. All rights reserved. +// Use of this source code is governed by a BSD-style license that can be +// found in the LICENSE file. + +#include "quic/load_balancer/load_balancer_decoder.h" + +#include "absl/base/macros.h" +#include "quic/load_balancer/load_balancer_server_id.h" +#include "quic/platform/api/quic_expect_bug.h" +#include "quic/platform/api/quic_test.h" +#include "quic/test_tools/quic_test_utils.h" + +namespace quic { + +namespace test { + +namespace { + +class LoadBalancerDecoderTest : public QuicTest {}; + +// Convenience function to shorten the code. Does not check if |array| is long +// enough or |length| is valid for a server ID. +inline LoadBalancerServerId MakeServerId(const uint8_t array[], + const uint8_t length) { + return *LoadBalancerServerId::Create( + absl::Span<const uint8_t>(array, length)); +} + +constexpr char kRawKey[] = {0x8f, 0x95, 0xf0, 0x92, 0x45, 0x76, 0x5f, 0x80, + 0x25, 0x69, 0x34, 0xe5, 0x0c, 0x66, 0x20, 0x7f}; +constexpr absl::string_view kKey(kRawKey, kLoadBalancerKeyLen); +constexpr uint8_t kServerId[] = {0xed, 0x79, 0x3a, 0x51, 0xd4, 0x9b, 0x8f, 0x5f, + 0xab, 0x65, 0xba, 0x04, 0xc3, 0x33, 0x0a}; + +struct LoadBalancerDecoderTestCase { + LoadBalancerConfig config; + QuicConnectionId connection_id; + LoadBalancerServerId server_id; +}; + +TEST_F(LoadBalancerDecoderTest, UnencryptedConnectionIdTestVectors) { + const struct LoadBalancerDecoderTestCase test_vectors[2] = { + { + *LoadBalancerConfig::CreateUnencrypted(0, 3, 4), + QuicConnectionId({0x07, 0xed, 0x79, 0x3a, 0x80, 0x49, 0x71, 0x8a}), + MakeServerId(kServerId, 3), + }, + { + *LoadBalancerConfig::CreateUnencrypted(1, 8, 5), + QuicConnectionId({0x4d, 0xed, 0x79, 0x3a, 0x51, 0xd4, 0x9b, 0x8f, + 0x5f, 0xee, 0x15, 0xda, 0x27, 0xc4}), + MakeServerId(kServerId, 8), + }}; + for (uint8_t i = 0; i < ABSL_ARRAYSIZE(test_vectors); i++) { + LoadBalancerDecoder decoder; + EXPECT_TRUE(decoder.AddConfig(test_vectors[i].config)); + EXPECT_EQ(decoder.GetServerId(test_vectors[i].connection_id), + test_vectors[i].server_id); + } +} + +// Compare test vectors from Appendix B of draft-ietf-quic-load-balancers-12. +TEST_F(LoadBalancerDecoderTest, DecoderTestVectors) { + // Try (1) the "standard" CID length of 8 + // (2) server_id_len > nonce_len, so there is a fourth decryption pass + // (3) the single-pass encryption case + // (4) An even total length. + const struct LoadBalancerDecoderTestCase test_vectors[4] = { + { + *LoadBalancerConfig::Create(0, 3, 4, kKey), + QuicConnectionId({0x07, 0xfb, 0xfe, 0x05, 0xf7, 0x31, 0xb4, 0x25}), + MakeServerId(kServerId, 3), + }, + { + *LoadBalancerConfig::Create(1, 10, 5, kKey), + QuicConnectionId({0x4f, 0x01, 0x09, 0x56, 0xfb, 0x5c, 0x1d, 0x4d, + 0x86, 0xe0, 0x10, 0x18, 0x3e, 0x0b, 0x7d, 0x1e}), + MakeServerId(kServerId, 10), + }, + { + *LoadBalancerConfig::Create(2, 8, 8, kKey), + QuicConnectionId({0x90, 0x4d, 0xd2, 0xd0, 0x5a, 0x7b, 0x0d, 0xe9, + 0xb2, 0xb9, 0x90, 0x7a, 0xfb, 0x5e, 0xcf, 0x8c, + 0xc3}), + MakeServerId(kServerId, 8), + }, + { + *LoadBalancerConfig::Create(0, 9, 9, kKey), + QuicConnectionId({0x12, 0x7a, 0x28, 0x5a, 0x09, 0xf8, 0x52, 0x80, + 0xf4, 0xfd, 0x6a, 0xbb, 0x43, 0x4a, 0x71, 0x59, + 0xe4, 0xd3, 0xeb}), + MakeServerId(kServerId, 9), + }, + }; + + for (uint8_t i = 0; i < ABSL_ARRAYSIZE(test_vectors); i++) { + LoadBalancerDecoder decoder; + EXPECT_TRUE(decoder.AddConfig(test_vectors[i].config)); + EXPECT_EQ(decoder.GetServerId(test_vectors[i].connection_id), + test_vectors[i].server_id); + } +} + +TEST_F(LoadBalancerDecoderTest, NoServerIdEntry) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id.has_value()); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(0, 3, 4))); + QuicConnectionId no_server_id_entry( + {0x00, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}); + EXPECT_TRUE(decoder.GetServerId(no_server_id_entry).has_value()); +} + +TEST_F(LoadBalancerDecoderTest, InvalidConfigId) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id.has_value()); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(1, 3, 4))); + QuicConnectionId wrong_config_id( + {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}); + EXPECT_FALSE(decoder + .GetServerId(QuicConnectionId( + {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07})) + .has_value()); +} + +TEST_F(LoadBalancerDecoderTest, UnroutableCodepoint) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id.has_value()); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(1, 3, 4))); + EXPECT_FALSE(decoder + .GetServerId(QuicConnectionId( + {0xc0, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07})) + .has_value()); +} + +TEST_F(LoadBalancerDecoderTest, UnroutableCodepointAnyLength) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id.has_value()); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(1, 3, 4))); + EXPECT_FALSE(decoder.GetServerId(QuicConnectionId({0xff})).has_value()); +} + +TEST_F(LoadBalancerDecoderTest, ConnectionIdTooShort) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id.has_value()); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(0, 3, 4))); + EXPECT_FALSE(decoder + .GetServerId(QuicConnectionId( + {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06})) + .has_value()); +} + +TEST_F(LoadBalancerDecoderTest, ConnectionIdTooLongIsOK) { + auto server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(0, 3, 4))); + auto server_id_result = decoder.GetServerId( + QuicConnectionId({0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08})); + EXPECT_TRUE(server_id_result.has_value()); + EXPECT_EQ(server_id_result, server_id); +} + +TEST_F(LoadBalancerDecoderTest, DeleteConfigBadId) { + LoadBalancerDecoder decoder; + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(2, 3, 4)); + decoder.DeleteConfig(0); + EXPECT_QUIC_BUG(decoder.DeleteConfig(3), + "Decoder deleting config with invalid config_id 3"); + EXPECT_TRUE(decoder + .GetServerId(QuicConnectionId( + {0x80, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07})) + .has_value()); +} + +TEST_F(LoadBalancerDecoderTest, DeleteConfigGoodId) { + LoadBalancerDecoder decoder; + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(2, 3, 4)); + decoder.DeleteConfig(2); + EXPECT_FALSE(decoder + .GetServerId(QuicConnectionId( + {0x80, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07})) + .has_value()); +} + +// Create two server IDs and make sure the decoder decodes the correct one. +TEST_F(LoadBalancerDecoderTest, TwoServerIds) { + auto server_id1 = LoadBalancerServerId::Create({0x01, 0x02, 0x03}); + EXPECT_TRUE(server_id1.has_value()); + auto server_id2 = LoadBalancerServerId::Create({0x04, 0x05, 0x06}); + LoadBalancerDecoder decoder; + EXPECT_TRUE( + decoder.AddConfig(*LoadBalancerConfig::CreateUnencrypted(0, 3, 4))); + EXPECT_EQ(decoder.GetServerId(QuicConnectionId( + {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07})), + server_id1); + EXPECT_EQ(decoder.GetServerId(QuicConnectionId( + {0x00, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a})), + server_id2); +} + +TEST_F(LoadBalancerDecoderTest, GetConfigId) { + 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)})); + EXPECT_TRUE(config_id.has_value()); + EXPECT_EQ(*config_id, i); + } + EXPECT_FALSE( + LoadBalancerDecoder::GetConfigId(QuicConnectionId({0xc0})).has_value()); +} + +} // namespace + +} // namespace test + +} // namespace quic