diff --git a/tensorstore/internal/aws/credentials/common.cc b/tensorstore/internal/aws/credentials/common.cc index b6c92e5c3..af31d44e4 100644 --- a/tensorstore/internal/aws/credentials/common.cc +++ b/tensorstore/internal/aws/credentials/common.cc @@ -36,6 +36,7 @@ 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); } @@ -43,6 +44,7 @@ inline AwsCredentialsProvider AsProvider(aws_credentials_provider* p) { } // 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); diff --git a/tensorstore/internal/aws/credentials/common.h b/tensorstore/internal/aws/credentials/common.h index 4245efbc3..31801b388 100644 --- a/tensorstore/internal/aws/credentials/common.h +++ b/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 diff --git a/tensorstore/kvstore/s3/BUILD b/tensorstore/kvstore/s3/BUILD index 571730b22..ca254e055 100644 --- a/tensorstore/kvstore/s3/BUILD +++ b/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", diff --git a/tensorstore/kvstore/s3/aws_credentials_resource.h b/tensorstore/kvstore/s3/aws_credentials_resource.h index b463c0a05..b65f536f7 100644 --- a/tensorstore/kvstore/s3/aws_credentials_resource.h +++ b/tensorstore/kvstore/s3/aws_credentials_resource.h @@ -18,9 +18,11 @@ #include #include +#include #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 { 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 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; } }; diff --git a/tensorstore/kvstore/s3/aws_credentials_resource_test.cc b/tensorstore/kvstore/s3/aws_credentials_resource_test.cc index fbcda9858..db6696e45 100644 --- a/tensorstore/kvstore/s3/aws_credentials_resource_test.cc +++ b/tensorstore/kvstore/s3/aws_credentials_resource_test.cc @@ -21,6 +21,7 @@ #include "absl/status/status.h" #include #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::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::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::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) { diff --git a/tensorstore/kvstore/s3/aws_credentials_spec.cc b/tensorstore/kvstore/s3/aws_credentials_spec.cc index ff610aa11..e4137fdb9 100644 --- a/tensorstore/kvstore/s3/aws_credentials_spec.cc +++ b/tensorstore/kvstore/s3/aws_credentials_spec.cc @@ -155,17 +155,24 @@ Result 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 MakeAwsCredentialsProvider(const Spec& spec) { aws_error_debug_str(err))); } - return MakeCache(std::move(credentials_provider)); + return credentials_provider; } } // namespace internal_kvstore_s3 diff --git a/tensorstore/kvstore/s3/aws_credentials_spec_test.cc b/tensorstore/kvstore/s3/aws_credentials_spec_test.cc index b54a9a166..cad1bf493 100644 --- a/tensorstore/kvstore/s3/aws_credentials_spec_test.cc +++ b/tensorstore/kvstore/s3/aws_credentials_spec_test.cc @@ -14,12 +14,20 @@ #include "tensorstore/kvstore/s3/aws_credentials_spec.h" +#include +#include +#include #include +#include +#include #include #include +#include "absl/cleanup/cleanup.h" #include "absl/status/status.h" #include +#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>> 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 diff --git a/tensorstore/kvstore/s3/s3_key_value_store.cc b/tensorstore/kvstore/s3/s3_key_value_store.cc index e26be8cdf..dffdb26bf 100644 --- a/tensorstore/kvstore/s3/s3_key_value_store.cc +++ b/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 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( - 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()) {