OHTTP key config: use first supported key config

With this change, clients will now select the first key they parsed instead of a random one.

PiperOrigin-RevId: 947502008
diff --git a/quiche/oblivious_http/common/oblivious_http_header_key_config.cc b/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
index 0386182..c1c7e86 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
@@ -308,20 +308,28 @@
     std::optional<absl::string_view> /*media_type*/) {
   ConfigMap configs;
   PublicKeyMap keys;
+  std::optional<uint8_t> first_key_id;
   // First, try to parse the keys using the length-prefixed format from RFC
   // 9458.
-  if (ReadKeyConfigsWithLengthPrefix(key_config, configs, keys).ok()) {
-    return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys));
+  if (ReadKeyConfigsWithLengthPrefix(key_config, configs, keys, first_key_id)
+          .ok()) {
+    return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys),
+                                   first_key_id);
   }
   // Otherwise, try parsing using the non-length-prefixed format from
   // draft-ietf-ohai-ohttp-08, a precursor to RFC 9458.
   configs.clear();
   keys.clear();
+  first_key_id.reset();
   QuicheDataReader reader(key_config);
   while (!reader.IsDoneReading()) {
     QUICHE_RETURN_IF_ERROR(ReadSingleKeyConfig(reader, configs, keys));
+    if (!first_key_id && !configs.empty()) {
+      first_key_id = configs.begin()->first;
+    }
   }
-  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys));
+  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys),
+                                 first_key_id);
 }
 
 absl::StatusOr<ObliviousHttpKeyConfigs> ObliviousHttpKeyConfigs::Create(
@@ -382,6 +390,12 @@
 ObliviousHttpHeaderKeyConfig ObliviousHttpKeyConfigs::PreferredConfig() const {
   // configs_ is forced to have at least one object during construction.
   QUICHE_CHECK(!configs_.empty());
+  if (first_key_id_) {
+    auto it = configs_.find(*first_key_id_);
+    if (it != configs_.end()) {
+      return it->second.front();
+    }
+  }
   return configs_.begin()->second.front();
 }
 
@@ -461,7 +475,8 @@
 
 // static
 absl::Status ObliviousHttpKeyConfigs::ReadKeyConfigsWithLengthPrefix(
-    absl::string_view key_configs, ConfigMap& configs, PublicKeyMap& keys) {
+    absl::string_view key_configs, ConfigMap& configs, PublicKeyMap& keys,
+    std::optional<uint8_t>& first_key_id) {
   QuicheDataReader reader(key_configs);
   while (!reader.IsDoneReading()) {
     absl::string_view single_key_config;
@@ -472,6 +487,9 @@
     QuicheDataReader single_reader(single_key_config);
     QUICHE_RETURN_IF_ERROR(ReadSingleKeyConfig(single_reader, configs, keys,
                                                /*skip_unknown_kems=*/true));
+    if (!first_key_id && !configs.empty()) {
+      first_key_id = configs.begin()->first;
+    }
   }
   if (configs.empty() || keys.empty()) {
     return absl::InvalidArgumentError("No supported key configs found");
diff --git a/quiche/oblivious_http/common/oblivious_http_header_key_config.h b/quiche/oblivious_http/common/oblivious_http_header_key_config.h
index 3e5aab0..1612f9f 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config.h
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config.h
@@ -214,12 +214,10 @@
 
   int NumKeys() const { return public_keys_.size(); }
 
-  // Returns a preferred config to use.  The preferred key is the key with
-  // the highest key_id.  If more than one configuration exists for the
-  // preferred key any configuration may be returned.
-  //
-  // These methods are useful in the (common) case where only one key
-  // configuration is supported by the server.
+  // Returns a preferred config to use. When the keys were parsed using
+  // `ParseConcatenatedKeys()`, the preferred key is the first suppported key.
+  // Otherwise, if more than one configuration exists for the preferred key,
+  // any configuration may be returned.
   ObliviousHttpHeaderKeyConfig PreferredConfig() const;
 
   absl::StatusOr<absl::string_view> GetPublicKeyForId(uint8_t key_id) const;
@@ -250,8 +248,11 @@
       absl::btree_map<uint8_t, std::vector<ObliviousHttpHeaderKeyConfig>,
                       std::greater<uint8_t>>;
 
-  ObliviousHttpKeyConfigs(ConfigMap cm, PublicKeyMap km)
-      : configs_(std::move(cm)), public_keys_(std::move(km)) {}
+  ObliviousHttpKeyConfigs(ConfigMap cm, PublicKeyMap km,
+                          std::optional<uint8_t> first_key_id = std::nullopt)
+      : configs_(std::move(cm)),
+        public_keys_(std::move(km)),
+        first_key_id_(first_key_id) {}
 
   static absl::Status ReadSingleKeyConfig(QuicheDataReader& reader,
                                           ConfigMap& configs,
@@ -262,13 +263,18 @@
   // https://www.rfc-editor.org/rfc/rfc9458.html#section-3.1-2. Note that this
   // can leave `configs` and `keys` in an invalid state when returning an error.
   static absl::Status ReadKeyConfigsWithLengthPrefix(
-      absl::string_view key_configs, ConfigMap& configs, PublicKeyMap& keys);
+      absl::string_view key_configs, ConfigMap& configs, PublicKeyMap& keys,
+      std::optional<uint8_t>& first_key_id);
 
   // A mapping from key_id to ObliviousHttpHeaderKeyConfig objects for that key.
   const ConfigMap configs_;
 
   // A mapping from key_id to the public key for that key_id.
   const PublicKeyMap public_keys_;
+
+  // The first supported key_id found when parsing. Only set when parsing a
+  // concatenated key string.
+  const std::optional<uint8_t> first_key_id_;
 };
 
 // Human-readable strings suitable for logging.
diff --git a/quiche/oblivious_http/common/oblivious_http_header_key_config_test.cc b/quiche/oblivious_http/common/oblivious_http_header_key_config_test.cc
index 7b315f7..abe0edc 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config_test.cc
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config_test.cc
@@ -467,8 +467,8 @@
               IsOkAndHolds(Property(&ObliviousHttpKeyConfigs::NumKeys, 2)));
   EXPECT_THAT(
       ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key)->PreferredConfig(),
-      AllOf(HasKeyId(0x4f), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
-            HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_128_GCM)));
+      AllOf(HasKeyId(0x4b), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+            HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_256_GCM)));
 }
 
 TEST(ObliviousHttpKeyConfigs, TwoSimilarKeyConfigsWithLengthPrefix) {
@@ -495,10 +495,108 @@
               IsOkAndHolds(Property(&ObliviousHttpKeyConfigs::NumKeys, 2)));
   EXPECT_THAT(
       ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key)->PreferredConfig(),
