OHTTP key config: maintain ordering

This CL changes ObliviousHttpKeyConfigs to maintain the ordering of key configs. If a config1 is passed in to Create before config2, it will now be serialized ahead as well. This allow to use ordering to communicate which config is preferred.

PiperOrigin-RevId: 948251030
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 c1c7e86..48ffce0 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config.cc
@@ -12,7 +12,6 @@
 
 #include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/escaping.h"
@@ -308,42 +307,49 @@
     std::optional<absl::string_view> /*media_type*/) {
   ConfigMap configs;
   PublicKeyMap keys;
-  std::optional<uint8_t> first_key_id;
+  std::vector<uint8_t> key_ids;
   // First, try to parse the keys using the length-prefixed format from RFC
   // 9458.
-  if (ReadKeyConfigsWithLengthPrefix(key_config, configs, keys, first_key_id)
-          .ok()) {
+  if (ReadKeyConfigsWithLengthPrefix(key_config, configs, keys, key_ids).ok()) {
     return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys),
-                                   first_key_id);
+                                   key_ids);
   }
   // 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();
+  key_ids.clear();
   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;
+    uint8_t key_id;
+    QUICHE_RETURN_IF_ERROR(ReadSingleKeyConfig(reader, configs, keys, key_id,
+                                               /*skip_unknown_kems=*/false));
+    if (!configs.empty() && (key_ids.empty() || key_ids.back() != key_id)) {
+      key_ids.push_back(key_id);
     }
   }
