Update config cache to use per proxy token cache managers

Add proxy layer parameter to TryGetAuthTokens

Ip Protection will involve two proxies: ProxyA and ProxyB. The IpProtectionProxyLayer will indicate which proxy the tokens are for.

Change-Id: I3d5e8a0ad835875f0aabc981c392bfc650cd1a24
Bug: 1486528
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4936872
Reviewed-by: Dustin Mitchell <djmitche@chromium.org>
Reviewed-by: Kinuko Yasuda <kinuko@chromium.org>
Commit-Queue: Ciara McMullin <ciaramcmullin@google.com>
Reviewed-by: Ken Buchanan <kenrb@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1211589}

chrome/browser/ip_protection/ip_protection_config_provider.cc

@@ -73,13 +73,15 @@ IpProtectionConfigProvider::~IpProtectionConfigProvider() = default;
 
 void IpProtectionConfigProvider::TryGetAuthTokens(
     uint32_t batch_size,
+    network::mojom::IpProtectionProxyLayer proxy_layer,
     TryGetAuthTokensCallback callback) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   CHECK(!is_shutting_down_);
   SetUp();
+  CHECK(proxy_layer == network::mojom::IpProtectionProxyLayer::kProxyA);
 
-  // The `batch_size` is cast to an `int` for use by BlindSignAuth, so check for
-  // overflow here.
+  // The `batch_size` is cast to an `int` for use by BlindSignAuth, so check
+  // for overflow here.
   if (batch_size == 0 || batch_size > INT_MAX) {
     mojo::ReportBadMessage("Invalid batch_size");
     return;

chrome/browser/ip_protection/ip_protection_config_provider.h

@@ -83,6 +83,7 @@ class IpProtectionConfigProvider
   // It is forbidden for two calls to this method to be outstanding at the same
   // time.
   void TryGetAuthTokens(uint32_t batch_size,
+                        network::mojom::IpProtectionProxyLayer proxy_layer,
                         TryGetAuthTokensCallback callback) override;
 
   // Get the list of IP Protection proxies.

chrome/browser/ip_protection/ip_protection_config_provider_browsertest.cc

@@ -73,6 +73,7 @@ class IpProtectionConfigGetterInterceptor
   }
 
   void TryGetAuthTokens(uint32_t batch_size,
+                        network::mojom::IpProtectionProxyLayer proxy_layer,
                         TryGetAuthTokensCallback callback) override {
     if (should_intercept_) {
       // NOTE: We'll ignore batch size and just return one token.
@@ -83,7 +84,8 @@ class IpProtectionConfigGetterInterceptor
       std::move(callback).Run(std::move(tokens), base::Time());
       return;
     }
-    GetForwardingInterface()->TryGetAuthTokens(batch_size, std::move(callback));
+    GetForwardingInterface()->TryGetAuthTokens(batch_size, proxy_layer,
+                                               std::move(callback));
   }
 
   void EnableInterception() { should_intercept_ = true; }

chrome/browser/ip_protection/ip_protection_config_provider_unittest.cc

@@ -174,7 +174,9 @@ class IpProtectionConfigProviderTest : public testing::Test {
       }
     }
 
