Split token cache from IdentityAPI and add tests

Implementing a subset matching cache will require a more complex cache
for the IdentityAPI class. At the moment, the existing cache is
contained inside the IdentityAPI class.

This CL separates the two, so all of the caching logic for IdentityAPI
is in its own class. Additionally, unit tests have been added to test
the basic functionality of this cache.

Bug: 1100535
Change-Id: I39f5357d3ce87ae4d958aa2c8585cc2deec624be
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2343318Reviewed-by: David Roger <droger@chromium.org>
Reviewed-by: Alex Ilin <alexilin@chromium.org>
Commit-Queue: William Lin <williamlin@google.com>
Cr-Commit-Position: refs/heads/master@{#797527}

chrome/browser/extensions/BUILD.gn

@@ -197,6 +197,8 @@ static_library("extensions") {
     "api/identity/identity_private_api.h",
     "api/identity/identity_remove_cached_auth_token_function.cc",
     "api/identity/identity_remove_cached_auth_token_function.h",
+    "api/identity/identity_token_cache.cc",
+    "api/identity/identity_token_cache.h",
     "api/identity/web_auth_flow.cc",
     "api/identity/web_auth_flow.h",
     "api/idltest/idltest_api.cc",

chrome/browser/extensions/api/identity/identity_api.cc

@@ -6,9 +6,7 @@
 
 #include <stddef.h>
 
-#include <algorithm>
 #include <memory>
-#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -23,7 +21,6 @@
 #include "chrome/browser/app_mode/app_mode_utils.h"
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/chrome_notification_types.h"
-#include "chrome/browser/extensions/api/identity/identity_constants.h"
 #include "chrome/browser/extensions/extension_service.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/signin/account_consistency_mode_manager.h"
@@ -47,99 +44,6 @@ namespace {
 const char kIdentityGaiaIdPref[] = "identity_gaia_id";
 }
 
-IdentityTokenCacheValue::IdentityTokenCacheValue() = default;
-IdentityTokenCacheValue::IdentityTokenCacheValue(
-    const IdentityTokenCacheValue& other) = default;
-IdentityTokenCacheValue::~IdentityTokenCacheValue() = default;
-
-// static
-IdentityTokenCacheValue IdentityTokenCacheValue::CreateIssueAdvice(
-    const IssueAdviceInfo& issue_advice) {
-  IdentityTokenCacheValue cache_value;
-  cache_value.status_ = CACHE_STATUS_ADVICE;
-  cache_value.issue_advice_ = issue_advice;
-  cache_value.expiration_time_ =
-      base::Time::Now() + base::TimeDelta::FromSeconds(
-                              identity_constants::kCachedIssueAdviceTTLSeconds);
-  return cache_value;
-}
-
-// static
-IdentityTokenCacheValue IdentityTokenCacheValue::CreateRemoteConsent(
-    const RemoteConsentResolutionData& resolution_data) {
-  IdentityTokenCacheValue cache_value;
-  cache_value.status_ = CACHE_STATUS_REMOTE_CONSENT;
-  cache_value.resolution_data_ = resolution_data;
-  cache_value.expiration_time_ =
-      base::Time::Now() + base::TimeDelta::FromSeconds(
-                              identity_constants::kCachedIssueAdviceTTLSeconds);
-  return cache_value;
-}
-
-// static
-IdentityTokenCacheValue IdentityTokenCacheValue::CreateRemoteConsentApproved(
-    const std::string& consent_result) {
-  IdentityTokenCacheValue cache_value;
-  cache_value.status_ = CACHE_STATUS_REMOTE_CONSENT_APPROVED;
-  cache_value.consent_result_ = consent_result;
-  cache_value.expiration_time_ =
-      base::Time::Now() + base::TimeDelta::FromSeconds(
-                              identity_constants::kCachedIssueAdviceTTLSeconds);
-  return cache_value;
-}
-
-// static
-IdentityTokenCacheValue IdentityTokenCacheValue::CreateToken(
-    const std::string& token,
-    base::TimeDelta time_to_live) {
-  IdentityTokenCacheValue cache_value;
-  cache_value.status_ = CACHE_STATUS_TOKEN;
-  cache_value.token_ = token;
-
-  // Remove 20 minutes from the ttl so cached tokens will have some time
-  // to live any time they are returned.
-  time_to_live -= base::TimeDelta::FromMinutes(20);
-
-  base::TimeDelta zero_delta;
-  if (time_to_live < zero_delta)
-    time_to_live = zero_delta;
-
-  cache_value.expiration_time_ = base::Time::Now() + time_to_live;
-  return cache_value;
-}
-
-IdentityTokenCacheValue::CacheValueStatus IdentityTokenCacheValue::status()
-    const {
-  if (is_expired())
-    return IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND;
-  else
-    return status_;
-}
-
-const IssueAdviceInfo& IdentityTokenCacheValue::issue_advice() const {
-  return issue_advice_;
-}
-
-const RemoteConsentResolutionData& IdentityTokenCacheValue::resolution_data()
-    const {
-  return resolution_data_;
-}
-
-const std::string& IdentityTokenCacheValue::consent_result() const {
-  return consent_result_;
-}
-
-const std::string& IdentityTokenCacheValue::token() const { return token_; }
-
-bool IdentityTokenCacheValue::is_expired() const {
-  return status_ == CACHE_STATUS_NOTFOUND ||
-         expiration_time_ < base::Time::Now();
-}
-
-const base::Time& IdentityTokenCacheValue::expiration_time() const {
-  return expiration_time_;
-}
-
 IdentityAPI::IdentityAPI(content::BrowserContext* context)
     : IdentityAPI(Profile::FromBrowserContext(context),
                   IdentityManagerFactory::GetForProfile(
@@ -151,37 +55,8 @@ IdentityAPI::~IdentityAPI() {}
 
 IdentityMintRequestQueue* IdentityAPI::mint_queue() { return &mint_queue_; }
 
-void IdentityAPI::SetCachedToken(const ExtensionTokenKey& key,
-                                 const IdentityTokenCacheValue& token_data) {
-  auto it = token_cache_.find(key);
-  if (it != token_cache_.end() && it->second.status() <= token_data.status())
-    token_cache_.erase(it);
-
-  token_cache_.insert(std::make_pair(key, token_data));
-}
-
-void IdentityAPI::EraseCachedToken(const std::string& extension_id,
-                                   const std::string& token) {
-  CachedTokens::iterator it;
-  for (it = token_cache_.begin(); it != token_cache_.end(); ++it) {
-    if (it->first.extension_id == extension_id &&
-        it->second.status() == IdentityTokenCacheValue::CACHE_STATUS_TOKEN &&
-        it->second.token() == token) {
-      token_cache_.erase(it);
-      break;
-    }
-  }
-}
-
-void IdentityAPI::EraseAllCachedTokens() { token_cache_.clear(); }
-
-const IdentityTokenCacheValue& IdentityAPI::GetCachedToken(
-    const ExtensionTokenKey& key) {
-  return token_cache_[key];
-}
-
-const IdentityAPI::CachedTokens& IdentityAPI::GetAllCachedTokens() {
-  return token_cache_;
+IdentityTokenCache* IdentityAPI::token_cache() {
+  return &token_cache_;
 }
 
 void IdentityAPI::SetGaiaIdForExtension(const std::string& extension_id,

chrome/browser/extensions/api/identity/identity_api.h

@@ -5,11 +5,8 @@
 #ifndef CHROME_BROWSER_EXTENSIONS_API_IDENTITY_IDENTITY_API_H_
 #define CHROME_BROWSER_EXTENSIONS_API_IDENTITY_IDENTITY_API_H_
 
-#include <map>
-#include <set>
 #include <string>
 #include <utility>
-#include <vector>
 
 #include "base/callback_list.h"
 #include "base/feature_list.h"
@@ -20,7 +17,6 @@
 #include "base/optional.h"
 #include "build/build_config.h"
 #include "build/buildflag.h"
-#include "chrome/browser/extensions/api/identity/extension_token_key.h"
 #include "chrome/browser/extensions/api/identity/gaia_web_auth_flow.h"
 #include "chrome/browser/extensions/api/identity/identity_get_accounts_function.h"
 #include "chrome/browser/extensions/api/identity/identity_get_auth_token_function.h"
@@ -28,6 +24,7 @@
 #include "chrome/browser/extensions/api/identity/identity_launch_web_auth_flow_function.h"
 #include "chrome/browser/extensions/api/identity/identity_mint_queue.h"
 #include "chrome/browser/extensions/api/identity/identity_remove_cached_auth_token_function.h"
+#include "chrome/browser/extensions/api/identity/identity_token_cache.h"
 #include "chrome/browser/extensions/api/identity/web_auth_flow.h"
 #include "components/signin/public/base/signin_buildflags.h"
 #include "components/signin/public/identity_manager/identity_manager.h"
@@ -42,53 +39,9 @@ class Profile;
 
 namespace extensions {
 
-class IdentityTokenCacheValue {
- public:
-  IdentityTokenCacheValue();
-  IdentityTokenCacheValue(const IdentityTokenCacheValue& other);
-  ~IdentityTokenCacheValue();
-
-  static IdentityTokenCacheValue CreateIssueAdvice(
-      const IssueAdviceInfo& issue_advice);
-  static IdentityTokenCacheValue CreateRemoteConsent(
-      const RemoteConsentResolutionData& resolution_data);
-  static IdentityTokenCacheValue CreateRemoteConsentApproved(
-      const std::string& consent_result);
-  static IdentityTokenCacheValue CreateToken(const std::string& token,
-                                             base::TimeDelta time_to_live);
-
-  // Order of these entries is used to determine whether or not new
-  // entries supersede older ones in SetCachedToken.
-  enum CacheValueStatus {
-    CACHE_STATUS_NOTFOUND,
-    CACHE_STATUS_ADVICE,
-    CACHE_STATUS_REMOTE_CONSENT,
-    CACHE_STATUS_REMOTE_CONSENT_APPROVED,
-    CACHE_STATUS_TOKEN
-  };
-
-  CacheValueStatus status() const;
-  const IssueAdviceInfo& issue_advice() const;
-  const RemoteConsentResolutionData& resolution_data() const;
-  const std::string& consent_result() const;
-  const std::string& token() const;
-  const base::Time& expiration_time() const;
-
- private:
-  bool is_expired() const;
-
-  CacheValueStatus status_ = CACHE_STATUS_NOTFOUND;
-  IssueAdviceInfo issue_advice_;
-  RemoteConsentResolutionData resolution_data_;
-  std::string consent_result_;
-  std::string token_;
-  base::Time expiration_time_;
-};
-
 class IdentityAPI : public BrowserContextKeyedAPI,
                     public signin::IdentityManager::Observer {
  public:
-  using CachedTokens = std::map<ExtensionTokenKey, IdentityTokenCacheValue>;
   using OnSetConsentResultSignature = void(const std::string&,
                                            const std::string&);
 
@@ -98,14 +51,7 @@ class IdentityAPI : public BrowserContextKeyedAPI,
   // Request serialization queue for getAuthToken.
   IdentityMintRequestQueue* mint_queue();
 
-  // Token cache.
-  void SetCachedToken(const ExtensionTokenKey& key,
-                      const IdentityTokenCacheValue& token_data);
-  void EraseCachedToken(const std::string& extension_id,
-                        const std::string& token);
-  void EraseAllCachedTokens();
-  const IdentityTokenCacheValue& GetCachedToken(const ExtensionTokenKey& key);
-  const CachedTokens& GetAllCachedTokens();
+  IdentityTokenCache* token_cache();
 
   // GAIA id cache.
   void SetGaiaIdForExtension(const std::string& extension_id,
@@ -180,7 +126,7 @@ class IdentityAPI : public BrowserContextKeyedAPI,
   EventRouter* const event_router_;
 
   IdentityMintRequestQueue mint_queue_;
-  CachedTokens token_cache_;
+  IdentityTokenCache token_cache_;
 
   OnSignInChangedCallback on_signin_changed_callback_for_testing_;
 

chrome/browser/extensions/api/identity/identity_apitest.cc

@@ -940,7 +940,7 @@ class GetAuthTokenFunctionTest
   void SetCachedTokenForAccount(const CoreAccountId account_id,
                                 const IdentityTokenCacheValue& token_data) {
     ExtensionTokenKey key(extension_id_, account_id, oauth_scopes_);
-    id_api()->SetCachedToken(key, token_data);
+    id_api()->token_cache()->SetToken(key, token_data);
   }
 
   void SetCachedGaiaId(const std::string& gaia_id) {
@@ -952,7 +952,7 @@ class GetAuthTokenFunctionTest
     ExtensionTokenKey key(
         extension_id_, account_id.empty() ? GetPrimaryAccountId() : account_id,
         oauth_scopes_);
-    return id_api()->GetCachedToken(key);
+    return id_api()->token_cache()->GetToken(key);
   }
 
   base::Optional<std::string> GetCachedGaiaId() {
@@ -1603,7 +1603,7 @@ IN_PROC_BROWSER_TEST_F(GetAuthTokenFunctionTest,
     func->set_extension(extension.get());
     // Make sure we don't get a cached issue_advice result, which would cause
     // flow to be leaked.
-    id_api()->EraseAllCachedTokens();
+    id_api()->token_cache()->EraseAllTokens();
     func->push_mint_token_result(TestOAuth2MintTokenFlow::ISSUE_ADVICE_SUCCESS);
     func->set_scope_ui_oauth_error(test_case.oauth_error);
     std::string error = utils::RunFunctionAndReturnError(
@@ -2965,11 +2965,11 @@ class RemoveCachedAuthTokenFunctionTest : public ExtensionBrowserTest {
   void SetCachedToken(const IdentityTokenCacheValue& token_data) {
     ExtensionTokenKey key(kExtensionId, CoreAccountId("test@example.com"),
                           std::set<std::string>());
-    id_api()->SetCachedToken(key, token_data);
+    id_api()->token_cache()->SetToken(key, token_data);
   }
 
   const IdentityTokenCacheValue& GetCachedToken() {
-    return id_api()->GetCachedToken(
+    return id_api()->token_cache()->GetToken(
         ExtensionTokenKey(kExtensionId, CoreAccountId("test@example.com"),
                           std::set<std::string>()));
   }

chrome/browser/extensions/api/identity/identity_get_auth_token_function.cc

@@ -376,7 +376,7 @@ void IdentityGetAuthTokenFunction::StartSigninFlow() {
   // All cached tokens are invalid because the user is not signed in.
   IdentityAPI* id_api =
       extensions::IdentityAPI::GetFactoryInstance()->Get(GetProfile());
-  id_api->EraseAllCachedTokens();
+  id_api->token_cache()->EraseAllTokens();
 
   // If the signin flow fails, don't display the login prompt again.
   should_prompt_for_signin_ = false;
@@ -475,7 +475,8 @@ void IdentityGetAuthTokenFunction::StartMintToken(
 
   const OAuth2Info& oauth2_info = OAuth2Info::GetOAuth2Info(extension());
   IdentityAPI* id_api = IdentityAPI::GetFactoryInstance()->Get(GetProfile());
-  IdentityTokenCacheValue cache_entry = id_api->GetCachedToken(token_key_);
+  IdentityTokenCacheValue cache_entry =
+      id_api->token_cache()->GetToken(token_key_);
   IdentityTokenCacheValue::CacheValueStatus cache_status = cache_entry.status();
 
   if (type == IdentityMintRequestQueue::MINT_TYPE_NONINTERACTIVE) {
@@ -577,7 +578,8 @@ void IdentityGetAuthTokenFunction::OnMintTokenSuccess(
       access_token, base::TimeDelta::FromSeconds(time_to_live));
   IdentityAPI::GetFactoryInstance()
       ->Get(GetProfile())
-      ->SetCachedToken(token_key_, token);
+      ->token_cache()
+      ->SetToken(token_key_, token);
 
   CompleteMintTokenFlow();
   CompleteFunctionWithResult(access_token, granted_scopes);
@@ -617,7 +619,7 @@ void IdentityGetAuthTokenFunction::OnIssueAdviceSuccess(
 
   IdentityAPI* identity_api =
       IdentityAPI::GetFactoryInstance()->Get(GetProfile());
-  identity_api->SetCachedToken(
+  identity_api->token_cache()->SetToken(
       token_key_, IdentityTokenCacheValue::CreateIssueAdvice(issue_advice));
   // IssueAdvice doesn't communicate back to Chrome which account has been
   // chosen by the user. Cached gaia id may contain incorrect information so
@@ -639,8 +641,9 @@ void IdentityGetAuthTokenFunction::OnRemoteConsentSuccess(
 
   IdentityAPI::GetFactoryInstance()
       ->Get(GetProfile())
-      ->SetCachedToken(token_key_, IdentityTokenCacheValue::CreateRemoteConsent(
-                                       resolution_data));
+      ->token_cache()
+      ->SetToken(token_key_,
+                 IdentityTokenCacheValue::CreateRemoteConsent(resolution_data));
   should_prompt_for_signin_ = false;
   resolution_data_ = resolution_data;
   CompleteMintTokenFlow();
@@ -759,7 +762,8 @@ void IdentityGetAuthTokenFunction::OnGaiaFlowCompleted(
         access_token, base::TimeDelta::FromSeconds(time_to_live));
     IdentityAPI::GetFactoryInstance()
         ->Get(GetProfile())
-        ->SetCachedToken(token_key_, token_value);
+        ->token_cache()
+        ->SetToken(token_key_, token_value);
   }
 
   CompleteMintTokenFlow();
@@ -842,7 +846,7 @@ void IdentityGetAuthTokenFunction::OnGaiaRemoteConsentFlowApproved(
   // as this call may start a new request synchronously and query the cache.
   ExtensionTokenKey new_token_key(token_key_);
   new_token_key.account_id = account->account_id;
-  id_api->SetCachedToken(
+  id_api->token_cache()->SetToken(
       new_token_key,
       IdentityTokenCacheValue::CreateRemoteConsentApproved(consent_result));
   CompleteMintTokenFlow();

chrome/browser/extensions/api/identity/identity_remove_cached_auth_token_function.cc

@@ -26,7 +26,8 @@ ExtensionFunction::ResponseAction IdentityRemoveCachedAuthTokenFunction::Run() {
   EXTENSION_FUNCTION_VALIDATE(params.get());
   IdentityAPI::GetFactoryInstance()
       ->Get(browser_context())
-      ->EraseCachedToken(extension()->id(), params->details.token);
+      ->token_cache()
+      ->EraseToken(extension()->id(), params->details.token);
   return RespondNow(NoArguments());
 }
 

chrome/browser/extensions/api/identity/identity_token_cache.cc

+// Copyright 2020 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 "chrome/browser/extensions/api/identity/identity_token_cache.h"
+
+#include <algorithm>
+
+#include "base/stl_util.h"
+#include "chrome/browser/extensions/api/identity/identity_constants.h"
+
+namespace extensions {
+
+IdentityTokenCacheValue::IdentityTokenCacheValue() = default;
+IdentityTokenCacheValue::IdentityTokenCacheValue(
+    const IdentityTokenCacheValue& other) = default;
+IdentityTokenCacheValue& IdentityTokenCacheValue::operator=(
+    const IdentityTokenCacheValue& other) = default;
+IdentityTokenCacheValue::~IdentityTokenCacheValue() = default;
+
+// static
+IdentityTokenCacheValue IdentityTokenCacheValue::CreateIssueAdvice(
+    const IssueAdviceInfo& issue_advice) {
+  IdentityTokenCacheValue cache_value;
+  cache_value.status_ = CACHE_STATUS_ADVICE;
+  cache_value.issue_advice_ = issue_advice;
+  cache_value.expiration_time_ =
+      base::Time::Now() + base::TimeDelta::FromSeconds(
+                              identity_constants::kCachedIssueAdviceTTLSeconds);
+  return cache_value;
+}
+
+// static
+IdentityTokenCacheValue IdentityTokenCacheValue::CreateRemoteConsent(
+    const RemoteConsentResolutionData& resolution_data) {
+  IdentityTokenCacheValue cache_value;
+  cache_value.status_ = CACHE_STATUS_REMOTE_CONSENT;
+  cache_value.resolution_data_ = resolution_data;
+  cache_value.expiration_time_ =
+      base::Time::Now() + base::TimeDelta::FromSeconds(
+                              identity_constants::kCachedIssueAdviceTTLSeconds);
+  return cache_value;
+}
+
+// static
+IdentityTokenCacheValue IdentityTokenCacheValue::CreateRemoteConsentApproved(
+    const std::string& consent_result) {
+  IdentityTokenCacheValue cache_value;
+  cache_value.status_ = CACHE_STATUS_REMOTE_CONSENT_APPROVED;
+  cache_value.consent_result_ = consent_result;
+  cache_value.expiration_time_ =
+      base::Time::Now() + base::TimeDelta::FromSeconds(
+                              identity_constants::kCachedIssueAdviceTTLSeconds);
+  return cache_value;
+}
+
+// static
+IdentityTokenCacheValue IdentityTokenCacheValue::CreateToken(
+    const std::string& token,
+    base::TimeDelta time_to_live) {
+  IdentityTokenCacheValue cache_value;
+  cache_value.status_ = CACHE_STATUS_TOKEN;
+  cache_value.token_ = token;
+
+  // Remove 20 minutes from the ttl so cached tokens will have some time
+  // to live any time they are returned.
+  time_to_live -= base::TimeDelta::FromMinutes(20);
+
+  base::TimeDelta zero_delta;
+  if (time_to_live < zero_delta)
+    time_to_live = zero_delta;
+
+  cache_value.expiration_time_ = base::Time::Now() + time_to_live;
+  return cache_value;
+}
+
+IdentityTokenCacheValue::CacheValueStatus IdentityTokenCacheValue::status()
+    const {
+  if (is_expired())
+    return IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND;
+  else
+    return status_;
+}
+
+bool IdentityTokenCacheValue::is_expired() const {
+  return status_ == CACHE_STATUS_NOTFOUND ||
+         expiration_time_ < base::Time::Now();
+}
+
+const base::Time& IdentityTokenCacheValue::expiration_time() const {
+  return expiration_time_;
+}
+
+const IssueAdviceInfo& IdentityTokenCacheValue::issue_advice() const {
+  return issue_advice_;
+}
+
+const RemoteConsentResolutionData& IdentityTokenCacheValue::resolution_data()
+    const {
+  return resolution_data_;
+}
+
+const std::string& IdentityTokenCacheValue::consent_result() const {
+  return consent_result_;
+}
+
+const std::string& IdentityTokenCacheValue::token() const {
+  return token_;
+}
+
+IdentityTokenCache::IdentityTokenCache() = default;
+IdentityTokenCache::~IdentityTokenCache() = default;
+
+void IdentityTokenCache::SetToken(const ExtensionTokenKey& key,
+                                  const IdentityTokenCacheValue& token_data) {
+  auto it = token_cache_.find(key);
+  if (it != token_cache_.end() && it->second.status() <= token_data.status())
+    token_cache_.erase(it);
+
+  token_cache_.insert(std::make_pair(key, token_data));
+}
+
+void IdentityTokenCache::EraseToken(const std::string& extension_id,
+                                    const std::string& token) {
+  CachedTokens::iterator it;
+  for (it = token_cache_.begin(); it != token_cache_.end(); ++it) {
+    if (it->first.extension_id == extension_id &&
+        it->second.status() == IdentityTokenCacheValue::CACHE_STATUS_TOKEN &&
+        it->second.token() == token) {
+      token_cache_.erase(it);
+      break;
+    }
+  }
+}
+
+void IdentityTokenCache::EraseAllTokens() {
+  token_cache_.clear();
+}
+
+const IdentityTokenCacheValue& IdentityTokenCache::GetToken(
+    const ExtensionTokenKey& key) {
+  return token_cache_[key];
+}
+
+const IdentityTokenCache::CachedTokens& IdentityTokenCache::GetAllTokens() {
+  return token_cache_;
+}
+
+}  // namespace extensions

chrome/browser/extensions/api/identity/identity_token_cache.h

+// Copyright 2020 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 CHROME_BROWSER_EXTENSIONS_API_IDENTITY_IDENTITY_TOKEN_CACHE_H_
+#define CHROME_BROWSER_EXTENSIONS_API_IDENTITY_IDENTITY_TOKEN_CACHE_H_
+
+#include <map>
+#include <set>
+#include <string>
+
+#include "base/time/time.h"
+#include "chrome/browser/extensions/api/identity/extension_token_key.h"
+#include "google_apis/gaia/oauth2_mint_token_flow.h"
+
+namespace extensions {
+
+class IdentityTokenCacheValue {
+ public:
+  IdentityTokenCacheValue();
+  IdentityTokenCacheValue(const IdentityTokenCacheValue& other);
+  IdentityTokenCacheValue& operator=(const IdentityTokenCacheValue& other);
+  ~IdentityTokenCacheValue();
+
+  static IdentityTokenCacheValue CreateIssueAdvice(
+      const IssueAdviceInfo& issue_advice);
+  static IdentityTokenCacheValue CreateRemoteConsent(
+      const RemoteConsentResolutionData& resolution_data);
+  static IdentityTokenCacheValue CreateRemoteConsentApproved(
+      const std::string& consent_result);
+  static IdentityTokenCacheValue CreateToken(const std::string& token,
+                                             base::TimeDelta time_to_live);
+
+  // Order of these entries is used to determine whether or not new
+  // entries supersede older ones in SetCachedToken.
+  enum CacheValueStatus {
+    CACHE_STATUS_NOTFOUND,
+    CACHE_STATUS_ADVICE,
+    CACHE_STATUS_REMOTE_CONSENT,
+    CACHE_STATUS_REMOTE_CONSENT_APPROVED,
+    CACHE_STATUS_TOKEN
+  };
+
+  CacheValueStatus status() const;
+  const base::Time& expiration_time() const;
+
+  const IssueAdviceInfo& issue_advice() const;
+  const RemoteConsentResolutionData& resolution_data() const;
+  const std::string& consent_result() const;
+  const std::string& token() const;
+
+ private:
+  bool is_expired() const;
+
+  CacheValueStatus status_ = CACHE_STATUS_NOTFOUND;
+  base::Time expiration_time_;
+
+  // TODO(alexilin): This class holds at any given time one of the several
+  // possible types. Consider rewriting using absl::variant
+  IssueAdviceInfo issue_advice_;
+  RemoteConsentResolutionData resolution_data_;
+  std::string consent_result_;
+  std::string token_;
+};
+
+// In-memory cache of OAuth2 access tokens that are requested by extensions
+// through the `getAuthToken` API. Also caches intermediate short-lived values
+// used at different stages of the `getAuthToken` flow before a token is
+// obtained. The cache automatically handles token expiration. Extensions can
+// manually remove tokens from the cache using `removeCachedAuthToken` API.
+//
+// chrome://identity-internals provides a view of cache's content for debugging.
+class IdentityTokenCache {
+ public:
+  IdentityTokenCache();
+  ~IdentityTokenCache();
+  IdentityTokenCache(const IdentityTokenCache& other) = delete;
+  IdentityTokenCache& operator=(const IdentityTokenCache& other) = delete;
+
+  using CachedTokens = std::map<ExtensionTokenKey, IdentityTokenCacheValue>;
+
+  void SetToken(const ExtensionTokenKey& key,
+                const IdentityTokenCacheValue& token_data);
+  void EraseToken(const std::string& extension_id, const std::string& token);
+  void EraseAllTokens();
+  const IdentityTokenCacheValue& GetToken(const ExtensionTokenKey& key);
+  const CachedTokens& GetAllTokens();
+
+ private:
+  CachedTokens token_cache_;
+};
+
+}  // namespace extensions
+
+#endif  // CHROME_BROWSER_EXTENSIONS_API_IDENTITY_IDENTITY_TOKEN_CACHE_H_

chrome/browser/extensions/api/identity/identity_token_cache_unittest.cc

+// Copyright 2020 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 "chrome/browser/extensions/api/identity/identity_token_cache.h"
+
+#include <set>
+#include <string>
+
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace extensions {
+
+namespace {
+
+const char kDefaultExtensionId[] = "ext_id";
+
+}  // namespace
+
+class IdentityTokenCacheTest : public testing::Test {
+ public:
+  void SetAccessToken(const std::string& ext_id,
+                      const std::string& token_string,
+                      const std::set<std::string>& scopes) {
+    ExtensionTokenKey key(ext_id, CoreAccountId(), scopes);
+
+    IdentityTokenCacheValue token = IdentityTokenCacheValue::CreateToken(
+        token_string, base::TimeDelta::FromSeconds(3600));
+    cache_.SetToken(key, token);
+  }
+
+  void SetRemoteConsentApprovedToken(const std::string& ext_id,
+                                     const std::string& consent_result,
+                                     const std::set<std::string>& scopes) {
+    ExtensionTokenKey key(ext_id, CoreAccountId(), scopes);
+    IdentityTokenCacheValue token =
+        IdentityTokenCacheValue::CreateRemoteConsentApproved(consent_result);
+    cache_.SetToken(key, token);
+  }
+
+  const IdentityTokenCacheValue& GetToken(const std::string& ext_id,
+                                          const std::set<std::string>& scopes) {
+    ExtensionTokenKey key(ext_id, CoreAccountId(), scopes);
+    return cache_.GetToken(key);
+  }
+
+  IdentityTokenCache& cache() { return cache_; }
+
+ private:
+  IdentityTokenCache cache_;
+};
+
+TEST_F(IdentityTokenCacheTest, AccessTokenCacheHit) {
+  std::string token_string = "token";
+  std::set<std::string> scopes = {"foo", "bar"};
+  SetAccessToken(kDefaultExtensionId, token_string, scopes);
+
+  const IdentityTokenCacheValue& cached_token =
+      GetToken(kDefaultExtensionId, scopes);
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN, cached_token.status());
+  EXPECT_EQ(token_string, cached_token.token());
+}
+
+TEST_F(IdentityTokenCacheTest, IntermediateValueCacheHit) {
+  std::string consent_result = "result";
+  std::set<std::string> scopes = {"foo", "bar"};
+  SetRemoteConsentApprovedToken(kDefaultExtensionId, consent_result, scopes);
+
+  const IdentityTokenCacheValue& cached_token =
+      GetToken(kDefaultExtensionId, scopes);
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_REMOTE_CONSENT_APPROVED,
+            cached_token.status());
+  EXPECT_EQ(consent_result, cached_token.consent_result());
+}
+
+TEST_F(IdentityTokenCacheTest, CacheHitPriority) {
+  std::string token_string = "token";
+  std::set<std::string> scopes = {"foo", "bar"};
+  SetAccessToken(kDefaultExtensionId, token_string, scopes);
+  SetRemoteConsentApprovedToken(kDefaultExtensionId, "result", scopes);
+
+  // Prioritize access tokens over immediate values.
+  const IdentityTokenCacheValue& cached_token =
+      GetToken(kDefaultExtensionId, scopes);
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN, cached_token.status());
+  EXPECT_EQ(token_string, cached_token.token());
+}
+
+TEST_F(IdentityTokenCacheTest, AccessTokenCacheMiss) {
+  std::string ext_1 = "ext_1";
+  std::set<std::string> scopes_1 = {"foo", "bar"};
+  SetAccessToken(ext_1, "token_1", scopes_1);
+
+  std::string ext_2 = "ext_2";
+  std::set<std::string> scopes_2 = {"foo", "foobar"};
+  SetAccessToken(ext_2, "token_2", scopes_2);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(ext_2, scopes_1).status());
+}
+
+TEST_F(IdentityTokenCacheTest, IntermediateValueCacheMiss) {
+  std::string ext_1 = "ext_1";
+  std::set<std::string> scopes_1 = {"foo", "bar"};
+  SetRemoteConsentApprovedToken(ext_1, "result_1", scopes_1);
+
+  std::string ext_2 = "ext_2";
+  std::set<std::string> scopes_2 = {"foo", "foobar"};
+  SetRemoteConsentApprovedToken(ext_2, "result_2", scopes_2);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(ext_2, scopes_1).status());
+}
+
+TEST_F(IdentityTokenCacheTest, EraseToken) {
+  std::string token_string = "token";
+  std::set<std::string> scopes = {"foo", "bar"};
+  SetAccessToken(kDefaultExtensionId, token_string, scopes);
+
+  cache().EraseToken(kDefaultExtensionId, token_string);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(kDefaultExtensionId, scopes).status());
+}
+
+TEST_F(IdentityTokenCacheTest, EraseTokenOthersUnaffected) {
+  std::string token_string = "token";
+  std::set<std::string> scopes = {"foo", "bar"};
+  SetAccessToken(kDefaultExtensionId, token_string, scopes);
+
+  std::string unrelated_token_string = "unrelated_token";
+  std::set<std::string> unrelated_scopes = {"foo", "foobar"};
+  SetAccessToken(kDefaultExtensionId, unrelated_token_string, unrelated_scopes);
+
+  cache().EraseToken(kDefaultExtensionId, token_string);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(kDefaultExtensionId, scopes).status());
+
+  const IdentityTokenCacheValue& cached_token =
+      GetToken(kDefaultExtensionId, unrelated_scopes);
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN, cached_token.status());
+  EXPECT_EQ(unrelated_token_string, cached_token.token());
+}
+
+TEST_F(IdentityTokenCacheTest, EraseAllTokens) {
+  std::string ext_1 = "ext_1";
+  std::string token_string = "token_1";
+  std::set<std::string> scopes_1 = {"foo", "bar"};
+  SetAccessToken(ext_1, token_string, scopes_1);
+
+  std::string ext_2 = "ext_2";
+  std::string remote_consent = "approved";
+  std::set<std::string> scopes_2 = {"foo", "foobar"};
+  SetRemoteConsentApprovedToken(ext_2, remote_consent, scopes_2);
+
+  cache().EraseAllTokens();
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(ext_1, scopes_1).status());
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(ext_2, scopes_2).status());
+}
+
+TEST_F(IdentityTokenCacheTest, GetAllTokens) {
+  std::string ext_1 = "ext_1";
+  std::string token_string_1 = "token_1";
+  std::set<std::string> scopes_1 = {"foo", "bar"};
+  SetAccessToken(ext_1, token_string_1, scopes_1);
+
+  std::string ext_2 = "ext_2";
+  std::string token_string_2 = "token_2";
+  std::set<std::string> scopes_2 = {"foobar"};
+  SetAccessToken(ext_2, token_string_2, scopes_2);
+
+  IdentityTokenCache::CachedTokens cached_tokens = cache().GetAllTokens();
+  EXPECT_EQ(2ul, cached_tokens.size());
+
+  ExtensionTokenKey key_1(ext_1, CoreAccountId(), scopes_1);
+  const IdentityTokenCacheValue& cached_token_1 = cached_tokens[key_1];
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
+            cached_token_1.status());
+  EXPECT_EQ(token_string_1, cached_token_1.token());
+
+  ExtensionTokenKey key_2(ext_2, CoreAccountId(), scopes_2);
+  const IdentityTokenCacheValue& cached_token_2 = cached_tokens[key_2];
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_TOKEN,
+            cached_token_2.status());
+  EXPECT_EQ(token_string_2, cached_token_2.token());
+}
+
+// Newly cached access tokens should override previously cached values with the
+// same scopes.
+TEST_F(IdentityTokenCacheTest, OverrideAccessToken) {
+  std::set<std::string> scopes = {"foo", "bar", "foobar"};
+  SetAccessToken(kDefaultExtensionId, "token1", scopes);
+
+  std::string override_token = "token_2";
+  SetAccessToken(kDefaultExtensionId, override_token, scopes);
+  cache().EraseToken(kDefaultExtensionId, override_token);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(kDefaultExtensionId, scopes).status());
+}
+
+TEST_F(IdentityTokenCacheTest, OverrideIntermediateToken) {
+  std::set<std::string> scopes = {"foo", "bar", "foobar"};
+  SetRemoteConsentApprovedToken(kDefaultExtensionId, "result", scopes);
+
+  std::string override_token = "token";
+  SetAccessToken(kDefaultExtensionId, override_token, scopes);
+  cache().EraseToken(kDefaultExtensionId, override_token);
+
+  EXPECT_EQ(IdentityTokenCacheValue::CACHE_STATUS_NOTFOUND,
+            GetToken(kDefaultExtensionId, scopes).status());
+}
+
+}  // namespace extensions

chrome/browser/ui/webui/identity_internals_ui.cc

@@ -140,8 +140,8 @@ void IdentityInternalsUIMessageHandler::OnTokenRevokerDone(
   CHECK(api);
 
   // Remove token from the cache.
-  api->EraseCachedToken(token_revoker->extension_id(),
-                        token_revoker->access_token());
+  api->token_cache()->EraseToken(token_revoker->extension_id(),
+                                 token_revoker->access_token());
 
   // Update view about the token being removed.
   base::ListValue result;
@@ -223,16 +223,15 @@ IdentityInternalsUIMessageHandler::GetInfoForToken(
 void IdentityInternalsUIMessageHandler::GetInfoForAllTokens(
     const base::ListValue* args) {
   base::ListValue results;
-  extensions::IdentityAPI::CachedTokens tokens;
+  extensions::IdentityTokenCache::CachedTokens tokens;
   // The API can be null in incognito.
   extensions::IdentityAPI* api =
       extensions::IdentityAPI::GetFactoryInstance()->Get(
           Profile::FromWebUI(web_ui()));
   if (api)
-    tokens = api->GetAllCachedTokens();
-  for (extensions::IdentityAPI::CachedTokens::const_iterator
-           iter = tokens.begin(); iter != tokens.end(); ++iter) {
-    results.Append(GetInfoForToken(iter->first, iter->second));
+    tokens = api->token_cache()->GetAllTokens();
+  for (const auto& token : tokens) {
+    results.Append(GetInfoForToken(token.first, token.second));
   }
 
   web_ui()->CallJavascriptFunctionUnsafe("identity_internals.returnTokens",

chrome/browser/ui/webui/identity_internals_ui_browsertest.cc

@@ -59,5 +59,6 @@ void IdentityInternalsUIBrowserTest::AddTokenToCache(
       std::set<std::string>(scopes.begin(), scopes.end()));
   extensions::IdentityAPI::GetFactoryInstance()
       ->Get(browser()->profile())
-      ->SetCachedToken(key, token_cache_value);
+      ->token_cache()
+      ->SetToken(key, token_cache_value);
 }

chrome/test/BUILD.gn

@@ -4794,6 +4794,7 @@ test("unit_tests") {
       "../browser/extensions/api/identity/gaia_web_auth_flow_unittest.cc",
       "../browser/extensions/api/identity/identity_api_unittest.cc",
       "../browser/extensions/api/identity/identity_mint_queue_unittest.cc",
+      "../browser/extensions/api/identity/identity_token_cache_unittest.cc",
       "../browser/extensions/api/image_writer_private/destroy_partitions_operation_unittest.cc",
       "../browser/extensions/api/image_writer_private/image_writer_private_api_unittest.cc",
       "../browser/extensions/api/image_writer_private/operation_manager_unittest.cc",