Share S3 credentials provider across Context and fix anonymous-fallback caching

tensorstore/internal/aws/credentials/common.cc

@@ -36,13 +36,15 @@ namespace internal_aws {
 namespace {
 
 inline AwsCredentialsProvider AsProvider(aws_credentials_provider* p) {
+  if (!p) return nullptr;
   assert(aws_atomic_load_int(&p->ref_count) == 1);
   return AwsCredentialsProvider(p, internal::adopt_object_ref);
 }
 
 }  // namespace
 
 AwsCredentialsProvider MakeCache(AwsCredentialsProvider provider) {
+  if (!provider) return nullptr;
   aws_allocator* allocator = GetAwsAllocator();
 
   aws_credentials_provider_cached_options options;
@@ -73,6 +75,7 @@ AwsCredentialsProvider MakeDefaultWithAnonymous(
 
   auto default_provider = MakeDefault(profile_name_override);
   auto anonymous_provider = MakeAnonymous();
+  if (!default_provider || !anonymous_provider) return nullptr;
 
   aws_credentials_provider* providers[2];
   AWS_ZERO_ARRAY(providers);

tensorstore/internal/aws/credentials/common.h

@@ -23,10 +23,13 @@ namespace tensorstore {
 namespace internal_aws {
 
 /// Returns a credentials provider that uses the AWS default credentials chain.
+/// The returned provider caches resolved credentials internally.
 AwsCredentialsProvider MakeDefault(std::string_view profile_name_override);
 
 /// Returns a credentials provider that uses the AWS default credentials chain
 /// with anonymous credentials as a fallback.
+/// The default chain caches resolved credentials internally, while the
+/// anonymous fallback is not cached; do not wrap this provider in `MakeCache`.
 AwsCredentialsProvider MakeDefaultWithAnonymous(
     std::string_view profile_name_override);
 
@@ -53,7 +56,10 @@ AwsCredentialsProvider MakeEcsRole(std::string_view endpoint,
 /// NOTE: Add additional credentials providers such as STS, x509, SSO, Cognito,
 /// etc.
 
-/// Returns a credentials provider that caches the credentials.
+/// Returns a credentials provider that caches resolved credentials (refreshing
+/// expiring credentials 5 minutes before expiration, and caching non-expiring
+/// credentials indefinitely). Returns nullptr if `provider` is nullptr. Must
+/// not wrap a provider chain with an anonymous fallback.
 AwsCredentialsProvider MakeCache(AwsCredentialsProvider provider);
 
 }  // namespace internal_aws

tensorstore/kvstore/s3/BUILD

@@ -84,6 +84,7 @@ tensorstore_cc_library(
         "//tensorstore:context",
         "//tensorstore:json_serialization_options",
         "//tensorstore/internal:data_copy_concurrency_resource",
+        "//tensorstore/internal/aws:aws_credentials",
         "//tensorstore/internal/cache_key",
         "//tensorstore/internal/cache_key:json",
         "//tensorstore/internal/json_binding",
@@ -100,6 +101,7 @@ tensorstore_cc_test(
     deps = [
         ":aws_credentials_resource",
         "//tensorstore:context",
+        "//tensorstore/internal/aws:aws_credentials",
         "//tensorstore/util:status_testutil",
         "@abseil-cpp//absl/status",
         "@googletest//:gtest_main",
@@ -419,8 +421,11 @@ tensorstore_cc_test(
     deps = [
         ":aws_credentials_spec",
         "//tensorstore:json_serialization_options_base",
+        "//tensorstore/internal:env",
+        "//tensorstore/internal/aws:aws_credentials",
         "//tensorstore/internal/testing:json_gtest",
         "//tensorstore/util:status_testutil",
+        "@abseil-cpp//absl/cleanup",
         "@abseil-cpp//absl/status",
         "@googletest//:gtest_main",
         "@nlohmann_json//:json",

tensorstore/kvstore/s3/aws_credentials_resource.h

@@ -18,9 +18,11 @@
 #include <stddef.h>
 
 #include <cassert>
+#include <utility>
 
 #include "tensorstore/context.h"
 #include "tensorstore/context_resource_provider.h"
+#include "tensorstore/internal/aws/aws_credentials.h"
 #include "tensorstore/internal/json_binding/json_binding.h"
 #include "tensorstore/json_serialization_options.h"
 #include "tensorstore/kvstore/s3/aws_credentials_spec.h"
@@ -33,10 +35,12 @@ namespace internal_kvstore_s3 {
 struct AwsCredentialsResource
     : public internal::ContextResourceTraits<AwsCredentialsResource> {
   static constexpr char id[] = "aws_credentials";
-  constexpr static bool config_only = true;
 
   using Spec = AwsCredentialsSpec;
-  using Resource = Spec;
+  struct Resource {
+    Spec spec;
+    internal_aws::AwsCredentialsProvider credentials_provider;
+  };
 
   static Spec Default() { return Spec{}; }
   static constexpr auto JsonBinder() {
@@ -45,12 +49,14 @@ struct AwsCredentialsResource
 
   static Result<Resource> Create(
       const Spec& spec, internal::ContextResourceCreationContext context) {
-    return spec;
+    TENSORSTORE_ASSIGN_OR_RETURN(auto credentials_provider,
+                                 MakeAwsCredentialsProvider(spec));
+    return Resource{spec, std::move(credentials_provider)};
   }
 
   static Spec GetSpec(const Resource& resource,
                       const internal::ContextSpecBuilder& builder) {
-    return resource;
+    return resource.spec;
   }
 };
 

tensorstore/kvstore/s3/aws_credentials_resource_test.cc

@@ -21,6 +21,7 @@
 #include "absl/status/status.h"
 #include <nlohmann/json.hpp>
 #include "tensorstore/context.h"
+#include "tensorstore/internal/aws/aws_credentials.h"
 #include "tensorstore/util/status_testutil.h"
 
 namespace {
@@ -55,6 +56,7 @@ TEST(AwsCredentialsResourceTest, Default) {
 
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto resource,
                                    context.GetResource(resource_spec));
+  EXPECT_THAT(resource->credentials_provider.get(), ::testing::NotNull());
 }
 
 TEST(AwsCredentialsResourceTest, ExplicitDefault) {
@@ -66,17 +68,34 @@ TEST(AwsCredentialsResourceTest, ExplicitDefault) {
 
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto resource,
                                    context.GetResource(resource_spec));
+  EXPECT_THAT(resource->credentials_provider.get(), ::testing::NotNull());
 }
 
 TEST(AwsCredentialsResourceTest, ValidSpec) {
+  auto context = Context::Default();
+  for (const auto& json : {::nlohmann::json{{"type", "default"},
+                                            {"profile", "my_profile"}},
+                           ::nlohmann::json{{"type", "environment"}}}) {
+    TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+        auto resource_spec,
+        Context::Resource<AwsCredentialsResource>::FromJson(json));
+    TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto resource,
+                                     context.GetResource(resource_spec));
+    EXPECT_THAT(resource->credentials_provider.get(), ::testing::NotNull());
+  }
+}
+
+TEST(AwsCredentialsResourceTest, ProfileSpecMissingFilesFails) {
   auto context = Context::Default();
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(
-      auto resource_spec,
+      auto profile_spec,
       Context::Resource<AwsCredentialsResource>::FromJson(
-          {{"type", "profile"}, {"profile", "my_profile"}}));
-
-  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto resource,
-                                   context.GetResource(resource_spec));
+          {{"type", "profile"},
+           {"profile", "my_profile"},
+           {"credentials_file", "/nonexistent/credentials"},
+           {"config_file", "/nonexistent/config"}}));
+  EXPECT_THAT(context.GetResource(profile_spec),
+              StatusIs(absl::StatusCode::kInternal));
 }
 
 TEST(AwsCredentialsResourceTest, ValidAnonymousSpec) {
@@ -87,6 +106,27 @@ TEST(AwsCredentialsResourceTest, ValidAnonymousSpec) {
 
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto resource,
                                    context.GetResource(resource_spec));
+  EXPECT_THAT(resource->credentials_provider.get(), ::testing::NotNull());
+}
+
+TEST(AwsCredentialsResourceTest, SharesProviderWithinContext) {
+  auto parent = Context::Default();
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+      auto resource_spec,
+      Context::Resource<AwsCredentialsResource>::FromJson("aws_credentials"));
+
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto r1, parent.GetResource(resource_spec));
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto r2, parent.GetResource(resource_spec));
+  EXPECT_THAT(r1->credentials_provider.get(), ::testing::NotNull());
+  EXPECT_EQ(r1->credentials_provider.get(), r2->credentials_provider.get());
+
+  auto child = Context(Context::Spec{}, parent);
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto r3, child.GetResource(resource_spec));
+  EXPECT_EQ(r1->credentials_provider.get(), r3->credentials_provider.get());
+
+  auto other = Context::Default();
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto r4, other.GetResource(resource_spec));
+  EXPECT_NE(r1->credentials_provider.get(), r4->credentials_provider.get());
 }
 
 TEST(AwsCredentialsResourceTest, InvalidSpecs) {

tensorstore/kvstore/s3/aws_credentials_spec.cc

@@ -155,17 +155,24 @@ Result<AwsCredentialsProvider> MakeAwsCredentialsProvider(const Spec& spec) {
   struct MakeCredentialsVisitor {
     using R = AwsCredentialsProvider;
 
+    // MakeDefault in MakeDefaultWithAnonymous already caches the default chain
+    // internally, and its anonymous fallback must not be cached; MakeAnonymous
+    // needs no cache.
     R operator()(const Spec::Default& v) {
       return MakeDefaultWithAnonymous(v.profile);
     }
     R operator()(const Spec::Anonymous&) { return MakeAnonymous(); }
-    R operator()(const Spec::Environment&) { return MakeEnvironment(); }
-    R operator()(const Spec::Imds&) { return MakeImds(); }
+    R operator()(const Spec::Environment&) {
+      return MakeCache(MakeEnvironment());
+    }
+    R operator()(const Spec::Imds&) { return MakeCache(MakeImds()); }
     R operator()(const Spec::Profile& v) {
-      return MakeProfile(v.profile, v.credentials_file, v.config_file);
+      return MakeCache(
+          MakeProfile(v.profile, v.credentials_file, v.config_file));
     }
     R operator()(const Spec::EcsRole& v) {
-      return MakeEcsRole(v.endpoint, v.auth_token_file, v.auth_token);
+      return MakeCache(
+          MakeEcsRole(v.endpoint, v.auth_token_file, v.auth_token));
     }
   };
 
@@ -182,7 +189,7 @@ Result<AwsCredentialsProvider> MakeAwsCredentialsProvider(const Spec& spec) {
                                   aws_error_debug_str(err)));
   }
 
-  return MakeCache(std::move(credentials_provider));
+  return credentials_provider;
 }
 
 }  // namespace internal_kvstore_s3

