Relocate QUICHE files into quiche/ directory within the quiche repo, and change the relative include paths accordingly.
PiperOrigin-RevId: 440164720
Change-Id: I64d8a975d08888a3a86f6c51908e63d5cd45fa35
diff --git a/quiche/quic/load_balancer/load_balancer_config.cc b/quiche/quic/load_balancer/load_balancer_config.cc
new file mode 100644
index 0000000..2061abb
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_config.cc
@@ -0,0 +1,195 @@
+// 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 "quiche/quic/load_balancer/load_balancer_config.h"
+
+#include <memory>
+#include <string_view>
+
+#include "third_party/boringssl/src/include/openssl/aes.h"
+#include "quiche/quic/platform/api/quic_bug_tracker.h"
+
+namespace quic {
+
+namespace {
+
+// Validates all non-key parts of the input.
+bool CommonValidation(const uint8_t config_id, const uint8_t server_id_len,
+ const uint8_t nonce_len) {
+ if (config_id > 2 || server_id_len == 0 ||
+ nonce_len < kLoadBalancerMinNonceLen ||
+ nonce_len > kLoadBalancerMaxNonceLen ||
+ server_id_len >
+ (kQuicMaxConnectionIdWithLengthPrefixLength - nonce_len - 1)) {
+ QUIC_BUG(quic_bug_433862549_01)
+ << "Invalid LoadBalancerConfig "
+ << "Config ID " << static_cast<int>(config_id) << " Server ID Length "
+ << static_cast<int>(server_id_len) << " Nonce Length "
+ << static_cast<int>(nonce_len);
+ return false;
+ }
+ return true;
+}
+
+// Initialize the key in the constructor
+absl::optional<AES_KEY> BuildKey(absl::string_view key, bool encrypt) {
+ if (key.empty()) {
+ return absl::optional<AES_KEY>();
+ }
+ AES_KEY raw_key;
+ if (encrypt) {
+ if (AES_set_encrypt_key(reinterpret_cast<const uint8_t *>(key.data()),
+ key.size() * 8, &raw_key) < 0) {
+ return absl::optional<AES_KEY>();
+ }
+ } else if (AES_set_decrypt_key(reinterpret_cast<const uint8_t *>(key.data()),
+ key.size() * 8, &raw_key) < 0) {
+ return absl::optional<AES_KEY>();
+ }
+ return raw_key;
+}
+
+// Functions to handle 4-pass encryption/decryption.
+// TakePlaintextFrom{Left,Right}() reads the left or right half of 'from' and
+// expands it into a full encryption block ('to') in accordance with the
+// internet-draft.
+void TakePlaintextFromLeft(uint8_t *to, uint8_t *from, uint8_t total_len,
+ uint8_t index) {
+ uint8_t half = total_len / 2;
+ memset(to, 0, kLoadBalancerBlockSize - 1);
+ memcpy(to, from, half);
+ if (total_len % 2) {
+ to[half] = from[half] & 0xf0;
+ }
+ to[kLoadBalancerBlockSize - 1] = index;
+}
+
+void TakePlaintextFromRight(uint8_t *to, uint8_t *from, uint8_t total_len,
+ uint8_t index) {
+ const uint8_t half = total_len / 2;
+ const uint8_t write_point = kLoadBalancerBlockSize - half;
+ const uint8_t read_point = total_len - half;
+ memset((to + 1), 0, kLoadBalancerBlockSize - 1);
+ memcpy(to + write_point, from + read_point, half);
+ if (total_len % 2) {
+ to[write_point - 1] = from[read_point - 1] & 0x0f;
+ }
+ to[0] = index;
+}
+
+// CiphertextXorWith{Left,Right}() takes the relevant end of the ciphertext in
+// 'from' and XORs it with half of the ConnectionId stored at 'to', in
+// accordance with the internet-draft.
+void CiphertextXorWithLeft(uint8_t *to, uint8_t *from, uint8_t total_len) {
+ uint8_t half = total_len / 2;
+ for (int i = 0; i < half; i++) {
+ *(to + i) ^= *(from + i);
+ }
+ if (total_len % 2) {
+ *(to + half) ^= (*(from + half) & 0xf0);
+ }
+}
+
+void CiphertextXorWithRight(uint8_t *to, uint8_t *from, uint8_t total_len) {
+ const uint8_t half = total_len / 2;
+ const uint8_t write_point = total_len - half;
+ const uint8_t read_point = kLoadBalancerBlockSize - half;
+ if (total_len % 2) {
+ *(to + write_point - 1) ^= (*(from + read_point - 1) & 0x0f);
+ }
+ for (int i = 0; i < half; i++) {
+ *(to + write_point + i) ^= *(from + read_point + i);
+ }
+}
+
+} // namespace
+
+absl::optional<LoadBalancerConfig> LoadBalancerConfig::Create(
+ const uint8_t config_id, const uint8_t server_id_len,
+ const uint8_t nonce_len, const absl::string_view key) {
+ // Check for valid parameters.
+ if (key.size() != kLoadBalancerKeyLen) {
+ QUIC_BUG(quic_bug_433862549_02)
+ << "Invalid LoadBalancerConfig Key Length: " << key.size();
+ return absl::optional<LoadBalancerConfig>();
+ }
+ if (!CommonValidation(config_id, server_id_len, nonce_len)) {
+ return absl::optional<LoadBalancerConfig>();
+ }
+ auto new_config =
+ LoadBalancerConfig(config_id, server_id_len, nonce_len, key);
+ if (!new_config.IsEncrypted()) {
+ // Something went wrong in assigning the key!
+ QUIC_BUG(quic_bug_433862549_03) << "Something went wrong in initializing "
+ "the load balancing key.";
+ return absl::optional<LoadBalancerConfig>();
+ }
+ return new_config;
+}
+
+// Creates an unencrypted config.
+absl::optional<LoadBalancerConfig> LoadBalancerConfig::CreateUnencrypted(
+ const uint8_t config_id, const uint8_t server_id_len,
+ const uint8_t nonce_len) {
+ return CommonValidation(config_id, server_id_len, nonce_len)
+ ? LoadBalancerConfig(config_id, server_id_len, nonce_len, "")
+ : absl::optional<LoadBalancerConfig>();
+}
+
+bool LoadBalancerConfig::EncryptionPass(uint8_t *target,
+ const uint8_t index) const {
+ uint8_t plaintext[kLoadBalancerBlockSize], ciphertext[kLoadBalancerBlockSize];
+ if (!key_.has_value() || target == nullptr) {
+ return false;
+ }
+ if (index % 2) { // Odd indices go from left to right
+ TakePlaintextFromLeft(plaintext, target, total_len(), index);
+ } else {
+ TakePlaintextFromRight(plaintext, target, total_len(), index);
+ }
+ if (!BlockEncrypt(plaintext, ciphertext)) {
+ return false;
+ }
+ // XOR bits over the correct half.
+ if (index % 2) {
+ CiphertextXorWithRight(target, ciphertext, total_len());
+ } else {
+ CiphertextXorWithLeft(target, ciphertext, total_len());
+ }
+ return true;
+}
+
+bool LoadBalancerConfig::BlockEncrypt(
+ const uint8_t plaintext[kLoadBalancerBlockSize],
+ uint8_t ciphertext[kLoadBalancerBlockSize]) const {
+ if (!key_.has_value()) {
+ return false;
+ }
+ AES_encrypt(plaintext, ciphertext, &key_.value());
+ return true;
+}
+
+bool LoadBalancerConfig::BlockDecrypt(
+ const uint8_t ciphertext[kLoadBalancerBlockSize],
+ uint8_t plaintext[kLoadBalancerBlockSize]) const {
+ if (!block_decrypt_key_.has_value()) {
+ return false;
+ }
+ AES_decrypt(ciphertext, plaintext, &block_decrypt_key_.value());
+ return true;
+}
+
+LoadBalancerConfig::LoadBalancerConfig(const uint8_t config_id,
+ const uint8_t server_id_len,
+ const uint8_t nonce_len,
+ const absl::string_view key)
+ : config_id_(config_id),
+ server_id_len_(server_id_len),
+ nonce_len_(nonce_len),
+ key_(BuildKey(key, /* encrypt = */ true)),
+ block_decrypt_key_((server_id_len + nonce_len == kLoadBalancerBlockSize)
+ ? BuildKey(key, /* encrypt = */ false)
+ : absl::optional<AES_KEY>()) {}
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_config.h b/quiche/quic/load_balancer/load_balancer_config.h
new file mode 100644
index 0000000..8b21a1f
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_config.h
@@ -0,0 +1,88 @@
+// 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_CONFIG_H_
+#define QUICHE_QUIC_LOAD_BALANCER_CONFIG_H_
+
+#include "third_party/boringssl/src/include/openssl/aes.h"
+#include "quiche/quic/core/quic_types.h"
+#include "quiche/quic/platform/api/quic_export.h"
+
+namespace quic {
+
+inline constexpr uint8_t kLoadBalancerKeyLen = 16;
+// Regardless of key length, the AES block size is always 16 Bytes.
+inline constexpr uint8_t kLoadBalancerBlockSize = 16;
+// The spec says nonces can be 18 bytes, but 16 lets it be a uint128.
+inline constexpr uint8_t kLoadBalancerMaxNonceLen = 16;
+inline constexpr uint8_t kLoadBalancerMinNonceLen = 4;
+
+// This the base class for QUIC-LB configuration. It contains configuration
+// elements usable by both encoders (servers) and decoders (load balancers).
+// Confusingly, it is called "LoadBalancerConfig" because it pertains to objects
+// that both servers and load balancers use to interact with each other.
+class QUIC_EXPORT_PRIVATE LoadBalancerConfig {
+ public:
+ // This factory function initializes an encrypted LoadBalancerConfig and
+ // returns it in absl::optional, which is empty if the config is invalid.
+ // config_id: The first two bits of the Connection Id. Must be no larger than
+ // 2.
+ // server_id_len: Expected length of the server ids associated with this
+ // config. Must be greater than 0 and less than 16.
+ // nonce_len: Length of the nonce. Must be at least 4 and no larger than 16.
+ // Further the server_id_len + nonce_len must be no larger than 19.
+ // key: The encryption key must be 16B long.
+ static absl::optional<LoadBalancerConfig> Create(const uint8_t config_id,
+ const uint8_t server_id_len,
+ const uint8_t nonce_len,
+ const absl::string_view key);
+
+ // Creates an unencrypted config.
+ static absl::optional<LoadBalancerConfig> CreateUnencrypted(
+ const uint8_t config_id, const uint8_t server_id_len,
+ const uint8_t nonce_len);
+
+ // Handles one pass of 4-pass encryption. Encoder and decoder use of this
+ // function varies substantially, so they are not implemented here.
+ // Returns false if the config is not encrypted.
+ ABSL_MUST_USE_RESULT bool EncryptionPass(uint8_t *target,
+ const uint8_t index) const;
+ // Use the key to do a block encryption, which is used both in all cases of
+ // encrypted configs. Returns false if there's no key.
+ ABSL_MUST_USE_RESULT bool BlockEncrypt(
+ const uint8_t plaintext[kLoadBalancerBlockSize],
+ uint8_t ciphertext[kLoadBalancerBlockSize]) const;
+ // Returns false if the config does not require block decryption.
+ ABSL_MUST_USE_RESULT bool BlockDecrypt(
+ const uint8_t ciphertext[kLoadBalancerBlockSize],
+ uint8_t plaintext[kLoadBalancerBlockSize]) const;
+
+ uint8_t config_id() const { return config_id_; }
+ uint8_t server_id_len() const { return server_id_len_; }
+ uint8_t nonce_len() const { return nonce_len_; }
+ uint8_t total_len() const { return server_id_len_ + nonce_len_; }
+ bool IsEncrypted() const { return key_.has_value(); }
+
+ private:
+ // Constructor is private because it doesn't validate input.
+ LoadBalancerConfig(uint8_t config_id, uint8_t server_id_len,
+ uint8_t nonce_len, absl::string_view key);
+
+ uint8_t config_id_;
+ uint8_t server_id_len_;
+ uint8_t nonce_len_;
+ // All Connection ID encryption and decryption uses the AES_encrypt function
+ // at root, so there is a single key for all of it. This is empty if the
+ // config is not encrypted.
+ absl::optional<AES_KEY> key_;
+ // The one exception is that when total_len == 16, connection ID decryption
+ // uses AES_decrypt. The bytes that comprise the key are the same, but
+ // AES_decrypt requires an AES_KEY that is initialized differently. In all
+ // other cases, block_decrypt_key_ is empty.
+ absl::optional<AES_KEY> block_decrypt_key_;
+};
+
+} // namespace quic
+
+#endif // QUICHE_QUIC_LOAD_BALANCER_CONFIG_H
diff --git a/quiche/quic/load_balancer/load_balancer_config_test.cc b/quiche/quic/load_balancer/load_balancer_config_test.cc
new file mode 100644
index 0000000..1dad91c
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_config_test.cc
@@ -0,0 +1,152 @@
+// 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 "quiche/quic/load_balancer/load_balancer_config.h"
+
+#include "quiche/quic/platform/api/quic_expect_bug.h"
+#include "quiche/quic/platform/api/quic_test.h"
+#include "quiche/quic/test_tools/quic_test_utils.h"
+
+namespace quic {
+
+namespace test {
+
+namespace {
+
+constexpr char raw_key[] = {
+ 0xfd, 0xf7, 0x26, 0xa9, 0x89, 0x3e, 0xc0, 0x5c,
+ 0x06, 0x32, 0xd3, 0x95, 0x66, 0x80, 0xba, 0xf0,
+};
+
+class LoadBalancerConfigTest : public QuicTest {};
+
+TEST_F(LoadBalancerConfigTest, InvalidParams) {
+ // Bogus config_id.
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerConfig::CreateUnencrypted(3, 4, 10).has_value()),
+ "Invalid LoadBalancerConfig Config ID 3 Server ID Length 4 "
+ "Nonce Length 10");
+ // Bad Server ID lengths.
+ EXPECT_QUIC_BUG(EXPECT_FALSE(LoadBalancerConfig::Create(
+ 2, 0, 10, absl::string_view(raw_key, 16))
+ .has_value()),
+ "Invalid LoadBalancerConfig Config ID 2 Server ID Length 0 "
+ "Nonce Length 10");
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerConfig::CreateUnencrypted(2, 16, 4).has_value()),
+ "Invalid LoadBalancerConfig Config ID 2 Server ID Length 16 "
+ "Nonce Length 4");
+ // Bad Nonce lengths.
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerConfig::CreateUnencrypted(2, 4, 2).has_value()),
+ "Invalid LoadBalancerConfig Config ID 2 Server ID Length 4 "
+ "Nonce Length 2");
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerConfig::CreateUnencrypted(2, 1, 17).has_value()),
+ "Invalid LoadBalancerConfig Config ID 2 Server ID Length 1 "
+ "Nonce Length 17");
+ // Bad key lengths.
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerConfig::Create(2, 3, 4, "").has_value()),
+ "Invalid LoadBalancerConfig Key Length: 0");
+ EXPECT_QUIC_BUG(EXPECT_FALSE(LoadBalancerConfig::Create(
+ 2, 3, 4, absl::string_view(raw_key, 10))
+ .has_value()),
+ "Invalid LoadBalancerConfig Key Length: 10");
+ EXPECT_QUIC_BUG(EXPECT_FALSE(LoadBalancerConfig::Create(
+ 0, 3, 4, absl::string_view(raw_key, 17))
+ .has_value()),
+ "Invalid LoadBalancerConfig Key Length: 17");
+}
+
+TEST_F(LoadBalancerConfigTest, ValidParams) {
+ // Test valid configurations and accessors
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_EQ(config->config_id(), 0);
+ EXPECT_EQ(config->server_id_len(), 3);
+ EXPECT_EQ(config->nonce_len(), 4);
+ EXPECT_EQ(config->total_len(), 7);
+ EXPECT_FALSE(config->IsEncrypted());
+ auto config2 =
+ LoadBalancerConfig::Create(2, 6, 7, absl::string_view(raw_key, 16));
+ EXPECT_TRUE(config.has_value());
+ EXPECT_EQ(config2->config_id(), 2);
+ EXPECT_EQ(config2->server_id_len(), 6);
+ EXPECT_EQ(config2->nonce_len(), 7);
+ EXPECT_EQ(config2->total_len(), 13);
+ EXPECT_TRUE(config2->IsEncrypted());
+}
+
+// Compare EncryptionPass() results to the example in
+// draft-ietf-quic-load-balancers-12, Section 5.3.2.
+TEST_F(LoadBalancerConfigTest, TestEncryptionPassExample) {
+ auto config =
+ LoadBalancerConfig::Create(0, 3, 4, absl::string_view(raw_key, 16));
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(config->IsEncrypted());
+ std::array<uint8_t, 7> bytes = {
+ 0x31, 0x44, 0x1a, 0x9c, 0x69, 0xc2, 0x75,
+ };
+ EXPECT_FALSE(config->EncryptionPass(nullptr, 0));
+ EXPECT_TRUE(config->EncryptionPass(bytes.data(), 1));
+ EXPECT_TRUE((bytes == std::array<uint8_t, 7>(
+ {0x31, 0x44, 0x1a, 0x9d, 0xbc, 0x04, 0x26})));
+ EXPECT_TRUE(config->EncryptionPass(bytes.data(), 2));
+ EXPECT_TRUE((bytes == std::array<uint8_t, 7>(
+ {0x4f, 0xdd, 0x0c, 0x9d, 0xbc, 0x04, 0x26})));
+ EXPECT_TRUE(config->EncryptionPass(bytes.data(), 3));
+ EXPECT_TRUE((bytes == std::array<uint8_t, 7>(
+ {0x4f, 0xdd, 0x0c, 0x9b, 0xba, 0x1e, 0xe0})));
+ EXPECT_TRUE(config->EncryptionPass(bytes.data(), 4));
+ EXPECT_TRUE((bytes == std::array<uint8_t, 7>(
+ {0xe2, 0x3c, 0xb4, 0x2b, 0xba, 0x1e, 0xe0})));
+}
+
+TEST_F(LoadBalancerConfigTest, EncryptionPassPlaintext) {
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ std::array<uint8_t, 7> bytes = {
+ 0x31, 0x44, 0x1a, 0x9c, 0x69, 0xc2, 0x75,
+ };
+ EXPECT_FALSE(config->EncryptionPass(bytes.data(), 1));
+}
+
+TEST_F(LoadBalancerConfigTest, InvalidBlockEncryption) {
+ uint8_t pt[kLoadBalancerBlockSize], ct[kLoadBalancerBlockSize];
+ auto pt_config = LoadBalancerConfig::CreateUnencrypted(0, 8, 8);
+ EXPECT_FALSE(pt_config->BlockEncrypt(pt, ct));
+ EXPECT_FALSE(pt_config->BlockDecrypt(ct, pt));
+ EXPECT_FALSE(pt_config->EncryptionPass(pt, 0));
+ auto small_cid_config =
+ LoadBalancerConfig::Create(0, 3, 4, absl::string_view(raw_key, 16));
+ EXPECT_TRUE(small_cid_config->BlockEncrypt(pt, ct));
+ EXPECT_FALSE(small_cid_config->BlockDecrypt(ct, pt));
+ auto block_config =
+ LoadBalancerConfig::Create(0, 8, 8, absl::string_view(raw_key, 16));
+ EXPECT_TRUE(block_config->BlockEncrypt(pt, ct));
+ EXPECT_TRUE(block_config->BlockDecrypt(ct, pt));
+}
+
+// Block decrypt test from the Test Vector in
+// draft-ietf-quic-load-balancers-12, Appendix B.
+TEST_F(LoadBalancerConfigTest, BlockEncryptionExample) {
+ const uint8_t ptext[] = {0xed, 0x79, 0x3a, 0x51, 0xd4, 0x9b, 0x8f, 0x5f,
+ 0xee, 0x08, 0x0d, 0xbf, 0x48, 0xc0, 0xd1, 0xe5};
+ const uint8_t ctext[] = {0x4d, 0xd2, 0xd0, 0x5a, 0x7b, 0x0d, 0xe9, 0xb2,
+ 0xb9, 0x90, 0x7a, 0xfb, 0x5e, 0xcf, 0x8c, 0xc3};
+ const char key[] = {0x8f, 0x95, 0xf0, 0x92, 0x45, 0x76, 0x5f, 0x80,
+ 0x25, 0x69, 0x34, 0xe5, 0x0c, 0x66, 0x20, 0x7f};
+ uint8_t result[sizeof(ptext)];
+ auto config = LoadBalancerConfig::Create(0, 8, 8, absl::string_view(key, 16));
+ EXPECT_TRUE(config->BlockEncrypt(ptext, result));
+ EXPECT_EQ(memcmp(result, ctext, sizeof(ctext)), 0);
+ EXPECT_TRUE(config->BlockDecrypt(ctext, result));
+ EXPECT_EQ(memcmp(result, ptext, sizeof(ptext)), 0);
+}
+
+} // namespace
+
+} // namespace test
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_decoder.cc b/quiche/quic/load_balancer/load_balancer_decoder.cc
new file mode 100644
index 0000000..ead70d8
--- /dev/null
+++ b/quiche/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 "quiche/quic/load_balancer/load_balancer_decoder.h"
+
+#include "quiche/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/quiche/quic/load_balancer/load_balancer_decoder.h b/quiche/quic/load_balancer/load_balancer_decoder.h
new file mode 100644
index 0000000..34c298c
--- /dev/null
+++ b/quiche/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 "quiche/quic/load_balancer/load_balancer_config.h"
+#include "quiche/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/quiche/quic/load_balancer/load_balancer_decoder_test.cc b/quiche/quic/load_balancer/load_balancer_decoder_test.cc
new file mode 100644
index 0000000..379bacd
--- /dev/null
+++ b/quiche/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 "quiche/quic/load_balancer/load_balancer_decoder.h"
+
+#include "absl/base/macros.h"
+#include "quiche/quic/load_balancer/load_balancer_server_id.h"
+#include "quiche/quic/platform/api/quic_expect_bug.h"
+#include "quiche/quic/platform/api/quic_test.h"
+#include "quiche/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
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.cc b/quiche/quic/load_balancer/load_balancer_encoder.cc
new file mode 100644
index 0000000..852e904
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_encoder.cc
@@ -0,0 +1,178 @@
+// 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 "quiche/quic/load_balancer/load_balancer_encoder.h"
+
+#include "absl/numeric/int128.h"
+#include "quiche/quic/core/quic_connection_id.h"
+#include "quiche/quic/core/quic_data_reader.h"
+#include "quiche/quic/core/quic_data_writer.h"
+#include "quiche/quic/core/quic_packet_number.h"
+#include "quiche/quic/core/quic_types.h"
+#include "quiche/quic/core/quic_utils.h"
+#include "quiche/quic/load_balancer/load_balancer_config.h"
+#include "quiche/quic/platform/api/quic_bug_tracker.h"
+
+namespace quic {
+
+namespace {
+
+// Returns the number of nonces given a certain |nonce_len|.
+absl::uint128 NumberOfNonces(uint8_t nonce_len) {
+ return (static_cast<absl::uint128>(1) << (nonce_len * 8));
+}
+
+// Writes the |size| least significant bytes from |in| to |out| in host byte
+// order. Returns false if |out| does not have enough space.
+bool WriteUint128(const absl::uint128 in, uint8_t size, QuicDataWriter &out) {
+ if (out.remaining() < size) {
+ QUIC_BUG(quic_bug_435375038_05)
+ << "Call to WriteUint128() does not have enough space in |out|";
+ return false;
+ }
+ uint64_t num64 = absl::Uint128Low64(in);
+ if (size <= sizeof(num64)) {
+ out.WriteBytes(&num64, size);
+ } else {
+ out.WriteBytes(&num64, sizeof(num64));
+ num64 = absl::Uint128High64(in);
+ out.WriteBytes(&num64, size - sizeof(num64));
+ }
+ return true;
+}
+
+} // 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) {
+ if (unroutable_connection_id_len == 0 ||
+ unroutable_connection_id_len >
+ kQuicMaxConnectionIdWithLengthPrefixLength) {
+ QUIC_BUG(quic_bug_435375038_01)
+ << "Invalid unroutable_connection_id_len = "
+ << static_cast<int>(unroutable_connection_id_len);
+ return absl::optional<LoadBalancerEncoder>();
+ }
+ return LoadBalancerEncoder(random, visitor, len_self_encoded,
+ unroutable_connection_id_len);
+}
+
+bool LoadBalancerEncoder::UpdateConfig(const LoadBalancerConfig &config,
+ const LoadBalancerServerId server_id) {
+ if (config_.has_value() && config_->config_id() == config.config_id()) {
+ QUIC_BUG(quic_bug_435375038_02)
+ << "Attempting to change config with same ID";
+ return false;
+ }
+ if (server_id.length() != config.server_id_len()) {
+ QUIC_BUG(quic_bug_435375038_03)
+ << "Server ID length " << static_cast<int>(server_id.length())
+ << " does not match configured value of "
+ << static_cast<int>(config.server_id_len());
+ return false;
+ }
+ if (visitor_ != nullptr) {
+ if (config_.has_value()) {
+ visitor_->OnConfigChanged(config_->config_id(), config.config_id());
+ } else {
+ visitor_->OnConfigAdded(config.config_id());
+ }
+ }
+ config_ = config;
+ server_id_ = server_id;
+
+ seed_ = absl::MakeUint128(random_.RandUint64(), random_.RandUint64()) %
+ NumberOfNonces(config.nonce_len());
+ num_nonces_left_ = NumberOfNonces(config.nonce_len());
+ return true;
+}
+
+void LoadBalancerEncoder::DeleteConfig() {
+ if (visitor_ != nullptr && config_.has_value()) {
+ visitor_->OnConfigDeleted(config_->config_id());
+ }
+ config_.reset();
+ server_id_.reset();
+ num_nonces_left_ = 0;
+}
+
+QuicConnectionId LoadBalancerEncoder::GenerateConnectionId() {
+ uint8_t length = (config_.has_value()) ? (config_->total_len() + 1)
+ : unroutable_connection_id_len_;
+ uint8_t config_id = config_.has_value() ? (config_->config_id() << 6)
+ : kLoadBalancerUnroutableConfigId;
+ if (config_.has_value() != server_id_.has_value()) {
+ QUIC_BUG(quic_bug_435375038_04)
+ << "Existence of config and server_id are out of sync";
+ return QuicConnectionId();
+ }
+ uint8_t first_byte;
+ // first byte
+ if (len_self_encoded_) {
+ first_byte = config_id | (length - 1);
+ } else {
+ random_.RandBytes(static_cast<void *>(&first_byte), 1);
+ first_byte = config_id | (first_byte & kLoadBalancerLengthMask);
+ }
+ if (config_id == kLoadBalancerUnroutableConfigId) {
+ return MakeUnroutableConnectionId(first_byte);
+ }
+ QuicConnectionId id;
+ id.set_length(length);
+ QuicDataWriter writer(length, id.mutable_data(), quiche::HOST_BYTE_ORDER);
+ writer.WriteUInt8(first_byte);
+ absl::uint128 next_nonce =
+ (seed_ + num_nonces_left_--) % NumberOfNonces(config_->nonce_len());
+ writer.WriteBytes(server_id_->data().data(), server_id_->length());
+ if (!WriteUint128(next_nonce, config_->nonce_len(), writer)) {
+ return QuicConnectionId();
+ }
+ uint8_t *block_start = reinterpret_cast<uint8_t *>(writer.data() + 1);
+ if (!config_->IsEncrypted()) {
+ // Fill the nonce field with a hash of the Connection ID to avoid the nonce
+ // visibly increasing by one. This would allow observers to correlate
+ // connection IDs as being sequential and likely from the same connection,
+ // not just the same server.
+ absl::uint128 nonce_hash =
+ QuicUtils::FNV1a_128_Hash(absl::string_view(writer.data(), length));
+ QuicDataWriter rewriter(config_->nonce_len(),
+ id.mutable_data() + config_->server_id_len() + 1,
+ quiche::HOST_BYTE_ORDER);
+ if (!WriteUint128(nonce_hash, config_->nonce_len(), rewriter)) {
+ return QuicConnectionId();
+ }
+ } else if (config_->total_len() == kLoadBalancerBlockSize) {
+ // Use one encryption pass.
+ if (!config_->BlockEncrypt(block_start, block_start)) {
+ QUIC_LOG(ERROR) << "Block encryption failed";
+ return QuicConnectionId();
+ }
+ } else {
+ for (uint8_t i = 1; i <= 4; i++) {
+ if (!config_->EncryptionPass(block_start, i)) {
+ QUIC_LOG(ERROR) << "Block encryption failed";
+ return QuicConnectionId();
+ }
+ }
+ }
+ if (num_nonces_left_ == 0) {
+ DeleteConfig();
+ }
+ return id;
+}
+
+QuicConnectionId LoadBalancerEncoder::MakeUnroutableConnectionId(
+ uint8_t first_byte) {
+ QuicConnectionId id;
+ id.set_length(unroutable_connection_id_len_);
+ id.mutable_data()[0] = first_byte;
+ random_.RandBytes(&id.mutable_data()[1], unroutable_connection_id_len_ - 1);
+ return id;
+}
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.h b/quiche/quic/load_balancer/load_balancer_encoder.h
new file mode 100644
index 0000000..98e1cdb
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_encoder.h
@@ -0,0 +1,123 @@
+// 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_ENCODER_H_
+#define QUICHE_QUIC_LOAD_BALANCER_ENCODER_H_
+
+#include "quiche/quic/core/crypto/quic_random.h"
+#include "quiche/quic/load_balancer/load_balancer_config.h"
+#include "quiche/quic/load_balancer/load_balancer_server_id.h"
+
+namespace quic {
+
+namespace test {
+class LoadBalancerEncoderPeer;
+}
+
+// Default length of a 4-tuple connection ID.
+inline constexpr uint8_t kLoadBalancerUnroutableLen = 8;
+
+// Interface which receives notifications when the current config is updated.
+class QUIC_EXPORT_PRIVATE LoadBalancerEncoderVisitorInterface {
+ public:
+ virtual ~LoadBalancerEncoderVisitorInterface() {}
+
+ // Called when a config is added where none existed.
+ //
+ // Connections that support address migration should retire unroutable
+ // connection IDs and replace them with routable ones using the new config,
+ // while avoiding sending a sudden storm of packets containing
+ // RETIRE_CONNECTION_ID and NEW_CONNECTION_ID frames.
+ virtual void OnConfigAdded(const uint8_t config_id) = 0;
+ // Called when the config is changed.
+ //
+ // Existing routable connection IDs should be retired before the decoder stops
+ // supporting that config. The timing of this event is deployment-dependent
+ // and might be tied to the arrival of a new config at the encoder.
+ virtual void OnConfigChanged(const uint8_t old_config_id,
+ const uint8_t new_config_id) = 0;
+ // Called when a config is deleted. The encoder will generate unroutable
+ // connection IDs from now on.
+ //
+ // New connections will not be able to support address migration until a new
+ // config arrives. Existing connections can retain connection IDs that use the
+ // deleted config, which will only become unroutable once the decoder also
+ // deletes it. The time of that deletion is deployment-dependent and might be
+ // tied to the arrival of a new config at the encoder.
+ virtual void OnConfigDeleted(const uint8_t config_id) = 0;
+};
+
+// Manages QUIC-LB configurations to properly encode a given server ID in a
+// QUIC Connection ID.
+class QUIC_EXPORT_PRIVATE LoadBalancerEncoder {
+ public:
+ // Returns a newly created encoder with no active config, if
+ // |unroutable_connection_id_length| is valid. |visitor| specifies an optional
+ // interface to receive callbacks when connection IDs need to be retired.
+ // If |encode_length| is true, then the first byte of any generated
+ // connection ids will encode the length. Otherwise, those bits will be
+ // random. |unroutable_connection_id_length| specifies the length of
+ // connection IDs to be generated when there is no active config. It must not
+ // be 0 and must not be larger than the RFC9000 maximum of 20.
+ static absl::optional<LoadBalancerEncoder> Create(
+ QuicRandom& random, LoadBalancerEncoderVisitorInterface* const visitor,
+ const bool len_self_encoded,
+ const uint8_t unroutable_connection_id_len = kLoadBalancerUnroutableLen);
+
+ // Attempts to replace the current config and server_id with |config| and
+ // |server_id|. If the length |server_id| does not match the server_id_length
+ // of |config| or the ID of |config| matches the ID of the current config,
+ // returns false and leaves the current config unchanged. Otherwise, returns
+ // true. When the encoder runs out of nonces, it will delete the config and
+ // begin generating unroutable connection IDs.
+ bool UpdateConfig(const LoadBalancerConfig& config,
+ const LoadBalancerServerId server_id);
+
+ // Delete the current config and generate unroutable connection IDs from now
+ // on.
+ void DeleteConfig();
+
+ // Returns the number of additional connections IDs that can be generated with
+ // the current config, or 0 if there is no current config.
+ absl::uint128 num_nonces_left() const { return num_nonces_left_; }
+
+ // Returns true if there is an active configuration.
+ bool IsEncoding() const { return config_.has_value(); }
+ // Returns true if there is an active configuration that uses encryption.
+ bool IsEncrypted() const {
+ return config_.has_value() && config_->IsEncrypted();
+ }
+
+ // If there's an active config, generates a connection ID using it. If not,
+ // generates an unroutable connection_id. If there's an error, returns a zero-
+ // length Connection ID.
+ QuicConnectionId GenerateConnectionId();
+
+ private:
+ friend class test::LoadBalancerEncoderPeer;
+
+ LoadBalancerEncoder(QuicRandom& random,
+ LoadBalancerEncoderVisitorInterface* const visitor,
+ const bool len_self_encoded,
+ 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) {}
+
+ QuicConnectionId MakeUnroutableConnectionId(uint8_t first_byte);
+
+ 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_;
+};
+
+} // namespace quic
+
+#endif // QUICHE_QUIC_LOAD_BALANCER_ENCODER_H_
diff --git a/quiche/quic/load_balancer/load_balancer_encoder_test.cc b/quiche/quic/load_balancer/load_balancer_encoder_test.cc
new file mode 100644
index 0000000..ed6db41
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_encoder_test.cc
@@ -0,0 +1,355 @@
+// 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 "quiche/quic/load_balancer/load_balancer_encoder.h"
+
+#include <cstdint>
+
+#include "absl/numeric/int128.h"
+#include "quiche/quic/platform/api/quic_expect_bug.h"
+#include "quiche/quic/platform/api/quic_test.h"
+#include "quiche/quic/test_tools/quic_test_utils.h"
+
+namespace quic {
+
+namespace test {
+
+class LoadBalancerEncoderPeer {
+ public:
+ static void SetNumNoncesLeft(LoadBalancerEncoder &encoder,
+ uint64_t nonces_remaining) {
+ encoder.num_nonces_left_ = absl::uint128(nonces_remaining);
+ }
+};
+
+namespace {
+
+class TestLoadBalancerEncoderVisitor
+ : public LoadBalancerEncoderVisitorInterface {
+ public:
+ ~TestLoadBalancerEncoderVisitor() override {}
+
+ void OnConfigAdded(const uint8_t config_id) override {
+ num_adds_++;
+ current_config_id_ = config_id;
+ }
+
+ void OnConfigChanged(const uint8_t old_config_id,
+ const uint8_t new_config_id) override {
+ num_adds_++;
+ num_deletes_++;
+ EXPECT_EQ(old_config_id, current_config_id_);
+ current_config_id_ = new_config_id;
+ }
+
+ void OnConfigDeleted(const uint8_t config_id) override {
+ EXPECT_EQ(config_id, current_config_id_);
+ current_config_id_.reset();
+ num_deletes_++;
+ }
+
+ uint32_t num_adds() const { return num_adds_; }
+ uint32_t num_deletes() const { return num_deletes_; }
+
+ private:
+ uint32_t num_adds_ = 0, num_deletes_ = 0;
+ absl::optional<uint8_t> current_config_id_ = absl::optional<uint8_t>();
+};
+
+// Allows the caller to specify the exact results in 64-bit chunks.
+class TestRandom : public QuicRandom {
+ public:
+ uint64_t RandUint64() override {
+ if (next_values_.empty()) {
+ return base_;
+ }
+ uint64_t value = next_values_.front();
+ next_values_.pop();
+ return value;
+ }
+
+ void RandBytes(void *data, size_t len) override {
+ size_t written = 0;
+ uint8_t *ptr = static_cast<uint8_t *>(data);
+ while (written < len) {
+ uint64_t result = RandUint64();
+ size_t to_write = (len - written > sizeof(uint64_t)) ? sizeof(uint64_t)
+ : (len - written);
+ memcpy(ptr + written, &result, to_write);
+ written += to_write;
+ }
+ }
+
+ void InsecureRandBytes(void *data, size_t len) override {
+ RandBytes(data, len);
+ }
+
+ uint64_t InsecureRandUint64() override { return RandUint64(); }
+
+ void AddNextValues(uint64_t hi, uint64_t lo) {
+ next_values_.push(hi);
+ next_values_.push(lo);
+ }
+
+ private:
+ std::queue<uint64_t> next_values_;
+ uint64_t base_ = 0xDEADBEEFDEADBEEF;
+};
+
+class LoadBalancerEncoderTest : public QuicTest {
+ public:
+ TestRandom random_;
+};
+
+// Convenience function to shorten the code. Does not check if |array| is long
+// enough or |length| is valid for a server ID.
+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 uint64_t kNonceLow = 0xe5d1c048bf0d08ee;
+constexpr uint64_t kNonceHigh = 0x9321e7e34dde525d;
+constexpr uint8_t kServerId[] = {0xed, 0x79, 0x3a, 0x51, 0xd4, 0x9b, 0x8f, 0x5f,
+ 0xab, 0x65, 0xba, 0x04, 0xc3, 0x33, 0x0a};
+
+TEST_F(LoadBalancerEncoderTest, BadUnroutableLength) {
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(
+ LoadBalancerEncoder::Create(random_, nullptr, false, 0).has_value()),
+ "Invalid unroutable_connection_id_len = 0");
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(
+ LoadBalancerEncoder::Create(random_, nullptr, false, 21).has_value()),
+ "Invalid unroutable_connection_id_len = 21");
+}
+
+TEST_F(LoadBalancerEncoderTest, BadServerIdLength) {
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true);
+ EXPECT_TRUE(encoder.has_value());
+ // Expects a 3 byte server ID and got 4.
+ auto config = LoadBalancerConfig::CreateUnencrypted(1, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 4))),
+ "Server ID length 4 does not match configured value of 3");
+ EXPECT_FALSE(encoder->IsEncoding());
+}
+
+TEST_F(LoadBalancerEncoderTest, FailToUpdateConfigWithSameId) {
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
+ EXPECT_TRUE(encoder.has_value());
+ EXPECT_TRUE(encoder.has_value());
+ auto config = LoadBalancerConfig::CreateUnencrypted(1, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
+ EXPECT_EQ(visitor.num_adds(), 1u);
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3))),
+ "Attempting to change config with same ID");
+ EXPECT_EQ(visitor.num_adds(), 1u);
+}
+
+TEST_F(LoadBalancerEncoderTest, UnencryptedConnectionIdTestVectors) {
+ const uint8_t server_id_lens[] = {3, 8};
+ const uint8_t nonce_lens[] = {4, 5};
+ static_assert(sizeof(server_id_lens) == sizeof(nonce_lens));
+ const std::vector<QuicConnectionId> expected_connection_ids{
+ QuicConnectionId({0x07, 0xed, 0x79, 0x3a, 0x80, 0x49, 0x71, 0x8a}),
+ QuicConnectionId({0x4d, 0xed, 0x79, 0x3a, 0x51, 0xd4, 0x9b, 0x8f, 0x5f,
+ 0xee, 0x15, 0xda, 0x27, 0xc4}),
+ };
+ EXPECT_EQ(sizeof(server_id_lens), expected_connection_ids.size());
+ for (uint8_t i = 0; i < sizeof(server_id_lens); i++) {
+ uint8_t config_id = i % 3;
+ random_.AddNextValues(kNonceHigh, kNonceLow);
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true, 8);
+ EXPECT_TRUE(encoder.has_value());
+ auto config = LoadBalancerConfig::CreateUnencrypted(
+ config_id, server_id_lens[i], nonce_lens[i]);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(encoder->UpdateConfig(
+ *config, MakeServerId(kServerId, server_id_lens[i])));
+ absl::uint128 nonces_left = encoder->num_nonces_left();
+ EXPECT_EQ(encoder->GenerateConnectionId(), expected_connection_ids[i]);
+ EXPECT_EQ(encoder->num_nonces_left(), nonces_left - 1);
+ }
+}
+
+// Follow example in draft-ietf-quic-load-balancers-12.
+TEST_F(LoadBalancerEncoderTest, FollowSpecExample) {
+ const uint8_t config_id = 0, server_id_len = 3, nonce_len = 4;
+ const uint8_t raw_server_id[] = {
+ 0x31,
+ 0x44,
+ 0x1a,
+ };
+ const char raw_key[] = {
+ 0xfd, 0xf7, 0x26, 0xa9, 0x89, 0x3e, 0xc0, 0x5c,
+ 0x06, 0x32, 0xd3, 0x95, 0x66, 0x80, 0xba, 0xf0,
+ };
+ random_.AddNextValues(0, 0x75c2699c);
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true, 8);
+ EXPECT_TRUE(encoder.has_value());
+ auto config = LoadBalancerConfig::Create(config_id, server_id_len, nonce_len,
+ absl::string_view(raw_key));
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(encoder->UpdateConfig(
+ *config, *LoadBalancerServerId::Create(raw_server_id)));
+ EXPECT_TRUE(encoder->IsEncoding());
+ const char raw_connection_id[] = {0x07, 0xe2, 0x3c, 0xb4,
+ 0x2b, 0xba, 0x1e, 0xe0};
+ auto expected =
+ QuicConnectionId(raw_connection_id, 1 + server_id_len + nonce_len);
+ EXPECT_EQ(encoder->GenerateConnectionId(), expected);
+}
+
+// Compare test vectors from Appendix B of draft-ietf-quic-load-balancers-12.
+TEST_F(LoadBalancerEncoderTest, EncoderTestVectors) {
+ // Try (1) the "standard" ConnectionId 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.
+ uint8_t server_id_lens[] = {3, 10, 8, 9};
+ uint8_t nonce_lens[] = {4, 5, 8, 9};
+ static_assert(sizeof(server_id_lens) == sizeof(nonce_lens));
+ const std::vector<QuicConnectionId> expected_connection_ids{
+ QuicConnectionId({0x07, 0xfb, 0xfe, 0x05, 0xf7, 0x31, 0xb4, 0x25}),
+ QuicConnectionId({0x4f, 0x01, 0x09, 0x56, 0xfb, 0x5c, 0x1d, 0x4d, 0x86,
+ 0xe0, 0x10, 0x18, 0x3e, 0x0b, 0x7d, 0x1e}),
+ QuicConnectionId({0x90, 0x4d, 0xd2, 0xd0, 0x5a, 0x7b, 0x0d, 0xe9, 0xb2,
+ 0xb9, 0x90, 0x7a, 0xfb, 0x5e, 0xcf, 0x8c, 0xc3}),
+ QuicConnectionId({0x12, 0x7a, 0x28, 0x5a, 0x09, 0xf8, 0x52, 0x80, 0xf4,
+ 0xfd, 0x6a, 0xbb, 0x43, 0x4a, 0x71, 0x59, 0xe4, 0xd3,
+ 0xeb}),
+ };
+ EXPECT_EQ(sizeof(server_id_lens), expected_connection_ids.size());
+ for (uint8_t i = 0; i < sizeof(server_id_lens); i++) {
+ uint8_t config_id = i % 3;
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true, 8);
+ EXPECT_TRUE(encoder.has_value());
+ auto config = LoadBalancerConfig::Create(config_id, server_id_lens[i],
+ nonce_lens[i], kKey);
+ EXPECT_TRUE(config.has_value());
+ random_.AddNextValues(kNonceHigh, kNonceLow);
+ EXPECT_TRUE(encoder->UpdateConfig(
+ *config, MakeServerId(kServerId, server_id_lens[i])));
+ EXPECT_EQ(encoder->GenerateConnectionId(), expected_connection_ids[i]);
+ }
+}
+
+TEST_F(LoadBalancerEncoderTest, RunOutOfNonces) {
+ const uint8_t config_id = 0, server_id_len = 3, nonce_len = 4;
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true, 8);
+ EXPECT_TRUE(encoder.has_value());
+ auto config =
+ LoadBalancerConfig::Create(config_id, server_id_len, nonce_len, kKey);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(
+ encoder->UpdateConfig(*config, MakeServerId(kServerId, server_id_len)));
+ EXPECT_EQ(visitor.num_adds(), 1u);
+ LoadBalancerEncoderPeer::SetNumNoncesLeft(*encoder, 2);
+ EXPECT_EQ(encoder->num_nonces_left(), 2);
+ EXPECT_EQ(encoder->GenerateConnectionId(),
+ QuicConnectionId({0x07, 0xba, 0x7f, 0x38, 0x2f, 0x5c, 0x22, 0x1d}));
+ EXPECT_EQ(encoder->num_nonces_left(), 1);
+ encoder->GenerateConnectionId();
+ EXPECT_EQ(encoder->IsEncoding(), false);
+ // No retire_calls except for the initial UpdateConfig.
+ EXPECT_EQ(visitor.num_deletes(), 1u);
+}
+
+TEST_F(LoadBalancerEncoderTest, UnroutableConnectionId) {
+ random_.AddNextValues(0x83, kNonceHigh);
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, false);
+ EXPECT_TRUE(encoder.has_value());
+ EXPECT_EQ(encoder->num_nonces_left(), 0);
+ auto connection_id = encoder->GenerateConnectionId();
+ // The first byte is the config_id (0xc0) xored with (0x83 & 0x3f).
+ // The remaining bytes are random, and therefore match kNonceHigh.
+ QuicConnectionId expected({0xc3, 0x5d, 0x52, 0xde, 0x4d, 0xe3, 0xe7, 0x21});
+ EXPECT_EQ(expected, connection_id);
+}
+
+TEST_F(LoadBalancerEncoderTest, NonDefaultUnroutableConnectionIdLength) {
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true, 9);
+ EXPECT_TRUE(encoder.has_value());
+ QuicConnectionId connection_id = encoder->GenerateConnectionId();
+ EXPECT_EQ(connection_id.length(), 9);
+}
+
+TEST_F(LoadBalancerEncoderTest, DeleteConfigWhenNoConfigExists) {
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
+ EXPECT_TRUE(encoder.has_value());
+ encoder->DeleteConfig();
+ EXPECT_EQ(visitor.num_deletes(), 0u);
+}
+
+TEST_F(LoadBalancerEncoderTest, AddConfig) {
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
+ EXPECT_EQ(visitor.num_adds(), 1u);
+ absl::uint128 left = encoder->num_nonces_left();
+ EXPECT_EQ(left, (0x1ull << 32));
+ EXPECT_TRUE(encoder->IsEncoding());
+ EXPECT_FALSE(encoder->IsEncrypted());
+ encoder->GenerateConnectionId();
+ EXPECT_EQ(encoder->num_nonces_left(), left - 1);
+ EXPECT_EQ(visitor.num_deletes(), 0u);
+}
+
+TEST_F(LoadBalancerEncoderTest, UpdateConfig) {
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
+ config = LoadBalancerConfig::Create(1, 4, 4, kKey);
+ EXPECT_TRUE(config.has_value());
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 4)));
+ EXPECT_EQ(visitor.num_adds(), 2u);
+ EXPECT_EQ(visitor.num_deletes(), 1u);
+ EXPECT_TRUE(encoder->IsEncoding());
+ EXPECT_TRUE(encoder->IsEncrypted());
+}
+
+TEST_F(LoadBalancerEncoderTest, DeleteConfig) {
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ TestLoadBalancerEncoderVisitor visitor;
+ auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
+ encoder->DeleteConfig();
+ EXPECT_EQ(visitor.num_adds(), 1u);
+ EXPECT_EQ(visitor.num_deletes(), 1u);
+ EXPECT_FALSE(encoder->IsEncoding());
+ EXPECT_FALSE(encoder->IsEncrypted());
+ EXPECT_EQ(encoder->num_nonces_left(), 0);
+}
+
+TEST_F(LoadBalancerEncoderTest, DeleteConfigNoVisitor) {
+ auto config = LoadBalancerConfig::CreateUnencrypted(0, 3, 4);
+ EXPECT_TRUE(config.has_value());
+ auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true);
+ EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
+ encoder->DeleteConfig();
+ EXPECT_FALSE(encoder->IsEncoding());
+ EXPECT_FALSE(encoder->IsEncrypted());
+ EXPECT_EQ(encoder->num_nonces_left(), 0);
+}
+
+} // namespace
+
+} // namespace test
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_server_id.cc b/quiche/quic/load_balancer/load_balancer_server_id.cc
new file mode 100644
index 0000000..5805b7f
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_server_id.cc
@@ -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.
+
+#include "quiche/quic/load_balancer/load_balancer_server_id.h"
+
+#include "absl/strings/escaping.h"
+#include "absl/types/span.h"
+#include "quiche/quic/core/quic_types.h"
+#include "quiche/quic/platform/api/quic_bug_tracker.h"
+
+namespace quic {
+
+namespace {
+
+// Helper to allow setting the const array during initialization.
+std::array<uint8_t, kLoadBalancerMaxServerIdLen> MakeArray(
+ const absl::Span<const uint8_t> data, const uint8_t length) {
+ std::array<uint8_t, kLoadBalancerMaxServerIdLen> array;
+ memcpy(array.data(), data.data(), length);
+ return array;
+}
+
+} // namespace
+
+absl::optional<LoadBalancerServerId> LoadBalancerServerId::Create(
+ const absl::Span<const uint8_t> data) {
+ if (data.length() == 0 || data.length() > kLoadBalancerMaxServerIdLen) {
+ QUIC_BUG(quic_bug_433312504_01)
+ << "Attempted to create LoadBalancerServerId with length "
+ << data.length();
+ return absl::optional<LoadBalancerServerId>();
+ }
+ return LoadBalancerServerId(data);
+}
+
+std::string LoadBalancerServerId::ToString() const {
+ return absl::BytesToHexString(
+ absl::string_view(reinterpret_cast<const char*>(data_.data()), length_));
+}
+
+LoadBalancerServerId::LoadBalancerServerId(const absl::Span<const uint8_t> data)
+ : data_(MakeArray(data, data.length())), length_(data.length()) {}
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_server_id.h b/quiche/quic/load_balancer/load_balancer_server_id.h
new file mode 100644
index 0000000..039d615
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_server_id.h
@@ -0,0 +1,65 @@
+// 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_SERVER_ID_H_
+#define QUICHE_QUIC_LOAD_BALANCER_SERVER_ID_H_
+
+#include <array>
+
+#include "quiche/quic/core/quic_types.h"
+#include "quiche/quic/platform/api/quic_export.h"
+
+namespace quic {
+
+// The maximum number of bytes in a LoadBalancerServerId.
+inline constexpr uint8_t kLoadBalancerMaxServerIdLen = 15;
+
+// LoadBalancerServerId is the globally understood identifier for a given pool
+// member. It is unique to any given QUIC-LB configuration. See
+// draft-ietf-quic-load-balancers-12.
+// Note: this has nothing to do with QuicServerID. It's an unfortunate collision
+// between an internal term for the destination identifiers for a particular
+// deployment (QuicServerID) and the object of a load balancing decision
+// (LoadBalancerServerId).
+class QUIC_EXPORT_PRIVATE LoadBalancerServerId {
+ public:
+ // Copies all the bytes from |data| into a new LoadBalancerServerId.
+ static absl::optional<LoadBalancerServerId> Create(
+ const absl::Span<const uint8_t> data);
+
+ // Server IDs are opaque bytes, but defining these operators allows us to sort
+ // them into a tree and define ranges.
+ bool operator<(const LoadBalancerServerId& other) const {
+ return data() < other.data();
+ }
+ bool operator==(const LoadBalancerServerId& other) const {
+ return data() == other.data();
+ }
+
+ // Hash function to allow use as a key in unordered maps.
+ template <typename H>
+ friend H AbslHashValue(H h, const LoadBalancerServerId& server_id) {
+ return H::combine_contiguous(std::move(h), server_id.data().data(),
+ server_id.length());
+ }
+
+ absl::Span<const uint8_t> data() const {
+ return absl::MakeConstSpan(data_.data(), length_);
+ }
+ uint8_t length() const { return length_; }
+
+ // Returns the server ID in hex format.
+ std::string ToString() const;
+
+ private:
+ // The constructor is private because it can't validate the input.
+ LoadBalancerServerId(const absl::Span<const uint8_t> data);
+
+ std::array<uint8_t, kLoadBalancerMaxServerIdLen> data_;
+ uint8_t length_;
+};
+
+} // namespace quic
+
+#endif // QUICHE_QUIC_LOAD_BALANCER_SERVER_ID_H_
diff --git a/quiche/quic/load_balancer/load_balancer_server_id_map.h b/quiche/quic/load_balancer/load_balancer_server_id_map.h
new file mode 100644
index 0000000..02671bb
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_server_id_map.h
@@ -0,0 +1,103 @@
+// 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_SERVER_ID_MAP_H_
+#define QUICHE_QUIC_LOAD_BALANCER_LOAD_BALANCER_SERVER_ID_MAP_H_
+
+#include "absl/container/flat_hash_map.h"
+#include "quiche/quic/load_balancer/load_balancer_server_id.h"
+#include "quiche/quic/platform/api/quic_bug_tracker.h"
+
+namespace quic {
+
+// This class wraps an absl::flat_hash_map which associates server IDs to an
+// arbitrary type T. It validates that all server ids are of the same fixed
+// length.
+template <typename T>
+class QUIC_EXPORT_PRIVATE LoadBalancerServerIdMap {
+ public:
+ // Returns a newly created pool for server IDs of length |server_id_len|, or
+ // nullptr if |server_id_len| is invalid.
+ static std::shared_ptr<LoadBalancerServerIdMap> Create(
+ const uint8_t server_id_len);
+
+ // Returns the entry associated with |server_id|, if present. For small |T|,
+ // use Lookup. For large |T|, use LookupNoCopy.
+ absl::optional<const T> Lookup(const LoadBalancerServerId server_id) const;
+ const T* LookupNoCopy(const LoadBalancerServerId server_id) const;
+
+ // Updates the table so that |value| is associated with |server_id|. Sets
+ // QUIC_BUG if the length is incorrect for this map.
+ void AddOrReplace(const LoadBalancerServerId server_id, T value);
+
+ // Removes the entry associated with |server_id|.
+ void Erase(const LoadBalancerServerId server_id) {
+ server_id_table_.erase(server_id);
+ }
+
+ uint8_t server_id_len() const { return server_id_len_; }
+
+ private:
+ LoadBalancerServerIdMap(uint8_t server_id_len)
+ : server_id_len_(server_id_len) {}
+
+ const uint8_t server_id_len_; // All server IDs must be of this length.
+ absl::flat_hash_map<LoadBalancerServerId, T> server_id_table_;
+};
+
+template <typename T>
+std::shared_ptr<LoadBalancerServerIdMap<T>> LoadBalancerServerIdMap<T>::Create(
+ const uint8_t server_id_len) {
+ if (server_id_len == 0 || server_id_len > kLoadBalancerMaxServerIdLen) {
+ QUIC_BUG(quic_bug_434893339_01)
+ << "Tried to configure map with server ID length "
+ << static_cast<int>(server_id_len);
+ return nullptr;
+ }
+ return std::make_shared<LoadBalancerServerIdMap<T>>(
+ LoadBalancerServerIdMap(server_id_len));
+}
+
+template <typename T>
+absl::optional<const T> LoadBalancerServerIdMap<T>::Lookup(
+ const LoadBalancerServerId server_id) const {
+ if (server_id.length() != server_id_len_) {
+ QUIC_BUG(quic_bug_434893339_02)
+ << "Lookup with a " << static_cast<int>(server_id.length())
+ << " byte server ID, map requires " << static_cast<int>(server_id_len_);
+ return absl::optional<T>();
+ }
+ auto it = server_id_table_.find(server_id);
+ return (it != server_id_table_.end()) ? it->second
+ : absl::optional<const T>();
+}
+
+template <typename T>
+const T* LoadBalancerServerIdMap<T>::LookupNoCopy(
+ const LoadBalancerServerId server_id) const {
+ if (server_id.length() != server_id_len_) {
+ QUIC_BUG(quic_bug_434893339_02)
+ << "Lookup with a " << static_cast<int>(server_id.length())
+ << " byte server ID, map requires " << static_cast<int>(server_id_len_);
+ return nullptr;
+ }
+ auto it = server_id_table_.find(server_id);
+ return (it != server_id_table_.end()) ? &it->second : nullptr;
+}
+
+template <typename T>
+void LoadBalancerServerIdMap<T>::AddOrReplace(
+ const LoadBalancerServerId server_id, T value) {
+ if (server_id.length() == server_id_len_) {
+ server_id_table_[server_id] = value;
+ } else {
+ QUIC_BUG(quic_bug_434893339_03)
+ << "Server ID of " << static_cast<int>(server_id.length())
+ << " bytes; this map requires " << static_cast<int>(server_id_len_);
+ }
+}
+
+} // namespace quic
+
+#endif // QUICHE_QUIC_LOAD_BALANCER_LOAD_BALANCER_SERVER_ID_MAP_H_
diff --git a/quiche/quic/load_balancer/load_balancer_server_id_map_test.cc b/quiche/quic/load_balancer/load_balancer_server_id_map_test.cc
new file mode 100644
index 0000000..d8fcdac
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_server_id_map_test.cc
@@ -0,0 +1,94 @@
+// 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 "quiche/quic/load_balancer/load_balancer_server_id_map.h"
+
+#include "quiche/quic/platform/api/quic_expect_bug.h"
+#include "quiche/quic/platform/api/quic_test.h"
+#include "quiche/quic/test_tools/quic_test_utils.h"
+
+namespace quic {
+
+namespace test {
+
+namespace {
+
+constexpr uint8_t kServerId[] = {0xed, 0x79, 0x3a, 0x51};
+
+class LoadBalancerServerIdMapTest : public QuicTest {
+ public:
+ const LoadBalancerServerId valid_server_id_ =
+ *LoadBalancerServerId::Create(kServerId);
+ const LoadBalancerServerId invalid_server_id_ =
+ *LoadBalancerServerId::Create(absl::Span<const uint8_t>(kServerId, 3));
+};
+
+TEST_F(LoadBalancerServerIdMapTest, CreateWithBadServerIdLength) {
+ EXPECT_QUIC_BUG(EXPECT_EQ(LoadBalancerServerIdMap<int>::Create(0), nullptr),
+ "Tried to configure map with server ID length 0");
+ EXPECT_QUIC_BUG(EXPECT_EQ(LoadBalancerServerIdMap<int>::Create(16), nullptr),
+ "Tried to configure map with server ID length 16");
+}
+
+TEST_F(LoadBalancerServerIdMapTest, AddOrReplaceWithBadServerIdLength) {
+ int record = 1;
+ auto pool = LoadBalancerServerIdMap<int>::Create(4);
+ EXPECT_NE(pool, nullptr);
+ EXPECT_QUIC_BUG(pool->AddOrReplace(invalid_server_id_, record),
+ "Server ID of 3 bytes; this map requires 4");
+}
+
+TEST_F(LoadBalancerServerIdMapTest, LookupWithBadServerIdLength) {
+ int record = 1;
+ auto pool = LoadBalancerServerIdMap<int>::Create(4);
+ EXPECT_NE(pool, nullptr);
+ pool->AddOrReplace(valid_server_id_, record);
+ EXPECT_QUIC_BUG(EXPECT_FALSE(pool->Lookup(invalid_server_id_).has_value()),
+ "Lookup with a 3 byte server ID, map requires 4");
+ EXPECT_QUIC_BUG(EXPECT_EQ(pool->LookupNoCopy(invalid_server_id_), nullptr),
+ "Lookup with a 3 byte server ID, map requires 4");
+}
+
+TEST_F(LoadBalancerServerIdMapTest, LookupWhenEmpty) {
+ auto pool = LoadBalancerServerIdMap<int>::Create(4);
+ EXPECT_NE(pool, nullptr);
+ EXPECT_EQ(pool->LookupNoCopy(valid_server_id_), nullptr);
+ absl::optional<int> result = pool->Lookup(valid_server_id_);
+ EXPECT_FALSE(result.has_value());
+}
+
+TEST_F(LoadBalancerServerIdMapTest, AddLookup) {
+ int record1 = 1, record2 = 2;
+ auto pool = LoadBalancerServerIdMap<int>::Create(4);
+ EXPECT_NE(pool, nullptr);
+ auto other_server_id = LoadBalancerServerId::Create({0x01, 0x02, 0x03, 0x04});
+ EXPECT_TRUE(other_server_id.has_value());
+ pool->AddOrReplace(valid_server_id_, record1);
+ pool->AddOrReplace(*other_server_id, record2);
+ absl::optional<int> result = pool->Lookup(valid_server_id_);
+ EXPECT_TRUE(result.has_value());
+ EXPECT_EQ(*result, record1);
+ auto result_ptr = pool->LookupNoCopy(valid_server_id_);
+ EXPECT_NE(result_ptr, nullptr);
+ EXPECT_EQ(*result_ptr, record1);
+ result = pool->Lookup(*other_server_id);
+ EXPECT_TRUE(result.has_value());
+ EXPECT_EQ(*result, record2);
+}
+
+TEST_F(LoadBalancerServerIdMapTest, AddErase) {
+ int record = 1;
+ auto pool = LoadBalancerServerIdMap<int>::Create(4);
+ EXPECT_NE(pool, nullptr);
+ pool->AddOrReplace(valid_server_id_, record);
+ EXPECT_EQ(*pool->LookupNoCopy(valid_server_id_), record);
+ pool->Erase(valid_server_id_);
+ EXPECT_EQ(pool->LookupNoCopy(valid_server_id_), nullptr);
+}
+
+} // namespace
+
+} // namespace test
+
+} // namespace quic
diff --git a/quiche/quic/load_balancer/load_balancer_server_id_test.cc b/quiche/quic/load_balancer/load_balancer_server_id_test.cc
new file mode 100644
index 0000000..ab15094
--- /dev/null
+++ b/quiche/quic/load_balancer/load_balancer_server_id_test.cc
@@ -0,0 +1,111 @@
+// 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 "quiche/quic/load_balancer/load_balancer_server_id.h"
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Weverything"
+
+#include "absl/hash/hash_testing.h"
+
+#pragma clang diagnostic pop
+
+#include "quiche/quic/platform/api/quic_expect_bug.h"
+#include "quiche/quic/platform/api/quic_test.h"
+#include "quiche/quic/test_tools/quic_test_utils.h"
+
+namespace quic {
+
+namespace test {
+
+namespace {
+
+class LoadBalancerServerIdTest : public QuicTest {};
+
+constexpr uint8_t kRawServerId[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
+ 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b,
+ 0x0c, 0x0d, 0x0e, 0x0f};
+
+TEST_F(LoadBalancerServerIdTest, CreateReturnsNullIfTooLong) {
+ EXPECT_QUIC_BUG(EXPECT_FALSE(LoadBalancerServerId::Create(
+ absl::Span<const uint8_t>(kRawServerId, 16))
+ .has_value()),
+ "Attempted to create LoadBalancerServerId with length 16");
+ EXPECT_QUIC_BUG(
+ EXPECT_FALSE(LoadBalancerServerId::Create(absl::Span<const uint8_t>())
+ .has_value()),
+ "Attempted to create LoadBalancerServerId with length 0");
+}
+
+TEST_F(LoadBalancerServerIdTest, CompareIdenticalExceptLength) {
+ auto server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 15));
+ ASSERT_TRUE(server_id.has_value());
+ EXPECT_EQ(server_id->length(), 15);
+ auto shorter_server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 5));
+ ASSERT_TRUE(shorter_server_id.has_value());
+ EXPECT_EQ(shorter_server_id->length(), 5);
+ // Shorter comes before longer if all bits match
+ EXPECT_TRUE(shorter_server_id < server_id);
+ EXPECT_FALSE(server_id < shorter_server_id);
+ // Different lengths are never equal.
+ EXPECT_FALSE(shorter_server_id == server_id);
+}
+
+TEST_F(LoadBalancerServerIdTest, AccessorFunctions) {
+ auto server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 5));
+ EXPECT_TRUE(server_id.has_value());
+ EXPECT_EQ(server_id->length(), 5);
+ EXPECT_EQ(memcmp(server_id->data().data(), kRawServerId, 5), 0);
+ EXPECT_EQ(server_id->ToString(), "0001020304");
+}
+
+TEST_F(LoadBalancerServerIdTest, CompareDifferentServerIds) {
+ auto server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 5));
+ ASSERT_TRUE(server_id.has_value());
+ auto reverse = LoadBalancerServerId::Create({0x0f, 0x0e, 0x0d, 0x0c, 0x0b});
+ ASSERT_TRUE(reverse.has_value());
+ EXPECT_TRUE(server_id < reverse);
+ auto long_server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 15));
+ EXPECT_TRUE(long_server_id < reverse);
+}
+
+TEST_F(LoadBalancerServerIdTest, EqualityOperators) {
+ auto server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 15));
+ ASSERT_TRUE(server_id.has_value());
+ auto shorter_server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 5));
+ ASSERT_TRUE(shorter_server_id.has_value());
+ EXPECT_FALSE(server_id == shorter_server_id);
+ auto server_id2 = server_id;
+ EXPECT_TRUE(server_id == server_id2);
+}
+
+TEST_F(LoadBalancerServerIdTest, SupportsHash) {
+ auto server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 15));
+ ASSERT_TRUE(server_id.has_value());
+ auto shorter_server_id =
+ LoadBalancerServerId::Create(absl::Span<const uint8_t>(kRawServerId, 5));
+ ASSERT_TRUE(shorter_server_id.has_value());
+ auto different_server_id =
+ LoadBalancerServerId::Create({0x0f, 0x0e, 0x0d, 0x0c, 0x0b});
+ ASSERT_TRUE(different_server_id.has_value());
+ EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({
+ *server_id,
+ *shorter_server_id,
+ *different_server_id,
+ }));
+}
+
+} // namespace
+
+} // namespace test
+
+} // namespace quic
