Class that contains QUIC-LB configuration common to both servers and load balancers.

PiperOrigin-RevId: 434877105
diff --git a/quic/load_balancer/load_balancer_config.cc b/quic/load_balancer/load_balancer_config.cc
new file mode 100644
index 0000000..a22bf3f
--- /dev/null
+++ b/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 "quic/load_balancer/load_balancer_config.h"
+
+#include <memory>
+#include <string_view>
+
+#include "third_party/boringssl/src/include/openssl/aes.h"
+#include "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/quic/load_balancer/load_balancer_config.h b/quic/load_balancer/load_balancer_config.h
new file mode 100644
index 0000000..ac9967d
--- /dev/null
+++ b/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 "quic/core/quic_types.h"
+#include "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);
+
+  const uint8_t config_id_;
+  const uint8_t server_id_len_;
+  const 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.
+  const 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.
+  const absl::optional<AES_KEY> block_decrypt_key_;
+};
+
+}  // namespace quic
+
+#endif  // QUICHE_QUIC_LOAD_BALANCER_CONFIG_H
diff --git a/quic/load_balancer/load_balancer_config_test.cc b/quic/load_balancer/load_balancer_config_test.cc
new file mode 100644
index 0000000..06891d2
--- /dev/null
+++ b/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 "quic/load_balancer/load_balancer_config.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 {
+
+inline 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