tensorstore/kvstore/s3/aws_credentials_spec_test.cc

@@ -14,12 +14,20 @@
 
 #include "tensorstore/kvstore/s3/aws_credentials_spec.h"
 
+#include <iterator>
+#include <optional>
+#include <string>
 #include <type_traits>
+#include <utility>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/cleanup/cleanup.h"
 #include "absl/status/status.h"
 #include <nlohmann/json.hpp>
+#include "tensorstore/internal/aws/aws_credentials.h"
+#include "tensorstore/internal/env.h"
 #include "tensorstore/internal/testing/json_gtest.h"
 #include "tensorstore/json_serialization_options_base.h"
 #include "tensorstore/util/status_testutil.h"
@@ -32,6 +40,11 @@ namespace jb = ::tensorstore::internal_json_binding;
 using ::tensorstore::IsOk;
 using ::tensorstore::MatchesJson;
 using ::tensorstore::StatusIs;
+using ::tensorstore::internal::GetEnv;
+using ::tensorstore::internal::SetEnv;
+using ::tensorstore::internal::UnsetEnv;
+using ::tensorstore::internal_aws::GetAwsCredentials;
+using ::tensorstore::internal_kvstore_s3::MakeAwsCredentialsProvider;
 using ::testing::HasSubstr;
 
 struct Params {
@@ -142,4 +155,60 @@ TEST(SpecTest, Error) {
                HasSubstr("Failed to parse AWS credentials spec")));
 }
 
