Add a class for retrieving and caching instance identity tokens

VM instance identity tokens are good for ~1 hour so we can reuse
the same token during that time for service requests. This class
will be used in the heartbeat and session authorization impls so
we don't request new tokens for every request to the CRD backend.

Bug: 388885661
Change-Id: I6e8f6d62eb6a307f9984b2e75af80fc92786050e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6299649
Auto-Submit: Joe Downing <joedow@chromium.org>
Reviewed-by: Yuwei Huang <yuweih@chromium.org>
Commit-Queue: Joe Downing <joedow@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1424285}

remoting/base/BUILD.gn

@@ -78,6 +78,8 @@ source_set("base") {
     "directory_service_client.h",
     "fqdn.cc",
     "fqdn.h",
+    "instance_identity_token_getter.cc",
+    "instance_identity_token_getter.h",
     "internal_headers.h",
     "is_google_email.cc",
     "is_google_email.h",
@@ -400,6 +402,7 @@ source_set("unit_tests") {
     "buffered_socket_writer_unittest.cc",
     "capabilities_unittest.cc",
     "compound_buffer_unittest.cc",
+    "instance_identity_token_getter_unittest.cc",
     "leaky_bucket_unittest.cc",
     "local_session_policies_provider_unittest.cc",
     "oauth_token_getter_proxy_unittest.cc",

remoting/base/compute_engine_service_client.cc

@@ -32,7 +32,7 @@ namespace {
 // https://cloud.google.com/compute/docs/metadata/querying-metadata#query-https-mds
 constexpr char kHttpMetadataBaseUrl[] =
     "http://metadata.google.internal/computeMetadata/v1/instance/"
-    "service-accounts/default/";
+    "service-accounts/default";
 
 constexpr size_t kMaxResponseSize = 4096;
 

remoting/base/instance_identity_token_getter.cc

@@ -0,0 +1,77 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "remoting/base/instance_identity_token_getter.h"
+
+#include <string>
+#include <utility>
+
+#include "base/functional/callback.h"
+#include "base/location.h"
+#include "base/memory/scoped_refptr.h"
+#include "base/memory/weak_ptr.h"
+#include "base/sequence_checker.h"
+#include "base/task/sequenced_task_runner.h"
+#include "base/time/time.h"
+#include "net/base/net_errors.h"
+#include "remoting/base/http_status.h"
+#include "services/network/public/cpp/shared_url_loader_factory.h"
+
+namespace remoting {
+
+namespace {
+// Per Cloud documentation, the identity token is valid for ~1 hour.
+// TODO: joedow - Parse the identity token to set the exact expiration time.
+constexpr base::TimeDelta kTokenLifetime = base::Minutes(50);
+}  // namespace
+
+InstanceIdentityTokenGetter::InstanceIdentityTokenGetter(
+    std::string_view audience,
+    scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory)
+    : audience_(audience), compute_engine_service_client_(url_loader_factory) {}
+
+InstanceIdentityTokenGetter::~InstanceIdentityTokenGetter() = default;
+
+void InstanceIdentityTokenGetter::RetrieveToken(TokenCallback on_token) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+
+  // Check expiration, clear token if no longer valid.
+  if ((last_fetch_time_ + kTokenLifetime) < base::Time::Now()) {
+    identity_token_.clear();
+  }
+
+  if (!identity_token_.empty()) {
+    base::SequencedTaskRunner::GetCurrentDefault()->PostTask(
+        FROM_HERE, base::BindOnce(std::move(on_token), identity_token_));
+    return;
+  }
+
+  queued_callbacks_.emplace_back(std::move(on_token));
+  // Only make a service request for the first caller, the rest will be queued
+  // and provided with the token when the request completes.
+  if (queued_callbacks_.size() == 1) {
+    compute_engine_service_client_.GetInstanceIdentityToken(
+        audience_,
+        base::BindOnce(&InstanceIdentityTokenGetter::OnTokenRetrieved,
+                       weak_ptr_factory_.GetWeakPtr()));
+  }
+}
+
+void InstanceIdentityTokenGetter::OnTokenRetrieved(const HttpStatus& response) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+
+  if (response.ok()) {
+    identity_token_ = response.response_body();
+    last_fetch_time_ = base::Time::Now();
+  }
+
+  // TODO: joedow - Add token validation and check expiration time.
+
+  auto callbacks = std::move(queued_callbacks_);
+  for (auto& callback : callbacks) {
+    std::move(callback).Run(identity_token_);
+  }
+}
+
+}  // namespace remoting

remoting/base/instance_identity_token_getter.h

@@ -0,0 +1,76 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef REMOTING_BASE_INSTANCE_IDENTITY_TOKEN_GETTER_H_
+#define REMOTING_BASE_INSTANCE_IDENTITY_TOKEN_GETTER_H_
+
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "base/functional/callback_forward.h"
+#include "base/memory/scoped_refptr.h"
+#include "base/memory/weak_ptr.h"
+#include "base/sequence_checker.h"
+#include "base/time/time.h"
+#include "remoting/base/compute_engine_service_client.h"
+#include "remoting/base/http_status.h"
+
+namespace network {
+class SharedURLLoaderFactory;
+}  // namespace network
+
+namespace remoting {
+
+// InstanceIdentityTokenGetter caches instance identity tokens for Compute
+// Engine service requests and refreshes them as needed.
+class InstanceIdentityTokenGetter {
+ public:
+  using TokenCallback = base::OnceCallback<void(std::string_view)>;
+
+  InstanceIdentityTokenGetter(
+      std::string_view audience,
+      scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory);
+
+  InstanceIdentityTokenGetter(const InstanceIdentityTokenGetter&) = delete;
+  InstanceIdentityTokenGetter& operator=(const InstanceIdentityTokenGetter&) =
+      delete;
+
+  ~InstanceIdentityTokenGetter();
+
+  // Calls |on_token| with an identity token, or empty in the case the request
+  // fails. The token returned has a lifetime of at least 10 minutes and should
+  // not be cached.
+  void RetrieveToken(TokenCallback on_token);
+
+ private:
+  void OnTokenRetrieved(const HttpStatus& response);
+
+  // An identifier which is embedded in the identity token. Typically this will
+  // be the URL of the API which the token used to access but could be an
+  // arbitrary identifier or string.
+  const std::string audience_;
+
+  // The instance identity token from the last successful fetch operation.
+  std::string identity_token_ GUARDED_BY_CONTEXT(sequence_checker_);
+
+  // The timestamp of the last successful identity token fetch operation.
+  base::Time last_fetch_time_ GUARDED_BY_CONTEXT(sequence_checker_);
+
+  // The set of callbacks to run after fetching an updated identity token.
+  std::vector<TokenCallback> queued_callbacks_
+      GUARDED_BY_CONTEXT(sequence_checker_);
+
+  // Used to request an instance identity token.
+  ComputeEngineServiceClient compute_engine_service_client_
+      GUARDED_BY_CONTEXT(sequence_checker_);
+
+  SEQUENCE_CHECKER(sequence_checker_);
+
+  base::WeakPtrFactory<InstanceIdentityTokenGetter> weak_ptr_factory_{this};
+};
+
+}  // namespace remoting
+
+#endif  // REMOTING_BASE_INSTANCE_IDENTITY_TOKEN_GETTER_H_

remoting/base/instance_identity_token_getter_unittest.cc

@@ -0,0 +1,240 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "remoting/base/instance_identity_token_getter.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <utility>
+
+#include "base/functional/bind.h"
+#include "base/memory/ref_counted.h"
+#include "base/run_loop.h"
+#include "base/test/task_environment.h"
+#include "net/http/http_status_code.h"
+#include "services/network/public/cpp/weak_wrapper_shared_url_loader_factory.h"
+#include "services/network/test/test_shared_url_loader_factory.h"
+#include "services/network/test/test_url_loader_factory.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace remoting {
+
+namespace {
+constexpr char kTestAudience[] = "audience_for_testing";
+constexpr char kTokenBodyResponse[] = "instance_identity_token";
+// Matches the URL generated for requests from ComputeEngineServiceClient.
+constexpr char kHttpMetadataRequestUrl[] =
+    "http://metadata.google.internal/computeMetadata/v1/instance/"
+    "service-accounts/default/identity?audience=audience_for_testing&"
+    "format=full";
+}  // namespace
+
+class InstanceIdentityTokenGetterTest : public testing::Test {
+ public:
+  InstanceIdentityTokenGetterTest();
+  ~InstanceIdentityTokenGetterTest() override;
+
+  void SetUp() override;
+
+  void OnTokenRetrieved(std::string_view token);
+
+ protected:
+  void RunUntilQuit();
+  void SetTokenResponse(std::string_view response_body);
+  void SetErrorResponse(net::HttpStatusCode status);
+  void ResetQuitClosure();
+  void ClearTokenResponse();
+  void FastForwardBy(base::TimeDelta duration);
+
+  InstanceIdentityTokenGetter& instance_identity_token_getter() {
+    return *instance_identity_token_getter_;
+  }
+
+  void set_pending_callback_count(int count) {
+    pending_callback_count_ = count;
+  }
+
+  const std::optional<std::string>& token() { return token_; }
+
+  size_t url_loader_request_count() {
+    return test_url_loader_factory_.total_requests();
+  }
+
+ private:
+  int pending_callback_count_ = 1;
+  std::optional<std::string> token_;
+
+  base::RepeatingClosure quit_closure_;
+
+  base::test::TaskEnvironment task_environment_{
+      base::test::TaskEnvironment::MainThreadType::IO,
+      base::test::TaskEnvironment::TimeSource::MOCK_TIME};
+
+  network::TestURLLoaderFactory test_url_loader_factory_;
+  scoped_refptr<network::SharedURLLoaderFactory> shared_url_loader_factory_;
+
+  std::unique_ptr<InstanceIdentityTokenGetter> instance_identity_token_getter_;
+};
+
+InstanceIdentityTokenGetterTest::InstanceIdentityTokenGetterTest() = default;
+InstanceIdentityTokenGetterTest::~InstanceIdentityTokenGetterTest() = default;
+
+void InstanceIdentityTokenGetterTest::SetUp() {
+  shared_url_loader_factory_ = test_url_loader_factory_.GetSafeWeakWrapper();
+  instance_identity_token_getter_ =
+      std::make_unique<InstanceIdentityTokenGetter>(kTestAudience,
+                                                    shared_url_loader_factory_);
+  quit_closure_ = task_environment_.QuitClosure();
+}
+
+void InstanceIdentityTokenGetterTest::RunUntilQuit() {
+  task_environment_.RunUntilQuit();
+}
+
+void InstanceIdentityTokenGetterTest::SetTokenResponse(
+    std::string_view response_body) {
+  ClearTokenResponse();
+  test_url_loader_factory_.AddResponse(kHttpMetadataRequestUrl, response_body);
+}
+
+void InstanceIdentityTokenGetterTest::ClearTokenResponse() {
+  test_url_loader_factory_.ClearResponses();
+}
+
+void InstanceIdentityTokenGetterTest::SetErrorResponse(
+    net::HttpStatusCode status) {
+  ClearTokenResponse();
+  test_url_loader_factory_.AddResponse(kHttpMetadataRequestUrl,
+                                       /*content=*/std::string(), status);
+}
+
+void InstanceIdentityTokenGetterTest::ResetQuitClosure() {
+  ASSERT_TRUE(quit_closure_.is_null());
+  quit_closure_ = task_environment_.QuitClosure();
+}
+
+void InstanceIdentityTokenGetterTest::OnTokenRetrieved(std::string_view token) {
+  // If the callback has been run previously, make sure each callback receives
+  // the same value.
+  if (token_.has_value()) {
+    ASSERT_EQ(*token_, token);
+  } else {
+    token_ = token;
+  }
+
+  pending_callback_count_--;
+  if (pending_callback_count_ == 0) {
+    std::move(quit_closure_).Run();
+  }
+}
+
+void InstanceIdentityTokenGetterTest::FastForwardBy(base::TimeDelta duration) {
+  task_environment_.FastForwardBy(duration);
+}
+
+TEST_F(InstanceIdentityTokenGetterTest, SingleRequest) {
+  SetTokenResponse(kTokenBodyResponse);
+
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+  ASSERT_EQ(url_loader_request_count(), 1U);
+}
+
+TEST_F(InstanceIdentityTokenGetterTest, MultipleRequests) {
+  const int kQueuedCallbackCount = 10;
+
+  set_pending_callback_count(kQueuedCallbackCount);
+  for (int i = 0; i < kQueuedCallbackCount; i++) {
+    instance_identity_token_getter().RetrieveToken(
+        base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                       base::Unretained(this)));
+  }
+
+  SetTokenResponse(kTokenBodyResponse);
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+  ASSERT_EQ(url_loader_request_count(), 1U);
+}
+
+TEST_F(InstanceIdentityTokenGetterTest, CachedTokenReturned) {
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  SetTokenResponse(kTokenBodyResponse);
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+
+  // Call a second time and verify a token is provided w/o calling the service.
+  ClearTokenResponse();
+  ResetQuitClosure();
+  set_pending_callback_count(1);
+
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+  ASSERT_EQ(url_loader_request_count(), 1U);
+}
+
+TEST_F(InstanceIdentityTokenGetterTest, CachedTokenIgnored) {
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  SetTokenResponse(kTokenBodyResponse);
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+
+  // Call again and verify a token is provided after calling the service.
+  FastForwardBy(base::Hours(1));
+  ResetQuitClosure();
+  set_pending_callback_count(1);
+
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_EQ(*token(), kTokenBodyResponse);
+  ASSERT_EQ(url_loader_request_count(), 2U);
+}
+
+TEST_F(InstanceIdentityTokenGetterTest, ServiceFailureReturnsEmptyString) {
+  instance_identity_token_getter().RetrieveToken(
+      base::BindOnce(&InstanceIdentityTokenGetterTest::OnTokenRetrieved,
+                     base::Unretained(this)));
+
+  SetErrorResponse(net::HTTP_BAD_REQUEST);
+
+  RunUntilQuit();
+
+  ASSERT_TRUE(token().has_value());
+  ASSERT_TRUE(token()->empty());
+}
+
+}  // namespace remoting