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