Update get retry config API.

abakiaydin · facebook-github-bot · commit c7532b85d91d · 2025-05-13T22:52:17.000-07:00
Summary:
Adds `sni` parameter to `getRetryConfigs` api.
Returns configs matching the sni in the front of retry config list, keeps the relative order of the configs.

Reviewed By: mingtaoy

Differential Revision: D73597518

fbshipit-source-id: 29042607e20064b99d740575ad910c6609097c08
diff --git a/third-party/fizz/src/fizz/protocol/ech/Decrypter.cpp b/third-party/fizz/src/fizz/protocol/ech/Decrypter.cpp
@@ -8,6 +8,8 @@
 
 #include <fizz/protocol/ech/Decrypter.h>
 
+#include <iterator>
+
 namespace fizz {
 namespace ech {
 
@@ -180,11 +182,30 @@ ClientHello ECHConfigManager::decryptClientHelloHRR(
       *factory_, chlo, encapsulatedKey, dummy, configs_);
 }
 
-std::vector<ech::ECHConfig> ECHConfigManager::getRetryConfigs() const {
+std::vector<ech::ECHConfig> ECHConfigManager::getRetryConfigs(
+    const folly::Optional<std::string>& maybeSni) const {
   std::vector<ech::ECHConfig> retryConfigs;
-  for (const auto& cfg : configs_) {
-    retryConfigs.push_back(cfg.echConfig);
+  std::vector<ech::ECHConfig> nonMatchingConfigs;
+  for (const auto& config : configs_) {
+    switch (config.echConfig.version) {
+      case ECHVersion::Draft15: {
+        auto cursor =
+            folly::io::Cursor(config.echConfig.ech_config_content.get());
+        auto currentConfigDraft = decode<ECHConfigContentDraft>(cursor);
+        if (maybeSni.hasValue() &&
+            maybeSni.value() ==
+                currentConfigDraft.public_name->to<std::string>()) {
+          retryConfigs.push_back(config.echConfig);
+        } else {
+          nonMatchingConfigs.push_back(config.echConfig);
+        }
+      } break;
+    }
   }
+  retryConfigs.insert(
+      retryConfigs.end(),
+      std::make_move_iterator(nonMatchingConfigs.begin()),
+      std::make_move_iterator(nonMatchingConfigs.end()));
   return retryConfigs;
 }
 
diff --git a/third-party/fizz/src/fizz/protocol/ech/Decrypter.h b/third-party/fizz/src/fizz/protocol/ech/Decrypter.h
@@ -42,7 +42,8 @@ class Decrypter {
   virtual ClientHello decryptClientHelloHRR(
       const ClientHello& chlo,
       const std::unique_ptr<folly::IOBuf>& encapsulatedKey) = 0;
-  virtual std::vector<ech::ECHConfig> getRetryConfigs() const = 0;
+  virtual std::vector<ech::ECHConfig> getRetryConfigs(
+      const folly::Optional<std::string>& maybeSni) const = 0;
 };
 
 class ECHConfigManager : public Decrypter {
@@ -58,9 +59,10 @@ class ECHConfigManager : public Decrypter {
   ClientHello decryptClientHelloHRR(
       const ClientHello& chlo,
       const std::unique_ptr<folly::IOBuf>& encapsulatedKey) override;
-  std::vector<ech::ECHConfig> getRetryConfigs() const override;
+  std::vector<ech::ECHConfig> getRetryConfigs(
+      const folly::Optional<std::string>& maybeSni) const override;
 
- private:
+ protected:
   std::shared_ptr<Factory> factory_;
   std::vector<DecrypterParams> configs_;
 };
diff --git a/third-party/fizz/src/fizz/protocol/ech/test/DecrypterTest.cpp b/third-party/fizz/src/fizz/protocol/ech/test/DecrypterTest.cpp
@@ -16,6 +16,7 @@
 #include <fizz/protocol/ech/test/TestUtil.h>
 #include <fizz/protocol/test/TestUtil.h>
 
+using namespace std::string_literals;
 using namespace fizz::test;
 
 namespace fizz {
@@ -102,6 +103,36 @@ ClientHello getChloOuterHRRWithExt(
   return chloOuter;
 }
 
+ECHConfig makeDummyConfig(uint8_t configId, const std::string& publicName) {
+  auto content = ECHConfigContentDraft{};
+  content.key_config.config_id = configId;
+  content.public_name = folly::IOBuf::copyBuffer(publicName);
+  content.key_config.public_key = folly::IOBuf::copyBuffer("public");
+  auto config = ECHConfig{};
+  config.version = ECHVersion::Draft15;
+  config.ech_config_content = fizz::encode(std::move(content));
+  return config;
+}
+
+struct RetryConfigExpectation {
+  uint8_t id;
+  std::string name;
+};
+
+void checkRetryConfigExpectation(
+    const std::vector<RetryConfigExpectation>& expectations,
+    const std::vector<ECHConfig>& configs) {
+  ASSERT_EQ(expectations.size(), configs.size());
+  for (size_t i = 0; i < expectations.size(); ++i) {
+    const auto& config = configs[i];
+    auto cursor = folly::io::Cursor(config.ech_config_content.get());
+    auto configContent = decode<ECHConfigContentDraft>(cursor);
+    EXPECT_EQ(expectations[i].id, configContent.key_config.config_id);
+    EXPECT_EQ(
+        expectations[i].name, configContent.public_name->to<std::string>());
+  }
+}
+
 TEST(DecrypterTest, TestDecodeSuccess) {
   auto kex = openssl::makeOpenSSLECKeyExchange<fizz::P256>();
   kex->setPrivateKey(getPrivateKey(kP256Key));
@@ -254,6 +285,57 @@ TEST(DecrypterTest, TestDecodeMultipleDecrypterParam) {
   EXPECT_TRUE(result->configId == 2);
 }
 
+TEST(DecrypterTest, TestGetRetryConfigs) {
+  auto decrypter = ECHConfigManager(std::make_shared<DefaultFactory>());
+  auto publicNames = {"a.com", "b.com", "a.com", "c.com", "b.com", "a.com"};
+  uint8_t configId = 0;
+  for (const auto& name : publicNames) {
+    decrypter.addDecryptionConfig(
+        {.echConfig = makeDummyConfig(configId++, name)});
+  }
+
+  checkRetryConfigExpectation(
+      {{0, "a.com"},
+       {2, "a.com"},
+       {5, "a.com"},
+       {1, "b.com"},
+       {3, "c.com"},
+       {4, "b.com"}},
+      decrypter.getRetryConfigs("a.com"s));
+  checkRetryConfigExpectation(
+      {{1, "b.com"},
+       {4, "b.com"},
+       {0, "a.com"},
+       {2, "a.com"},
+       {3, "c.com"},
+       {5, "a.com"}},
+      decrypter.getRetryConfigs("b.com"s));
+  checkRetryConfigExpectation(
+      {{3, "c.com"},
+       {0, "a.com"},
+       {1, "b.com"},
+       {2, "a.com"},
+       {4, "b.com"},
+       {5, "a.com"}},
+      decrypter.getRetryConfigs("c.com"s));
+  checkRetryConfigExpectation(
+      {{0, "a.com"},
+       {1, "b.com"},
+       {2, "a.com"},
+       {3, "c.com"},
+       {4, "b.com"},
+       {5, "a.com"}},
+      decrypter.getRetryConfigs("d.com"s));
+  checkRetryConfigExpectation(
+      {{0, "a.com"},
+       {1, "b.com"},
+       {2, "a.com"},
+       {3, "c.com"},
+       {4, "b.com"},
+       {5, "a.com"}},
+      decrypter.getRetryConfigs(folly::none));
+}
+
 } // namespace test
 } // namespace ech
 } // namespace fizz
diff --git a/third-party/fizz/src/fizz/protocol/test/Mocks.h b/third-party/fizz/src/fizz/protocol/test/Mocks.h
@@ -483,7 +483,7 @@ class MockECHDecrypter : public ech::Decrypter {
   MOCK_METHOD(
       std::vector<ech::ECHConfig>,
       getRetryConfigs,
-      (),
+      (const folly::Optional<std::string>& maybeSni),
       (const, override));
 };
 } // namespace test
diff --git a/third-party/fizz/src/fizz/server/ServerProtocol.cpp b/third-party/fizz/src/fizz/server/ServerProtocol.cpp
@@ -1073,9 +1073,9 @@ static std::pair<std::vector<ExtensionType>, Buf> getCertificateRequest(
       std::move(certReqExtensions), std::move(encodedCertificateRequest));
 }
 
-static std::string getSNI(const ClientHello& chlo) {
+static folly::Optional<std::string> getSNI(const ClientHello& chlo) {
+  folly::Optional<std::string> sni;
   auto serverNameList = getExtension<ServerNameList>(chlo.extensions);
-  std::string sni;
   if (serverNameList && !serverNameList->server_name_list.empty()) {
     sni = serverNameList->server_name_list.front().hostname->to<std::string>();
   }
@@ -1185,6 +1185,7 @@ static std::pair<ECHStatus, folly::Optional<ECHState>> processECH(
     if (state.echState().has_value()) {
       echState->hpkeContext = std::move(state.echState()->hpkeContext);
       echState->cipherSuite = state.echState()->cipherSuite;
+      echState->outerSni = state.echState()->outerSni;
     }
   } else {
     bool requestedECH =
@@ -1196,14 +1197,12 @@ static std::pair<ECHStatus, folly::Optional<ECHState>> processECH(
     if (requestedECH) {
       auto echExt = getExtension<ech::OuterECHClientHello>(chlo.extensions);
       echState = ECHState{
-          echExt->cipher_suite, echExt->config_id, nullptr, folly::none};
+          echExt->cipher_suite, echExt->config_id, nullptr, getSNI(chlo)};
       if (decrypter) {
         auto gotChlo = decrypter->decryptClientHello(chlo);
         if (gotChlo.has_value()) {
-          auto outerSni = getSNI(chlo);
           echStatus = ECHStatus::Accepted;
           echState->hpkeContext = std::move(gotChlo->context);
-          echState->outerSni = outerSni;
           chlo = std::move(gotChlo->chlo);
         } else {
           echStatus = ECHStatus::Rejected;
@@ -1636,7 +1635,10 @@ EventHandler<ServerTypes, StateEnum::ExpectingClientHello, Event::ClientHello>::
                     std::move(echScheduler));
               } else if (echStatus == ECHStatus::Rejected) {
                 auto decrypter = state.context()->getECHDecrypter();
-                echRetryConfigs = decrypter->getRetryConfigs();
+                DCHECK(decrypter);
+                DCHECK(echState);
+                echRetryConfigs =
+                    decrypter->getRetryConfigs(echState->outerSni);
               }
 
               auto encodedServerHello = encodeHandshake(std::move(serverHello));
diff --git a/third-party/fizz/src/fizz/server/test/ServerProtocolTest.cpp b/third-party/fizz/src/fizz/server/test/ServerProtocolTest.cpp
@@ -1488,7 +1488,7 @@ TEST_F(ServerProtocolTest, TestECHDecryptionFailure) {
   auto decrypter = std::make_shared<MockECHDecrypter>();
   EXPECT_CALL(*decrypter, decryptClientHello(_))
       .WillOnce(InvokeWithoutArgs([=]() { return folly::none; }));
-  EXPECT_CALL(*decrypter, getRetryConfigs())
+  EXPECT_CALL(*decrypter, getRetryConfigs(_))
       .WillOnce(InvokeWithoutArgs([]() -> std::vector<ech::ECHConfig> {
         ech::ECHConfig cfg;
         cfg.version = ech::ECHVersion::Draft15;
@@ -3343,8 +3343,15 @@ TEST_F(ServerProtocolTest, TestRetryECHMissingInnerExtension) {
 TEST_F(ServerProtocolTest, TestRetryClientHelloECHRejectedFlow) {
   setUpExpectingClientHelloRetry();
   state_.echStatus() = ECHStatus::Rejected;
+  fizz::server::ECHState echState;
+  echState.cipherSuite.kdf_id = fizz::hpke::KDFId::Sha256;
+  echState.cipherSuite.aead_id = fizz::hpke::AeadId::TLS_AES_128_GCM_SHA256;
+  echState.configId = 0xFB;
+  echState.hpkeContext = nullptr;
+  echState.outerSni = "test.net";
+  state_.echState() = std::move(echState);
   auto decrypter = std::make_shared<MockECHDecrypter>();
-  EXPECT_CALL(*decrypter, getRetryConfigs())
+  EXPECT_CALL(*decrypter, getRetryConfigs(_))
       .WillOnce(InvokeWithoutArgs([]() -> std::vector<ech::ECHConfig> {
         ech::ECHConfig cfg;
         cfg.version = ech::ECHVersion::Draft15;