-  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys),
-                                 first_key_id);
+  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys), key_ids);
 }
 
 absl::StatusOr<ObliviousHttpKeyConfigs> ObliviousHttpKeyConfigs::Create(
-    absl::flat_hash_set<OhttpKeyConfig> ohttp_key_configs) {
+    std::vector<OhttpKeyConfig> ohttp_key_configs) {
   if (ohttp_key_configs.empty()) {
     return absl::InvalidArgumentError("Empty input");
   }
   ConfigMap configs_map;
   PublicKeyMap keys_map;
+  std::vector<uint8_t> key_ids;
+  key_ids.reserve(ohttp_key_configs.size());
   for (OhttpKeyConfig ohttp_key_config : ohttp_key_configs) {
+    const uint8_t key_id = ohttp_key_config.key_id;
     QUICHE_RETURN_IF_ERROR(StoreKeyConfigIfValid(std::move(ohttp_key_config),
                                                  configs_map, keys_map));
+    if (key_ids.empty() || key_ids.back() != key_id) {
+      key_ids.push_back(key_id);
+    }
   }
-  return ObliviousHttpKeyConfigs(std::move(configs_map), std::move(keys_map));
+  return ObliviousHttpKeyConfigs(std::move(configs_map), std::move(keys_map),
+                                 std::move(key_ids));
 }
 
 absl::StatusOr<ObliviousHttpKeyConfigs> ObliviousHttpKeyConfigs::Create(
@@ -361,15 +367,22 @@
   ConfigMap configs;
   PublicKeyMap keys;
   uint8_t key_id = single_key_config.GetKeyId();
+  std::vector<uint8_t> key_ids = {key_id};
   keys.emplace(key_id, public_key);
   configs[key_id].emplace_back(std::move(single_key_config));
-  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys));
+  return ObliviousHttpKeyConfigs(std::move(configs), std::move(keys),
+                                 std::move(key_ids));
 }
 
 absl::StatusOr<std::string> ObliviousHttpKeyConfigs::GenerateConcatenatedKeys(
     bool with_length_prefix) const {
   std::string concatenated_keys;
-  for (const auto& [key_id, ohttp_configs] : configs_) {
+  for (const uint8_t key_id : key_ids_) {
+    const auto it = configs_.find(key_id);
+    if (it == configs_.end()) {
+      continue;
+    }
+    const std::vector<ObliviousHttpHeaderKeyConfig>& ohttp_configs = it->second;
     QUICHE_ASSIGN_OR_RETURN(absl::string_view public_key,
                             GetPublicKeyForId(key_id));
     QUICHE_ASSIGN_OR_RETURN(
@@ -390,8 +403,8 @@
 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 (!key_ids_.empty()) {
+    auto it = configs_.find(key_ids_.front());
     if (it != configs_.end()) {
       return it->second.front();
     }
@@ -411,11 +424,12 @@
 
 absl::Status ObliviousHttpKeyConfigs::ReadSingleKeyConfig(
     QuicheDataReader& reader, ConfigMap& configs, PublicKeyMap& keys,
-    bool skip_unknown_kems) {
-  uint8_t key_id;
-  if (!reader.ReadUInt8(&key_id)) {
+    uint8_t& key_id, bool skip_unknown_kems) {
+  uint8_t key_id2;
+  if (!reader.ReadUInt8(&key_id2)) {
     return absl::InvalidArgumentError("Failed to read key_id");
   }
+  key_id = key_id2;
   uint16_t kem_id;
   if (!reader.ReadUInt16(&kem_id)) {
     return absl::InvalidArgumentError("Failed to read kem_id");
@@ -476,7 +490,7 @@
 // static
 absl::Status ObliviousHttpKeyConfigs::ReadKeyConfigsWithLengthPrefix(
     absl::string_view key_configs, ConfigMap& configs, PublicKeyMap& keys,
-    std::optional<uint8_t>& first_key_id) {
+    std::vector<uint8_t>& key_ids) {
   QuicheDataReader reader(key_configs);
   while (!reader.IsDoneReading()) {
     absl::string_view single_key_config;
@@ -485,10 +499,12 @@
           "Failed to read length-prefixed key config");
     }
     QuicheDataReader single_reader(single_key_config);
+    uint8_t key_id;
     QUICHE_RETURN_IF_ERROR(ReadSingleKeyConfig(single_reader, configs, keys,
+                                               key_id,
                                                /*skip_unknown_kems=*/true));
-    if (!first_key_id && !configs.empty()) {
-      first_key_id = configs.begin()->first;
+    if (!configs.empty() && (key_ids.empty() || key_ids.back() != key_id)) {
+      key_ids.push_back(key_id);
     }
   }
   if (configs.empty() || keys.empty()) {
@@ -569,7 +585,16 @@
 
 std::string ObliviousHttpKeyConfigs::DebugString() const {
   std::string s;
-  for (const auto& [key_id, ohttp_configs] : configs_) {
+  for (const uint8_t key_id : key_ids_) {
+    const auto kit = configs_.find(key_id);
+    if (kit == configs_.end()) {
+      continue;
+    }
+    const std::vector<ObliviousHttpHeaderKeyConfig>& ohttp_configs =
+        kit->second;
+    if (!s.empty()) {
+      absl::StrAppend(&s, "\n");
+    }
     absl::StrAppend(&s, "[key_id: ", static_cast<uint16_t>(key_id), ", {");
     for (const ObliviousHttpHeaderKeyConfig& ohttp_config : ohttp_configs) {
       absl::StrAppend(&s, "\n  ", ohttp_config.DebugString());
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 1612f9f..8517e29 100644
--- a/quiche/oblivious_http/common/oblivious_http_header_key_config.h
+++ b/quiche/oblivious_http/common/oblivious_http_header_key_config.h
@@ -11,7 +11,6 @@
 
 #include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -155,7 +154,7 @@
     uint8_t key_id;
     uint16_t kem_id;
     std::string public_key;  // Raw byte string.
-    absl::flat_hash_set<SymmetricAlgorithmsConfig> symmetric_algorithms;
+    std::vector<SymmetricAlgorithmsConfig> symmetric_algorithms;
 
     bool operator==(const OhttpKeyConfig& other) const {
       return key_id == other.key_id && kem_id == other.kem_id &&
@@ -197,7 +196,7 @@
   // keys], use `GenerateConcatenatedKeys()`. This output can inturn be parsed
   // by `ObliviousHttpKeyConfigs::ParseConcatenatedKeys` on client side.
   static absl::StatusOr<ObliviousHttpKeyConfigs> Create(
-      absl::flat_hash_set<OhttpKeyConfig> ohttp_key_configs);
+      std::vector<OhttpKeyConfig> ohttp_key_configs);
 
   // Builds `ObliviousHttpKeyConfigs` with given public_key and Single key
   // configuration specified in `ObliviousHttpHeaderKeyConfig` object. After
@@ -249,14 +248,14 @@
                       std::greater<uint8_t>>;
 
   ObliviousHttpKeyConfigs(ConfigMap cm, PublicKeyMap km,
-                          std::optional<uint8_t> first_key_id = std::nullopt)
+                          std::vector<uint8_t> key_ids)
       : configs_(std::move(cm)),
         public_keys_(std::move(km)),
-        first_key_id_(first_key_id) {}
+        key_ids_(key_ids) {}
 
   static absl::Status ReadSingleKeyConfig(QuicheDataReader& reader,
                                           ConfigMap& configs,
-                                          PublicKeyMap& keys,
+                                          PublicKeyMap& keys, uint8_t& key_id,
                                           bool skip_unknown_kems = false);
 
   // Reads key configs from a byte string formatted according to
@@ -264,7 +263,7 @@
   // 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,
-      std::optional<uint8_t>& first_key_id);
+      std::vector<uint8_t>& key_ids);
 
   // A mapping from key_id to ObliviousHttpHeaderKeyConfig objects for that key.
   const ConfigMap configs_;
@@ -272,9 +271,8 @@
   // 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_;
+  // The ordered list of key_ids in these configs.
+  const std::vector<uint8_t> key_ids_;
 };
 
 // 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 abe0edc..2aa0cf8 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
@@ -685,17 +685,17 @@
 }
 
 TEST(ObliviousHttpHeaderKeyConfigs, TestCreateWithWithMultipleKeys) {
-  std::string expected_preferred_public_key(32, 'b');
+  std::string expected_preferred_public_key(32, 'a');
   ObliviousHttpKeyConfigs::OhttpKeyConfig config1 = {
       100,
       EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
-      std::string(32, 'a'),
+      expected_preferred_public_key,
       {{EVP_HPKE_HKDF_SHA256, EVP_HPKE_AES_256_GCM}}};
   EXPECT_THAT(config1.DebugString(), HasSubstr("AES-256-GCM"));
   ObliviousHttpKeyConfigs::OhttpKeyConfig config2 = {
       200,
       EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
-      expected_preferred_public_key,
+      std::string(32, 'b'),
       {{EVP_HPKE_HKDF_SHA256, EVP_HPKE_CHACHA20_POLY1305}}};
   EXPECT_THAT(config2.DebugString(), HasSubstr("CHACHA20-POLY1305"));
   auto configs = ObliviousHttpKeyConfigs::Create({config1, config2});
@@ -703,17 +703,17 @@
   EXPECT_THAT(configs->DebugString(), HasSubstr("CHACHA20-POLY1305"));
   auto serialized_key = configs->GenerateConcatenatedKeys();
   ASSERT_THAT(serialized_key,
-              IsOkAndHolds(absl::StrCat(GetSerializedKeyConfig(config2),
-                                        GetSerializedKeyConfig(config1))));
+              IsOkAndHolds(absl::StrCat(GetSerializedKeyConfig(config1),
+                                        GetSerializedKeyConfig(config2))));
   auto ohttp_configs =
       ObliviousHttpKeyConfigs::ParseConcatenatedKeys(*serialized_key);
   QUICHE_ASSERT_OK(ohttp_configs);
   EXPECT_THAT(ohttp_configs->DebugString(), HasSubstr("CHACHA20-POLY1305"));
   ASSERT_EQ(ohttp_configs->NumKeys(), 2);
-  EXPECT_THAT(configs->PreferredConfig(),
-              AllOf(HasKeyId(200), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
-                    HasKdfId(EVP_HPKE_HKDF_SHA256),
-                    HasAeadId(EVP_HPKE_CHACHA20_POLY1305)));
+  EXPECT_THAT(
+      configs->PreferredConfig(),
+      AllOf(HasKeyId(100), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+            HasKdfId(EVP_HPKE_HKDF_SHA256), HasAeadId(EVP_HPKE_AES_256_GCM)));
   EXPECT_THAT(ohttp_configs->GetPublicKeyForId(
                   ohttp_configs->PreferredConfig().GetKeyId()),
               IsOkAndHolds(expected_preferred_public_key));
@@ -801,7 +801,7 @@
   QUICHE_ASSERT_OK(parsed_configs);
   EXPECT_EQ(parsed_configs->NumKeys(), 2);
   EXPECT_THAT(parsed_configs->PreferredConfig(),
-              AllOf(HasKeyId(200), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+              AllOf(HasKeyId(100), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
                     HasKdfId(EVP_HPKE_HKDF_SHA256)));
   EXPECT_THAT(parsed_configs->GetPublicKeyForId(100),
               IsOkAndHolds(std::string(32, 'a')));
@@ -835,7 +835,7 @@
   QUICHE_ASSERT_OK(parsed_configs);
   EXPECT_EQ(parsed_configs->NumKeys(), 2);
   EXPECT_THAT(parsed_configs->PreferredConfig(),
-              AllOf(HasKeyId(200), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
+              AllOf(HasKeyId(100), HasKemId(EVP_HPKE_DHKEM_X25519_HKDF_SHA256),
                     HasKdfId(EVP_HPKE_HKDF_SHA256)));
   EXPECT_THAT(parsed_configs->GetPublicKeyForId(100),
               IsOkAndHolds(std::string(32, 'a')));
diff --git a/quiche/quic/masque/masque_tcp_server_bin.cc b/quiche/quic/masque/masque_tcp_server_bin.cc
index f78054e..d28560a 100644
--- a/quiche/quic/masque/masque_tcp_server_bin.cc
+++ b/quiche/quic/masque/masque_tcp_server_bin.cc
@@ -22,7 +22,6 @@
 
 #include "absl/cleanup/cleanup.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -167,16 +166,16 @@
 
   static absl::StatusOr<std::unique_ptr<MasqueOhttpGateway>> Create() {
     auto ohttp_gateway = absl::WrapUnique(new MasqueOhttpGateway());
-    if (!quiche::GetQuicheCommandLineFlag(FLAGS_no_classic)) {
-      QUICHE_RETURN_IF_ERROR(ohttp_gateway->AddKeyConfig(
-          /*key_id=*/0x01, EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
-          quiche::GetQuicheCommandLineFlag(FLAGS_ohttp_key)));
-    }
     if (quiche::GetQuicheCommandLineFlag(FLAGS_pq)) {
       QUICHE_RETURN_IF_ERROR(ohttp_gateway->AddKeyConfig(
           /*key_id=*/0x02, EVP_HPKE_XWING,
           quiche::GetQuicheCommandLineFlag(FLAGS_pq_key)));
     }
+    if (!quiche::GetQuicheCommandLineFlag(FLAGS_no_classic)) {
+      QUICHE_RETURN_IF_ERROR(ohttp_gateway->AddKeyConfig(
+          /*key_id=*/0x01, EVP_HPKE_DHKEM_X25519_HKDF_SHA256,
+          quiche::GetQuicheCommandLineFlag(FLAGS_ohttp_key)));
+    }
     return ohttp_gateway;
   }
 
@@ -316,9 +315,9 @@
         kem_id,
         hpke_public_key,
         {{EVP_HPKE_HKDF_SHA256, EVP_HPKE_AES_128_GCM}}};
-    config_set_.insert(config);
+    configs_.push_back(config);
     QUICHE_ASSIGN_OR_RETURN(ObliviousHttpKeyConfigs ohttp_key_configs,
-                            ObliviousHttpKeyConfigs::Create(config_set_));
+                            ObliviousHttpKeyConfigs::Create(configs_));
     ohttp_key_configs_.emplace(std::move(ohttp_key_configs));
     QUICHE_LOG(INFO) << "Using OHTTP key configs: " << std::endl
                      << ohttp_key_configs_->DebugString();
@@ -328,7 +327,7 @@
   }
 
   Visitor* visitor_ = nullptr;
-  absl::flat_hash_set<ObliviousHttpKeyConfigs::OhttpKeyConfig> config_set_;
+  std::vector<ObliviousHttpKeyConfigs::OhttpKeyConfig> configs_;
   absl::flat_hash_map<uint8_t, std::string> hpke_private_keys_;
   std::optional<ObliviousHttpKeyConfigs> ohttp_key_configs_;
   std::string concatenated_keys_;