+TEST(SpecTest, DefaultDoesNotCacheAnonymousFallback) {
+  constexpr const char* kAwsEnvVars[] = {
+      "AWS_EC2_METADATA_DISABLED",
+      "AWS_WEB_IDENTITY_TOKEN_FILE",
+      "AWS_ROLE_ARN",
+      "AWS_ACCESS_KEY_ID",
+      "AWS_SECRET_ACCESS_KEY",
+      "AWS_SESSION_TOKEN",
+      "AWS_CONFIG_FILE",
+      "AWS_SHARED_CREDENTIALS_FILE",
+      "AWS_PROFILE",
+      "AWS_CONTAINER_CREDENTIALS_FULL_URI",
+      "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
+  };
+  std::vector<std::pair<const char*, std::optional<std::string>>> saved_env;
+  saved_env.reserve(std::size(kAwsEnvVars));
+  for (const char* var : kAwsEnvVars) {
+    saved_env.emplace_back(var, GetEnv(var));
+    UnsetEnv(var);
+  }
+  absl::Cleanup restore_env = [&saved_env] {
+    for (const auto& [var, val] : saved_env) {
+      if (val.has_value()) {
+        SetEnv(var, val->c_str());
+      } else {
+        UnsetEnv(var);
+      }
+    }
+  };
+  // Keep the test hermetic: no ~/.aws files and no EC2 IMDS requests.
+  SetEnv("AWS_EC2_METADATA_DISABLED", "true");
+  SetEnv("AWS_CONFIG_FILE", "/nonexistent/aws_config");
+  SetEnv("AWS_SHARED_CREDENTIALS_FILE", "/nonexistent/aws_credentials");
+
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+      auto provider, MakeAwsCredentialsProvider(Spec{Spec::Default{}}));
+
+  // No credentials configured: falls back to anonymous.
+  {
+    auto creds = GetAwsCredentials(provider.get()).result();
+    ASSERT_THAT(creds, IsOk());
+    EXPECT_TRUE(creds->IsAnonymous());
+  }
+
+  // Credentials appear; the same provider must pick them up.
+  SetEnv("AWS_ACCESS_KEY_ID", "AKIDEXAMPLE");
+  SetEnv("AWS_SECRET_ACCESS_KEY", "secretexample");
+  {
+    auto creds = GetAwsCredentials(provider.get()).result();
+    ASSERT_THAT(creds, IsOk());
+    EXPECT_FALSE(creds->IsAnonymous());
+    EXPECT_EQ(creds->GetAccessKeyId(), "AKIDEXAMPLE");
+    EXPECT_EQ(creds->GetSecretAccessKey(), "secretexample");
+  }
+}
+
 }  // namespace