-    getter_->TryGetAuthTokens(num_tokens, tokens_future_.GetCallback());
+    getter_->TryGetAuthTokens(num_tokens,
+                              network::mojom::IpProtectionProxyLayer::kProxyA,
+                              tokens_future_.GetCallback());
 
     switch (primary_account_behavior_) {
       case PrimaryAccountBehavior::kNone:
@@ -485,7 +487,8 @@ TEST_F(IpProtectionConfigProviderTest, SessionRefreshTriggersBackoffReset) {
       absl::optional<std::vector<network::mojom::BlindSignedAuthTokenPtr>>,
       absl::optional<base::Time>>
       tokens_future;
-  getter_->TryGetAuthTokens(1, tokens_future.GetCallback());
+  getter_->TryGetAuthTokens(1, network::mojom::IpProtectionProxyLayer::kProxyA,
+                            tokens_future.GetCallback());
   const absl::optional<base::Time>& try_again_after =
       tokens_future.Get<absl::optional<base::Time>>();
   ASSERT_TRUE(try_again_after);
@@ -497,7 +500,8 @@ TEST_F(IpProtectionConfigProviderTest, SessionRefreshTriggersBackoffReset) {
 
   bsa_->tokens_ = {{"single-use-1", absl_expiration_time_}};
   tokens_future.Clear();
-  getter_->TryGetAuthTokens(1, tokens_future.GetCallback());
+  getter_->TryGetAuthTokens(1, network::mojom::IpProtectionProxyLayer::kProxyA,
+                            tokens_future.GetCallback());
   identity_test_env_.WaitForAccessTokenRequestIfNecessaryAndRespondWithToken(
       "access_token", base::Time::Now());
   const absl::optional<std::vector<network::mojom::BlindSignedAuthTokenPtr>>&

services/network/ip_protection_config_cache.h

@@ -27,31 +27,33 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) IpProtectionConfigCache {
   // Initializes the proxy list and token managers for the cache.
   virtual void SetUp() = 0;
 
-  // Check whether tokens are available.
+  // Check whether tokens are available in all token caches.
   //
   // This function is called on every URL load, so it should complete quickly.
-  virtual bool IsAuthTokenAvailable() = 0;
+  virtual bool AreAuthTokensAvailable() = 0;
 
   // Get a token, if one is available.
   //
   // Returns `nullopt` if no token is available, whether for a transient or
   // permanent reason. This method may return `nullopt` even if
   // `IsAuthTokenAvailable()` recently returned `true`.
-  virtual absl::optional<network::mojom::BlindSignedAuthTokenPtr>
-  GetAuthToken() = 0;
+  virtual absl::optional<network::mojom::BlindSignedAuthTokenPtr> GetAuthToken(
+      network::mojom::IpProtectionProxyLayer proxy_layer) = 0;
 
   // Invalidate any previous instruction that token requests should not be
   // made until after a specified time.
   virtual void InvalidateTryAgainAfterTime() = 0;
 
   // Set the token cache manager for the cache.
   virtual void SetIpProtectionTokenCacheManagerForTesting(
+      network::mojom::IpProtectionProxyLayer proxy_layer,
       std::unique_ptr<IpProtectionTokenCacheManager>
           ipp_token_cache_manager) = 0;
 
   // Fetch the token cache manager.
   virtual IpProtectionTokenCacheManager*
-  GetIpProtectionTokenCacheManagerForTesting() = 0;
+  GetIpProtectionTokenCacheManagerForTesting(
+      network::mojom::IpProtectionProxyLayer proxy_layer) = 0;
 
   // Set the proxy list manager for the cache.
   virtual void SetIpProtectionProxyListManagerForTesting(

services/network/ip_protection_config_cache_impl.cc

@@ -12,15 +12,15 @@
 #include "services/network/ip_protection_proxy_list_manager_impl.h"
 #include "services/network/ip_protection_token_cache_manager.h"
 #include "services/network/ip_protection_token_cache_manager_impl.h"
+#include "services/network/public/mojom/network_context.mojom.h"
 
 namespace network {
 
 IpProtectionConfigCacheImpl::IpProtectionConfigCacheImpl(
     mojo::PendingRemote<network::mojom::IpProtectionConfigGetter>
         config_getter) {
-  // Proxy list and token cache managers are null upon cache creation.
+  // Proxy list is null upon cache creation.
   ipp_proxy_list_manager_ = nullptr;
-  ipp_token_cache_manager_ = nullptr;
 
   if (config_getter.is_valid()) {
     config_getter_.Bind(std::move(config_getter));
@@ -32,39 +32,48 @@ IpProtectionConfigCacheImpl::~IpProtectionConfigCacheImpl() = default;
 void IpProtectionConfigCacheImpl::SetUp() {
   ipp_proxy_list_manager_ =
       std::make_unique<IpProtectionProxyListManagerImpl>(&config_getter_);
-  ipp_token_cache_manager_ =
-      std::make_unique<IpProtectionTokenCacheManagerImpl>(&config_getter_);
+  // TODO(@ciaramcmullin) Initialize kProxyB cache manager when fetching tokens
+  // for proxy B is supported.
+  ipp_token_cache_managers_[network::mojom::IpProtectionProxyLayer::kProxyA] =
+      std::make_unique<IpProtectionTokenCacheManagerImpl>(
+          &config_getter_, network::mojom::IpProtectionProxyLayer::kProxyA);
 }
 
-bool IpProtectionConfigCacheImpl::IsAuthTokenAvailable() {
-  return ipp_token_cache_manager_ != nullptr
-             ? ipp_token_cache_manager_->IsAuthTokenAvailable()
-             : false;
+bool IpProtectionConfigCacheImpl::AreAuthTokensAvailable() {
+  for (const auto& manager : ipp_token_cache_managers_) {
+    if (!manager.second->IsAuthTokenAvailable()) {
+      return false;
+    }
+  }
+  return !ipp_token_cache_managers_.empty();
 }
 
 absl::optional<network::mojom::BlindSignedAuthTokenPtr>
-IpProtectionConfigCacheImpl::GetAuthToken() {
+IpProtectionConfigCacheImpl::GetAuthToken(
+    network::mojom::IpProtectionProxyLayer proxy_layer) {
   absl::optional<network::mojom::BlindSignedAuthTokenPtr> result;
-  if (ipp_token_cache_manager_ != nullptr) {
-    result = ipp_token_cache_manager_->GetAuthToken();
+  if (ipp_token_cache_managers_.count(proxy_layer) > 0) {
+    result = ipp_token_cache_managers_[proxy_layer]->GetAuthToken();
   }
   return result;
 }
 
 void IpProtectionConfigCacheImpl::InvalidateTryAgainAfterTime() {
-  if (ipp_token_cache_manager_ != nullptr) {
-    ipp_token_cache_manager_->InvalidateTryAgainAfterTime();
+  for (const auto& manager : ipp_token_cache_managers_) {
+    manager.second->InvalidateTryAgainAfterTime();
   }
 }
 
 void IpProtectionConfigCacheImpl::SetIpProtectionTokenCacheManagerForTesting(
+    network::mojom::IpProtectionProxyLayer proxy_layer,
     std::unique_ptr<IpProtectionTokenCacheManager> ipp_token_cache_manager) {
-  ipp_token_cache_manager_ = std::move(ipp_token_cache_manager);
+  ipp_token_cache_managers_[proxy_layer] = std::move(ipp_token_cache_manager);
 }
 
 IpProtectionTokenCacheManager*
-IpProtectionConfigCacheImpl::GetIpProtectionTokenCacheManagerForTesting() {
-  return ipp_token_cache_manager_.get();
+IpProtectionConfigCacheImpl::GetIpProtectionTokenCacheManagerForTesting(
+    network::mojom::IpProtectionProxyLayer proxy_layer) {
+  return ipp_token_cache_managers_[proxy_layer].get();
 }
 
 void IpProtectionConfigCacheImpl::SetIpProtectionProxyListManagerForTesting(

services/network/ip_protection_config_cache_impl.h

@@ -6,6 +6,7 @@
 #define SERVICES_NETWORK_IP_PROTECTION_CONFIG_CACHE_IMPL_H_
 
 #include <deque>
+#include <map>
 
 #include "base/component_export.h"
 #include "base/functional/callback.h"
@@ -33,15 +34,16 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) IpProtectionConfigCacheImpl
 
   // IpProtectionConfigCache implementation.
   void SetUp() override;
-  bool IsAuthTokenAvailable() override;
-  absl::optional<network::mojom::BlindSignedAuthTokenPtr> GetAuthToken()
-      override;
+  bool AreAuthTokensAvailable() override;
+  absl::optional<network::mojom::BlindSignedAuthTokenPtr> GetAuthToken(
+      network::mojom::IpProtectionProxyLayer proxy_layer) override;
   void InvalidateTryAgainAfterTime() override;
   void SetIpProtectionTokenCacheManagerForTesting(
+      network::mojom::IpProtectionProxyLayer proxy_layer,
       std::unique_ptr<IpProtectionTokenCacheManager> ipp_token_cache_manager)
       override;
-  IpProtectionTokenCacheManager* GetIpProtectionTokenCacheManagerForTesting()
-      override;
+  IpProtectionTokenCacheManager* GetIpProtectionTokenCacheManagerForTesting(
+      network::mojom::IpProtectionProxyLayer proxy_layer) override;
   void SetIpProtectionProxyListManagerForTesting(
       std::unique_ptr<IpProtectionProxyListManager> ipp_proxy_list_manager)
       override;
@@ -56,8 +58,10 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) IpProtectionConfigCacheImpl
   // A manager for the list of currently cached proxy hostnames.
   std::unique_ptr<IpProtectionProxyListManager> ipp_proxy_list_manager_;
 
-  // A manager for cache of blind-signed auth tokens.
-  std::unique_ptr<IpProtectionTokenCacheManager> ipp_token_cache_manager_;
+  // Proxy layer managers for cache of blind-signed auth tokens.
+  std::map<network::mojom::IpProtectionProxyLayer,
+           std::unique_ptr<IpProtectionTokenCacheManager>>
+      ipp_token_cache_managers_;
 
   base::WeakPtrFactory<IpProtectionConfigCacheImpl> weak_ptr_factory_{this};
 };

services/network/ip_protection_config_cache_impl_unittest.cc

@@ -94,10 +94,12 @@ TEST_F(IpProtectionConfigCacheImplTest, GetAuthTokenFromManager) {
       std::make_unique<MockIpProtectionTokenCacheManager>();
   ipp_token_cache_manager_->SetAuthToken(std::move(exp_token));
   ipp_config_cache_->SetIpProtectionTokenCacheManagerForTesting(
+      network::mojom::IpProtectionProxyLayer::kProxyA,
       std::move(ipp_token_cache_manager_));
 
-  ASSERT_TRUE(ipp_config_cache_->IsAuthTokenAvailable());
-  ASSERT_TRUE(ipp_config_cache_->GetAuthToken());
+  ASSERT_TRUE(ipp_config_cache_->AreAuthTokensAvailable());
+  ASSERT_TRUE(ipp_config_cache_->GetAuthToken(
+      network::mojom::IpProtectionProxyLayer::kProxyA));
 }
 
 // Proxy list manager returns currently cached proxy hostnames.

services/network/ip_protection_proxy_list_manager_impl_unittest.cc

@@ -11,6 +11,7 @@
 #include "mojo/public/cpp/bindings/receiver.h"
 #include "services/network/ip_protection_proxy_list_manager.h"
 #include "services/network/ip_protection_proxy_list_manager_impl.h"
+#include "services/network/public/mojom/network_context.mojom.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
 
@@ -56,6 +57,7 @@ class MockIpProtectionConfigGetter
   void Reset() { expected_get_proxy_list_calls_.clear(); }
 
   void TryGetAuthTokens(uint32_t batch_size,
+                        network::mojom::IpProtectionProxyLayer proxy_layer,
                         TryGetAuthTokensCallback callback) override {
     NOTREACHED_NORETURN();
   }

services/network/ip_protection_token_cache_manager_impl.cc

@@ -8,6 +8,7 @@
 #include "base/task/task_traits.h"
 #include "base/task/thread_pool.h"
 #include "base/time/time.h"
+#include "services/network/public/mojom/network_context.mojom.h"
 
 namespace network {
 
@@ -24,11 +25,13 @@ const base::TimeDelta kTokenRateMeasurementInterval = base::Minutes(5);
 
 IpProtectionTokenCacheManagerImpl::IpProtectionTokenCacheManagerImpl(
     mojo::Remote<network::mojom::IpProtectionConfigGetter>* config_getter,
+    network::mojom::IpProtectionProxyLayer proxy_layer,
     bool disable_cache_management_for_testing)
     : batch_size_(net::features::kIpPrivacyAuthTokenCacheBatchSize.Get()),
       cache_low_water_mark_(
           net::features::kIpPrivacyAuthTokenCacheLowWaterMark.Get()),
       config_getter_(config_getter),
+      proxy_layer_(proxy_layer),
       disable_cache_management_for_testing_(
           disable_cache_management_for_testing) {
   last_token_rate_measurement_ = base::TimeTicks::Now();
@@ -77,7 +80,7 @@ void IpProtectionTokenCacheManagerImpl::MaybeRefillCache() {
     fetching_auth_tokens_ = true;
     VLOG(2) << "IPPATC::MaybeRefillCache calling TryGetAuthTokens";
     config_getter_->get()->TryGetAuthTokens(
-        batch_size_,
+        batch_size_, proxy_layer_,
         base::BindOnce(&IpProtectionTokenCacheManagerImpl::OnGotAuthTokens,
                        weak_ptr_factory_.GetWeakPtr()));
   }
@@ -234,7 +237,7 @@ void IpProtectionTokenCacheManagerImpl::CallTryGetAuthTokensForTesting() {
   CHECK(config_getter_);
   CHECK(on_try_get_auth_tokens_completed_for_testing_);
   config_getter_->get()->TryGetAuthTokens(
-      batch_size_,
+      batch_size_, proxy_layer_,
       base::BindOnce(&IpProtectionTokenCacheManagerImpl::OnGotAuthTokens,
                      weak_ptr_factory_.GetWeakPtr()));
 }

services/network/ip_protection_token_cache_manager_impl.h

@@ -26,6 +26,7 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) IpProtectionTokenCacheManagerImpl
  public:
   explicit IpProtectionTokenCacheManagerImpl(
       mojo::Remote<network::mojom::IpProtectionConfigGetter>* config_getter,
+      network::mojom::IpProtectionProxyLayer proxy_layer,
       bool disable_cache_management_for_testing = false);
   ~IpProtectionTokenCacheManagerImpl() override;
 
@@ -96,6 +97,9 @@ class COMPONENT_EXPORT(NETWORK_SERVICE) IpProtectionTokenCacheManagerImpl
   raw_ptr<mojo::Remote<network::mojom::IpProtectionConfigGetter>>
       config_getter_;
 
+  // The proxy layer which the cache of tokens will be used for.
+  network::mojom::IpProtectionProxyLayer proxy_layer_;
+
   // True if an invocation of `config_getter_.TryGetAuthTokens()` is
   // outstanding.
   bool fetching_auth_tokens_ = false;

services/network/ip_protection_token_cache_manager_impl_unittest.cc

@@ -13,6 +13,7 @@
 #include "services/network/ip_protection_config_cache_impl.h"
 #include "services/network/ip_protection_token_cache_manager.h"
 #include "services/network/ip_protection_token_cache_manager_impl.h"
+#include "services/network/public/mojom/network_context.mojom.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
 
@@ -76,6 +77,7 @@ class MockIpProtectionConfigGetter
   void Reset() { expected_try_get_auth_token_calls_.clear(); }
 
   void TryGetAuthTokens(uint32_t batch_size,
+                        network::mojom::IpProtectionProxyLayer proxy_layer,
                         TryGetAuthTokensCallback callback) override {
     ASSERT_FALSE(expected_try_get_auth_token_calls_.empty())
         << "Unexpected call to TryGetAuthTokens";
@@ -112,7 +114,7 @@ class IpProtectionTokenCacheManagerImplTest : public testing::Test {
     remote_.Bind(receiver_.BindNewPipeAndPassRemote());
     ipp_token_cache_manager_ =
         std::make_unique<IpProtectionTokenCacheManagerImpl>(
-            &remote_,
+            &remote_, network::mojom::IpProtectionProxyLayer::kProxyA,
             /* disable_cache_management_for_testing=*/true);
   }
 
@@ -273,7 +275,7 @@ TEST_F(IpProtectionTokenCacheManagerImplTest, SkipExpiredTokens) {
 // but things don't crash.
 TEST_F(IpProtectionTokenCacheManagerImplTest, NullGetter) {
   auto ipp_token_cache_manager = IpProtectionTokenCacheManagerImpl(
-      nullptr,
+      nullptr, network::mojom::IpProtectionProxyLayer::kProxyA,
       /* disable_cache_management_for_testing=*/true);
   EXPECT_FALSE(ipp_token_cache_manager_->IsAuthTokenAvailable());
   auto token = ipp_token_cache_manager.GetAuthToken();

services/network/network_context.cc

@@ -2040,7 +2040,8 @@ void NetworkContext::VerifyIpProtectionConfigGetterForTesting(
   auto* ipp_token_cache_manager_impl =
       static_cast<IpProtectionTokenCacheManagerImpl*>(
           ipp_config_cache
-              ->GetIpProtectionTokenCacheManagerForTesting());  // IN-TEST
+              ->GetIpProtectionTokenCacheManagerForTesting(  // IN-TEST
+                  network::mojom::IpProtectionProxyLayer::kProxyA));
   CHECK(ipp_token_cache_manager_impl);
 
   // If active cache management is enabled (the default), disable it and do a
@@ -2061,8 +2062,9 @@ void NetworkContext::VerifyIpProtectionConfigGetterForTesting(
               auto* ipp_config_cache =
                   weak_ptr->proxy_delegate_->GetIpProtectionConfigCache();
               ipp_config_cache->InvalidateTryAgainAfterTime();
-              while (ipp_config_cache->IsAuthTokenAvailable()) {
-                ipp_config_cache->GetAuthToken();
+              while (ipp_config_cache->AreAuthTokensAvailable()) {
+                ipp_config_cache->GetAuthToken(
+                    network::mojom::IpProtectionProxyLayer::kProxyA);
               }
               // Call `PostTask()` instead of invoking the Verify method again
               // directly so that if `DisableCacheManagementForTesting()` needed
@@ -2103,10 +2105,12 @@ void NetworkContext::OnIpProtectionConfigAvailableForTesting(
   auto* ipp_token_cache_manager_impl =
       static_cast<IpProtectionTokenCacheManagerImpl*>(
           ipp_config_cache
-              ->GetIpProtectionTokenCacheManagerForTesting());  // IN-TEST
+              ->GetIpProtectionTokenCacheManagerForTesting(  // IN-TEST
+                  network::mojom::IpProtectionProxyLayer::kProxyA));
 
   absl::optional<network::mojom::BlindSignedAuthTokenPtr> result =
-      ipp_config_cache->GetAuthToken();
+      ipp_config_cache->GetAuthToken(
+          network::mojom::IpProtectionProxyLayer::kProxyA);
   if (result.has_value()) {
     std::move(callback).Run(std::move(result).value(), absl::nullopt);
     return;