+      AllOf(HasKeyId(0x4b), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+            HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_256_GCM)));
+}
+
+TEST(ObliviousHttpKeyConfigs, TwoSimilarKeyConfigsReverseOrder) {
+  std::string key;
+  ASSERT_TRUE(absl::HexStringToBytes(
+      // First key config.
+      "4f"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010001"                          // HKDF_SHA256, AES_128_GCM
+      // Second key config.
+      "4b"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010002",                         // HKDF_SHA256, AES_256_GCM
+      &key));
+  EXPECT_THAT(ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key),
+              IsOkAndHolds(Property(&ObliviousHttpKeyConfigs::NumKeys, 2)));
+  EXPECT_THAT(
+      ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key)->PreferredConfig(),
       AllOf(HasKeyId(0x4f), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
             HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_128_GCM)));
 }
 
+TEST(ObliviousHttpKeyConfigs,
+     TwoSimilarKeyConfigsReverseOrderWithLengthPrefix) {
+  std::string key;
+  ASSERT_TRUE(absl::HexStringToBytes(
+      // First key config.
+      "0029"                              // length of this key config
+      "4f"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010001"                          // HKDF_SHA256, AES_128_GCM
+      // Second key config.
+      "0029"                              // length of this key config
+      "4b"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010002",                         // HKDF_SHA256, AES_256_GCM
+      &key));
+  EXPECT_THAT(ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key),
+              IsOkAndHolds(Property(&ObliviousHttpKeyConfigs::NumKeys, 2)));
+  EXPECT_THAT(
+      ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key)->PreferredConfig(),
+      AllOf(HasKeyId(0x4f), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+            HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_128_GCM)));
+}
+
+TEST(ObliviousHttpKeyConfigs, UnsupportedKemWithLengthPrefix) {
+  std::string key;
+  ASSERT_TRUE(absl::HexStringToBytes(
+      // First key config.
+      "0039"                              // length of this key config
+      "33"                                // key_id
+      "FFFF"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "808182838485868888898a8b8c8d8e8f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010001"                          // HKDF_SHA256, AES_128_GCM
+      // Second key config.
+      "003d"                              // length of this key config
+      "4b"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0018"                              // len(symmetric_algorithms)
+      "DEAD0002"                          // Unsupported, AES_256_GCM
+      "0001BEEF"                          // HKDF_SHA256, Unsupported
+      "DEADBEEF"                          // Unsupported, Unsupported
+      "00020003"                          // HKDF_SHA384, CHACHA20_POLY1305
+      "00010001"                          // HKDF_SHA256, AES_128_GCM
+      "00010002"                          // HKDF_SHA256, AES_256_GCM
+      // Third key config.
+      "0029"                              // length of this key config
+      "4f"                                // key_id
+      "0020"                              // kem_id
+      "606162636465666768696a6b6c6d6e6f"  // public_key
+      "707172737475767778797a7b7c7d7e7f"  // public_key
+      "0004"                              // len(symmetric_algorithms)
+      "00010001",                         // HKDF_SHA256, AES_128_GCM
+      &key));
+  EXPECT_THAT(ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key),
+              IsOkAndHolds(Property(&ObliviousHttpKeyConfigs::NumKeys, 2)));
+  EXPECT_THAT(
+      ObliviousHttpKeyConfigs::ParseConcatenatedKeys(key)->PreferredConfig(),
+      AllOf(HasKeyId(0x4b), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+            HasKdfId(EVP_HPKE_HKDF_SHA384),
+            HasAeadId(EVP_HPKE_CHACHA20_POLY1305)));
+}
+
 TEST(ObliviousHttpKeyConfigs, RFCExample) {
   std::string key;
   ASSERT_TRUE(absl::HexStringToBytes(