tensorstore/kvstore/s3/s3_key_value_store.cc

@@ -118,7 +118,6 @@ using ::tensorstore::internal_kvstore_s3::GetNodeText;
 using ::tensorstore::internal_kvstore_s3::IsValidBucketName;
 using ::tensorstore::internal_kvstore_s3::IsValidObjectName;
 using ::tensorstore::internal_kvstore_s3::IsValidStorageGeneration;
-using ::tensorstore::internal_kvstore_s3::MakeAwsCredentialsProvider;
 using ::tensorstore::internal_kvstore_s3::S3ConcurrencyResource;
 using ::tensorstore::internal_kvstore_s3::S3EndpointRegion;
 using ::tensorstore::internal_kvstore_s3::S3RateLimiterResource;
@@ -1413,11 +1412,9 @@ Future<kvstore::DriverPtr> S3KeyValueStoreSpec::DoOpen() const {
   (void)internal_aws::GetAwsAllocator();  // Force AWS library initialization.
   // TODO: The transport should support the AWS_CA_BUNDLE environment
   // variable.
-  TENSORSTORE_ASSIGN_OR_RETURN(
-      auto provider, MakeAwsCredentialsProvider(*data_.aws_credentials));
-
   auto driver = internal::MakeIntrusivePtr<S3KeyValueStore>(
-      internal_http::GetDefaultHttpTransport(), data_, std::move(provider));
+      internal_http::GetDefaultHttpTransport(), data_,
+      data_.aws_credentials->credentials_provider);
 
   // NOTE: Remove temporary logging use of experimental feature.
   if (data_.rate_limiter.has_value()) {