diff --git a/tensorstore/kvstore/s3/BUILD b/tensorstore/kvstore/s3/BUILD index a957d27af..f0cedd24c 100644 --- a/tensorstore/kvstore/s3/BUILD +++ b/tensorstore/kvstore/s3/BUILD @@ -21,7 +21,6 @@ tensorstore_cc_library( ], hdrs = ["s3_metadata.h"], deps = [ - ":aws_credential_provider", ":s3_endpoint", ":s3_request_builder", ":s3_resource", @@ -47,6 +46,8 @@ tensorstore_cc_library( "//tensorstore/kvstore:key_range", "//tensorstore/kvstore/gcs:validate", "//tensorstore/kvstore/gcs_http:rate_limiter", + "//tensorstore/kvstore/s3/credentials:aws_credentials", + "//tensorstore/kvstore/s3/credentials:default_credential_provider", "//tensorstore/serialization", "//tensorstore/util:executor", "//tensorstore/util:future", @@ -129,9 +130,9 @@ tensorstore_cc_test( "//tensorstore/kvstore:read_result_testutil", "//tensorstore/kvstore:test_util", "//tensorstore/util:future", + "//tensorstore/util:result", "//tensorstore/util:status_testutil", "//tensorstore/util:str_cat", - "@com_github_nlohmann_json//:nlohmann_json", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:absl_log", "@com_google_absl//absl/status", @@ -150,13 +151,13 @@ tensorstore_cc_library( "s3_request_builder.h", ], deps = [ - ":aws_credential_provider", ":s3_uri_utils", "//tensorstore/internal:uri_utils", "//tensorstore/internal/digest:sha256", "//tensorstore/internal/http", "//tensorstore/internal/log:verbose_flag", "//tensorstore/kvstore:byte_range", + "//tensorstore/kvstore/s3/credentials:aws_credentials", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", @@ -172,9 +173,9 @@ tensorstore_cc_test( size = "small", srcs = ["s3_request_builder_test.cc"], deps = [ - ":aws_credential_provider", ":s3_request_builder", "//tensorstore/internal/http", + "//tensorstore/kvstore/s3/credentials:aws_credentials", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/time", @@ -208,73 +209,9 @@ tensorstore_cc_test( ], ) -tensorstore_cc_library( - name = "aws_credential_provider", - srcs = [ - "aws_credential_provider.cc", - "aws_metadata_credential_provider.cc", - ], - hdrs = [ - "aws_credential_provider.h", - "aws_metadata_credential_provider.h", - ], - deps = [ - "//tensorstore/internal:env", - "//tensorstore/internal:no_destructor", - "//tensorstore/internal:path", - "//tensorstore/internal/http", - "//tensorstore/internal/json", - "//tensorstore/internal/json_binding", - "//tensorstore/internal/json_binding:absl_time", - "//tensorstore/internal/json_binding:bindable", - "//tensorstore/util:result", - "//tensorstore/util:status", - "//tensorstore/util:str_cat", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/log:absl_log", - "@com_google_absl//absl/status", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:cord", - "@com_google_absl//absl/synchronization", - "@com_google_absl//absl/time", - ], -) - -tensorstore_cc_test( - name = "aws_credential_provider_test", - srcs = ["aws_credential_provider_test.cc"], - deps = [ - ":aws_credential_provider", - "//tensorstore/internal:env", - "//tensorstore/internal:path", - "//tensorstore/internal:test_util", - "//tensorstore/internal/http:curl_transport", - "//tensorstore/util:result", - "//tensorstore/util:status_testutil", - "@com_google_googletest//:gtest_main", - ], -) - -tensorstore_cc_test( - name = "aws_metadata_credential_provider_test", - srcs = ["aws_metadata_credential_provider_test.cc"], - deps = [ - ":aws_credential_provider", - "//tensorstore/internal/http", - "//tensorstore/util:result", - "//tensorstore/util:status_testutil", - "//tensorstore/util:str_cat", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log:absl_log", - "@com_google_absl//absl/status", - "@com_google_absl//absl/strings:cord", - "@com_google_absl//absl/time", - "@com_google_googletest//:gtest_main", - ], -) - tensorstore_cc_test( name = "localstack_test", + size = "small", srcs = ["localstack_test.cc"], args = [ "--localstack_binary=$(location :moto_server)", @@ -287,7 +224,6 @@ tensorstore_cc_test( "skip-cmake", ], deps = [ - ":aws_credential_provider", ":s3", ":s3_request_builder", "//tensorstore:context", @@ -300,10 +236,12 @@ tensorstore_cc_test( "//tensorstore/internal/os:subprocess", "//tensorstore/kvstore", "//tensorstore/kvstore:test_util", + "//tensorstore/kvstore/s3/credentials:aws_credentials", "//tensorstore/util:future", "//tensorstore/util:result", "//tensorstore/util:status_testutil", "@com_github_nlohmann_json//:nlohmann_json", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", @@ -336,6 +274,7 @@ tensorstore_cc_library( tensorstore_cc_test( name = "s3_endpoint_test", + size = "small", srcs = ["s3_endpoint_test.cc"], deps = [ ":s3_endpoint", diff --git a/tensorstore/kvstore/s3/aws_credential_provider.cc b/tensorstore/kvstore/s3/aws_credential_provider.cc deleted file mode 100644 index c334ca3dc..000000000 --- a/tensorstore/kvstore/s3/aws_credential_provider.cc +++ /dev/null @@ -1,257 +0,0 @@ -// Copyright 2023 The TensorStore Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "tensorstore/kvstore/s3/aws_credential_provider.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/log/absl_log.h" -#include "absl/status/status.h" -#include "absl/strings/ascii.h" -#include "absl/strings/str_cat.h" -#include "absl/synchronization/mutex.h" -#include "tensorstore/internal/env.h" -#include "tensorstore/internal/http/http_transport.h" -#include "tensorstore/internal/no_destructor.h" -#include "tensorstore/internal/path.h" -#include "tensorstore/kvstore/s3/aws_metadata_credential_provider.h" -#include "tensorstore/util/result.h" - -using ::tensorstore::Result; -using ::tensorstore::internal::GetEnv; -using ::tensorstore::internal::JoinPath; - -namespace tensorstore { -namespace internal_kvstore_s3 { -namespace { - -// For reference, see the latest AWS environment variables used by the cli: -// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html - -// AWS user identifier -constexpr char kEnvAwsAccessKeyId[] = "AWS_ACCESS_KEY_ID"; -constexpr char kCfgAwsAccessKeyId[] = "aws_access_key_id"; - -// AWS user password -constexpr char kEnvAwsSecretAccessKey[] = "AWS_SECRET_ACCESS_KEY"; -constexpr char kCfgAwsSecretAccessKeyId[] = "aws_secret_access_key"; - -// AWS session token -constexpr char kEnvAwsSessionToken[] = "AWS_SESSION_TOKEN"; -constexpr char kCfgAwsSessionToken[] = "aws_session_token"; - -// AWS Profile environment variables -constexpr char kEnvAwsProfile[] = "AWS_PROFILE"; -constexpr char kDefaultProfile[] = "default"; - -// Credentials file environment variable -constexpr char kEnvAwsCredentialsFile[] = "AWS_SHARED_CREDENTIALS_FILE"; - -// Default path to the AWS credentials file, relative to the home folder -constexpr char kDefaultAwsCredentialsFilePath[] = ".aws/credentials"; - -/// Returns whether the given path points to a readable file. -bool IsFile(const std::string& filename) { - std::ifstream fstream(filename.c_str()); - return fstream.good(); -} - -Result GetAwsCredentialsFileName() { - std::string result; - - auto credentials_file = GetEnv(kEnvAwsCredentialsFile); - if (!credentials_file) { - auto home_dir = GetEnv("HOME"); - if (!home_dir) { - return absl::NotFoundError("Could not read $HOME"); - } - result = JoinPath(*home_dir, kDefaultAwsCredentialsFilePath); - } else { - result = *credentials_file; - } - if (!IsFile(result)) { - return absl::NotFoundError( - absl::StrCat("Could not find the credentials file at " - "location [", - result, "]")); - } - return result; -} - -Result> GetDefaultAwsCredentialProvider( - std::string_view profile, - std::shared_ptr transport) { - // 1. Obtain credentials from environment variables - if (auto access_key = GetEnv(kEnvAwsAccessKeyId); access_key) { - ABSL_LOG_FIRST_N(INFO, 1) - << "Using Environment Variable " << kEnvAwsAccessKeyId; - AwsCredentials credentials; - credentials.access_key = *access_key; - auto secret_key = GetEnv(kEnvAwsSecretAccessKey); - - if (secret_key.has_value()) { - credentials.secret_key = *secret_key; - } - - auto session_token = GetEnv(kEnvAwsSessionToken); - - if (session_token.has_value()) { - credentials.session_token = *session_token; - } - - return std::make_unique( - std::move(credentials)); - } - - // 2. Obtain credentials from AWS_SHARED_CREDENTIALS_FILE or - // ~/.aws/credentials - if (auto credentials_file = GetAwsCredentialsFileName(); - credentials_file.ok()) { - std::string env_profile; // value must not outlive view - if (profile.empty()) { - env_profile = GetEnv(kEnvAwsProfile).value_or(kDefaultProfile); - profile = std::string_view(env_profile); - } - ABSL_LOG(INFO) << "Using File AwsCredentialProvider with profile " - << profile; - return std::make_unique( - std::move(credentials_file).value(), std::string(profile)); - } - - // 3. Obtain credentials from EC2 Metadata server - if (IsEC2MetadataServiceAvailable(*transport)) { - ABSL_LOG(INFO) << "Using EC2 Metadata Service AwsCredentialProvider"; - return std::make_unique(transport); - } - - return absl::NotFoundError( - "No credentials provided in environment variables, " - "credentials file not found and not running on AWS."); -} - -struct AwsCredentialProviderRegistry { - std::vector> providers; - absl::Mutex mutex; -}; - -AwsCredentialProviderRegistry& GetAwsProviderRegistry() { - static internal::NoDestructor registry; - return *registry; -} - -} // namespace - -/// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-format -Result FileCredentialProvider::GetCredentials() { - absl::ReaderMutexLock lock(&mutex_); - std::ifstream ifs(filename_); - - if (!ifs) { - return absl::NotFoundError( - absl::StrCat("Could not open the credentials file [", filename_, "]")); - } - - AwsCredentials credentials; - std::string section_name; - std::string line; - bool profile_found = false; - - while (std::getline(ifs, line)) { - auto sline = absl::StripAsciiWhitespace(line); - // Ignore empty and commented out lines - if (sline.empty() || sline[0] == '#') continue; - - // A configuration section name has been encountered - if (sline[0] == '[' && sline[sline.size() - 1] == ']') { - section_name = - absl::StripAsciiWhitespace(sline.substr(1, sline.size() - 2)); - continue; - } - - // Look for key=value pairs if we're in the appropriate profile - if (section_name == profile_) { - profile_found = true; - if (auto pos = sline.find('='); pos != std::string::npos) { - auto key = absl::StripAsciiWhitespace(sline.substr(0, pos)); - auto value = absl::StripAsciiWhitespace(sline.substr(pos + 1)); - - if (key == kCfgAwsAccessKeyId) { - credentials.access_key = value; - } else if (key == kCfgAwsSecretAccessKeyId) { - credentials.secret_key = value; - } else if (key == kCfgAwsSessionToken) { - credentials.session_token = value; - } - } - } - } - - if (!profile_found) { - return absl::NotFoundError(absl::StrCat("Profile [", profile_, - "] not found " - "in credentials file [", - filename_, "]")); - } - - return credentials; -} - -void RegisterAwsCredentialProviderProvider(AwsCredentialProviderFn provider, - int priority) { - auto& registry = GetAwsProviderRegistry(); - absl::WriterMutexLock lock(®istry.mutex); - registry.providers.emplace_back(priority, std::move(provider)); - std::sort(registry.providers.begin(), registry.providers.end(), - [](const auto& a, const auto& b) { return a.first < b.first; }); -} - -/// @brief Obtain a credential provider from a series of registered and default -/// providers -/// -/// Providers are returned in the following order: -/// 1. Any registered providers that supply valid credentials -/// 2. Environment variable provider if valid credential can be obtained from -/// AWS_* environment variables -/// 3. File provider containing credentials from an ~/.aws/credentials file -/// 4. EC2 Metadata server -/// -/// @param profile The profile to use when retrieving credentials from a -/// credentials file. -/// @param transport Optionally specify the http transport used to retreive S3 -/// credentials -/// from the EC2 metadata server. -/// @return Provider that supplies S3 Credentials -Result> GetAwsCredentialProvider( - std::string_view profile, - std::shared_ptr transport) { - auto& registry = GetAwsProviderRegistry(); - absl::WriterMutexLock lock(®istry.mutex); - for (const auto& provider : registry.providers) { - auto credentials = provider.second(); - if (credentials.ok()) return credentials; - } - - return internal_kvstore_s3::GetDefaultAwsCredentialProvider(profile, - transport); -} - -} // namespace internal_kvstore_s3 -} // namespace tensorstore diff --git a/tensorstore/kvstore/s3/aws_credential_provider.h b/tensorstore/kvstore/s3/aws_credential_provider.h deleted file mode 100644 index ee946ec26..000000000 --- a/tensorstore/kvstore/s3/aws_credential_provider.h +++ /dev/null @@ -1,104 +0,0 @@ -// Copyright 2023 The TensorStore Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#ifndef TENSORSTORE_KVSTORE_S3_AWS_CREDENTIAL_PROVIDER_H -#define TENSORSTORE_KVSTORE_S3_AWS_CREDENTIAL_PROVIDER_H - -#include -#include -#include -#include -#include - -#include "absl/synchronization/mutex.h" -#include "tensorstore/internal/http/http_transport.h" -#include "tensorstore/util/result.h" - -namespace tensorstore { -namespace internal_kvstore_s3 { - -/// Holds S3 credentials -/// -/// Contains the access key, secret key and session token. -/// An empty access key implies anonymous access, -/// while the presence of a session token implies the use of -/// short-lived STS credentials -/// https://docs.aws.amazon.com/STS/latest/APIReference/welcome.html -struct AwsCredentials { - /// AWS_ACCESS_KEY_ID - std::string access_key; - /// AWS_SECRET_KEY_ID - std::string secret_key; - /// AWS_SESSION_TOKEN - std::string session_token; - - bool IsAnonymous() const { return access_key.empty(); } -}; - -/// Base class for S3 Credential Providers -/// -/// Implementers should override GetCredentials -class AwsCredentialProvider { - public: - virtual ~AwsCredentialProvider() = default; - virtual Result GetCredentials() = 0; -}; - -/// Provides credentials from the following environment variables: -/// AWS_ACCESS_KEY_ID, AWS_SECRET_KEY_ID, AWS_SESSION_TOKEN -class EnvironmentCredentialProvider : public AwsCredentialProvider { - private: - AwsCredentials credentials_; - - public: - EnvironmentCredentialProvider(const AwsCredentials& credentials) - : credentials_(credentials) {} - - Result GetCredentials() override { return credentials_; } -}; - -/// Obtains S3 credentials from a profile in a file, usually -/// `~/.aws/credentials` or a file specified in AWS_SHARED_CREDENTIALS_FILE. A -/// desired profile may be specified in the constructor: This value should be -/// derived from the s3 json spec. -/// However, if profile is passed as an empty string, the profile is obtained -/// from AWS_DEFAULT_PROFILE, AWS_PROFILE before finally defaulting to -/// "default". -class FileCredentialProvider : public AwsCredentialProvider { - private: - absl::Mutex mutex_; - std::string filename_; - std::string profile_; - - public: - FileCredentialProvider(std::string filename, std::string profile) - : filename_(std::move(filename)), profile_(std::move(profile)) {} - - Result GetCredentials() override; -}; - -using AwsCredentialProviderFn = - std::function>()>; - -void RegisterAwsCredentialProviderProvider(AwsCredentialProviderFn provider, - int priority); - -Result> GetAwsCredentialProvider( - std::string_view profile, - std::shared_ptr transport); - -} // namespace internal_kvstore_s3 -} // namespace tensorstore - -#endif // TENSORSTORE_KVSTORE_S3_AWS_CREDENTIAL_PROVIDER_H diff --git a/tensorstore/kvstore/s3/aws_credential_provider_test.cc b/tensorstore/kvstore/s3/aws_credential_provider_test.cc deleted file mode 100644 index 04553aefb..000000000 --- a/tensorstore/kvstore/s3/aws_credential_provider_test.cc +++ /dev/null @@ -1,158 +0,0 @@ -// Copyright 2023 The TensorStore Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "tensorstore/kvstore/s3/aws_credential_provider.h" - -#include -#include - -#include -#include -#include "tensorstore/internal/env.h" -#include "tensorstore/internal/http/curl_transport.h" -#include "tensorstore/internal/path.h" -#include "tensorstore/internal/test_util.h" -#include "tensorstore/util/result.h" -#include "tensorstore/util/status_testutil.h" - -namespace { - -using ::tensorstore::internal::GetEnv; -using ::tensorstore::internal::JoinPath; -using ::tensorstore::internal::SetEnv; -using ::tensorstore::internal::UnsetEnv; -using ::tensorstore::internal_http::GetDefaultHttpTransport; -using ::tensorstore::internal_kvstore_s3::GetAwsCredentialProvider; - -class TestData : public tensorstore::internal::ScopedTemporaryDirectory { - public: - std::string WriteCredentialsFile() { - auto p = JoinPath(path(), "aws_config"); - std::ofstream ofs(p); - ofs << "discarded_value = 500\n" - "\n" - "[default]\n" - "aws_access_key_id =AKIAIOSFODNN7EXAMPLE\n" - "aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\n" - "aws_session_token= abcdef1234567890 \n" - "\n" - "[alice]\n" - "aws_access_key_id = AKIAIOSFODNN6EXAMPLE\n" - "aws_secret_access_key = " - "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY\n" - "\n"; - ofs.close(); - return p; - } -}; - -class AwsCredentialProviderTest : public ::testing::Test { - protected: - void SetUp() override { - // Make sure that env vars are not set. - for (const char* var : - {"AWS_SHARED_CREDENTIALS_FILE", "AWS_ACCESS_KEY_ID", - "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_PROFILE"}) { - UnsetEnv(var); - } - } -}; - -TEST_F(AwsCredentialProviderTest, ProviderNoCredentials) { - ASSERT_FALSE(GetAwsCredentialProvider("", GetDefaultHttpTransport()).ok()); - SetEnv("AWS_ACCESS_KEY_ID", "foo"); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, GetAwsCredentialProvider("", GetDefaultHttpTransport())); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, - provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "foo"); - ASSERT_TRUE(credentials.secret_key.empty()); - ASSERT_TRUE(credentials.session_token.empty()); -} - -TEST_F(AwsCredentialProviderTest, ProviderAwsCredentialsFromEnv) { - SetEnv("AWS_ACCESS_KEY_ID", "foo"); - SetEnv("AWS_SECRET_ACCESS_KEY", "bar"); - SetEnv("AWS_SESSION_TOKEN", "qux"); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, GetAwsCredentialProvider("", GetDefaultHttpTransport())); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, - provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "foo"); - ASSERT_EQ(credentials.secret_key, "bar"); - ASSERT_EQ(credentials.session_token, "qux"); -} - -TEST_F(AwsCredentialProviderTest, ProviderAwsCredentialsFromFileDefault) { - TestData test_data; - std::string credentials_filename = test_data.WriteCredentialsFile(); - - SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, GetAwsCredentialProvider("", GetDefaultHttpTransport())); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, - provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN7EXAMPLE"); - ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"); - ASSERT_EQ(credentials.session_token, "abcdef1234567890"); -} - -TEST_F(AwsCredentialProviderTest, - ProviderAwsCredentialsFromFileProfileOverride) { - TestData test_data; - std::string credentials_filename = test_data.WriteCredentialsFile(); - - SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, - GetAwsCredentialProvider("alice", GetDefaultHttpTransport())); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, - provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); - ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); - ASSERT_EQ(credentials.session_token, ""); -} - -TEST_F(AwsCredentialProviderTest, ProviderAwsCredentialsFromFileProfileEnv) { - TestData test_data; - std::string credentials_filename = test_data.WriteCredentialsFile(); - - SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); - SetEnv("AWS_PROFILE", "alice"); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, GetAwsCredentialProvider("", GetDefaultHttpTransport())); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, - provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); - ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); - ASSERT_EQ(credentials.session_token, ""); -} - -TEST_F(AwsCredentialProviderTest, - ProviderAwsCredentialsFromFileInvalidProfileEnv) { - TestData test_data; - std::string credentials_filename = test_data.WriteCredentialsFile(); - - SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); - SetEnv("AWS_PROFILE", "bob"); - TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto provider, GetAwsCredentialProvider("", GetDefaultHttpTransport())); - auto result = provider->GetCredentials(); - ASSERT_FALSE(result.ok()); - EXPECT_THAT( - result.status().message(), - ::testing::HasSubstr("Profile [bob] not found in credentials file")); -} - -} // namespace diff --git a/tensorstore/kvstore/s3/aws_metadata_credential_provider_test.cc b/tensorstore/kvstore/s3/aws_metadata_credential_provider_test.cc deleted file mode 100644 index 24c88c6f8..000000000 --- a/tensorstore/kvstore/s3/aws_metadata_credential_provider_test.cc +++ /dev/null @@ -1,154 +0,0 @@ -// Copyright 2023 The TensorStore Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "tensorstore/kvstore/s3/aws_metadata_credential_provider.h" - -#include -#include - -#include -#include -#include "absl/container/flat_hash_map.h" -#include "absl/log/absl_log.h" -#include "absl/status/status.h" -#include "absl/strings/cord.h" -#include "absl/time/time.h" -#include "tensorstore/internal/http/http_request.h" -#include "tensorstore/internal/http/http_response.h" -#include "tensorstore/internal/http/http_transport.h" -#include "tensorstore/util/result.h" -#include "tensorstore/util/status_testutil.h" -#include "tensorstore/util/str_cat.h" - -namespace { - -using ::tensorstore::Future; -using ::tensorstore::MatchesStatus; -using ::tensorstore::internal_http::HttpRequest; -using ::tensorstore::internal_http::HttpResponse; -using ::tensorstore::internal_http::HttpTransport; -using ::tensorstore::internal_kvstore_s3::EC2MetadataCredentialProvider; - -class EC2MetadataMockTransport : public HttpTransport { - public: - EC2MetadataMockTransport( - const absl::flat_hash_map& url_to_response) - : url_to_response_(url_to_response) {} - - Future IssueRequest(const HttpRequest& request, - absl::Cord payload, - absl::Duration request_timeout, - absl::Duration connect_timeout) override { - ABSL_LOG(INFO) << request; - auto it = url_to_response_.find( - tensorstore::StrCat(request.method, " ", request.url)); - if (it != url_to_response_.end()) { - return it->second; - } - return HttpResponse{404, absl::Cord(), {}}; - } - - const absl::flat_hash_map& url_to_response_; -}; - -TEST(EC2MetadataCredentialProviderTest, CredentialRetrievalFlow) { - auto url_to_response = absl::flat_hash_map{ - {"POST http://169.254.169.254/latest/api/token", - HttpResponse{200, absl::Cord{"1234567890"}}}, - {"GET http://169.254.169.254/latest/meta-data/iam/", - HttpResponse{200, - absl::Cord{"info"}, - {{"x-aws-ec2-metadata-token", "1234567890"}}}}, - {"GET http://169.254.169.254/latest/meta-data/iam/security-credentials/", - HttpResponse{200, - absl::Cord{"mock-iam-role\nmock-iam-role2"}, - {{"x-aws-ec2-metadata-token", "1234567890"}}}}, - {"GET " - "http://169.254.169.254/latest/meta-data/iam/security-credentials/" - "mock-iam-role", - HttpResponse{200, - absl::Cord(R"({ - "Code": "Success", - "LastUpdated": "2023-09-21T12:42:12Z", - "Type": "AWS-HMAC", - "AccessKeyId": "ASIA1234567890", - "SecretAccessKey": "1234567890abcdef", - "Token": "abcdef123456790", - "Expiration": "2023-09-21T12:42:12Z" - })"), - {{"x-aws-ec2-metadata-token", "1234567890"}}}}}; - - auto mock_transport = - std::make_shared(url_to_response); - auto provider = - std::make_shared(mock_transport); - TENSORSTORE_CHECK_OK_AND_ASSIGN(auto credentials, provider->GetCredentials()); - ASSERT_EQ(credentials.access_key, "ASIA1234567890"); - ASSERT_EQ(credentials.secret_key, "1234567890abcdef"); - ASSERT_EQ(credentials.session_token, "abcdef123456790"); -} - -TEST(EC2MetadataCredentialProviderTest, NoIamRolesInSecurityCredentials) { - auto url_to_response = absl::flat_hash_map{ - {"POST http://169.254.169.254/latest/api/token", - HttpResponse{200, absl::Cord{"1234567890"}}}, - {"GET http://169.254.169.254/latest/meta-data/iam/security-credentials/", - HttpResponse{ - 200, absl::Cord{""}, {{"x-aws-ec2-metadata-token", "1234567890"}}}}, - }; - - auto mock_transport = - std::make_shared(url_to_response); - auto provider = - std::make_shared(mock_transport); - ASSERT_FALSE(provider->GetCredentials()); - EXPECT_THAT(provider->GetCredentials().status().ToString(), - ::testing::HasSubstr("Empty EC2 Role list")); -} - -TEST(EC2MetadataCredentialProviderTest, UnsuccessfulJsonResponse) { - // Test that "Code" != "Success" parsing succeeds - auto url_to_response = absl::flat_hash_map{ - {"POST http://169.254.169.254/latest/api/token", - HttpResponse{200, absl::Cord{"1234567890"}}}, - {"GET http://169.254.169.254/latest/meta-data/iam/", - HttpResponse{200, - absl::Cord{"info"}, - {{"x-aws-ec2-metadata-token", "1234567890"}}}}, - {"GET http://169.254.169.254/latest/meta-data/iam/security-credentials/", - HttpResponse{200, - absl::Cord{"mock-iam-role"}, - {{"x-aws-ec2-metadata-token", "1234567890"}}}}, - {"GET " - "http://169.254.169.254/latest/meta-data/iam/security-credentials/" - "mock-iam-role", - HttpResponse{200, - absl::Cord(R"({"Code": "EntirelyUnsuccessful"})"), - {{"x-aws-ec2-metadata-token", "1234567890"}}}}}; - - auto mock_transport = - std::make_shared(url_to_response); - auto provider = - std::make_shared(mock_transport); - auto credentials = provider->GetCredentials(); - - EXPECT_THAT(credentials.status(), MatchesStatus(absl::StatusCode::kNotFound)); - EXPECT_THAT(credentials.status().ToString(), - ::testing::AllOf( - ::testing::HasSubstr("EC2Metadata request"), - ::testing::HasSubstr( - "failed with {\"Code\": \"EntirelyUnsuccessful\"}"))); -} - -} // namespace diff --git a/tensorstore/kvstore/s3/credentials/BUILD b/tensorstore/kvstore/s3/credentials/BUILD new file mode 100644 index 000000000..d8eea218e --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/BUILD @@ -0,0 +1,170 @@ +load("//bazel:tensorstore.bzl", "tensorstore_cc_library", "tensorstore_cc_test") + +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) + +tensorstore_cc_library( + name = "aws_credentials", + hdrs = ["aws_credentials.h"], + deps = [ + "//tensorstore/util:result", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_library( + name = "environment_credential_provider", + srcs = ["environment_credential_provider.cc"], + hdrs = ["environment_credential_provider.h"], + deps = [ + ":aws_credentials", + "//tensorstore/internal:env", + "//tensorstore/util:result", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_library( + name = "file_credential_provider", + srcs = ["file_credential_provider.cc"], + hdrs = ["file_credential_provider.h"], + deps = [ + ":aws_credentials", + "//tensorstore/internal:env", + "//tensorstore/internal:path", + "//tensorstore/util:result", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_library( + name = "ec2_credential_provider", + srcs = ["ec2_credential_provider.cc"], + hdrs = ["ec2_credential_provider.h"], + deps = [ + ":aws_credentials", + "//tensorstore/internal:env", + "//tensorstore/internal/http", + "//tensorstore/internal/json", + "//tensorstore/internal/json_binding", + "//tensorstore/internal/json_binding:absl_time", + "//tensorstore/internal/json_binding:bindable", + "//tensorstore/util:result", + "//tensorstore/util:status", + "//tensorstore/util:str_cat", + "@com_github_nlohmann_json//:nlohmann_json", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:cord", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_library( + name = "default_credential_provider", + srcs = ["default_credential_provider.cc"], + hdrs = ["default_credential_provider.h"], + deps = [ + ":aws_credentials", + ":ec2_credential_provider", + ":environment_credential_provider", + ":file_credential_provider", + "//tensorstore/internal:no_destructor", + "//tensorstore/internal/http", + "//tensorstore/internal/http:curl_transport", + "//tensorstore/util:result", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/functional:function_ref", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_library( + name = "test_utils", + testonly = 1, + srcs = ["test_utils.cc"], + hdrs = ["test_utils.h"], + deps = [ + "//tensorstore/internal/http", + "//tensorstore/util:future", + "//tensorstore/util:str_cat", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/strings:cord", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/time", + ], +) + +tensorstore_cc_test( + name = "default_credential_provider_test", + size = "small", + srcs = ["default_credential_provider_test.cc"], + deps = [ + ":default_credential_provider", + ":test_utils", + "//tensorstore/internal:env", + "//tensorstore/internal:path", + "//tensorstore/internal:test_util", + "//tensorstore/internal/http", + "//tensorstore/util:status_testutil", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings:cord", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) + +tensorstore_cc_test( + name = "environmental_credential_provider_test", + size = "small", + srcs = ["environment_credential_provider_test.cc"], + deps = [ + ":environment_credential_provider", + "//tensorstore/internal:env", + "//tensorstore/util:status_testutil", + "@com_google_googletest//:gtest_main", + ], +) + +tensorstore_cc_test( + name = "file_credential_provider_test", + size = "small", + srcs = ["file_credential_provider_test.cc"], + deps = [ + ":file_credential_provider", + "//tensorstore/internal:env", + "//tensorstore/internal:path", + "//tensorstore/internal:test_util", + "//tensorstore/util:status_testutil", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) + +tensorstore_cc_test( + name = "ec2_credential_provider_test", + size = "small", + srcs = ["ec2_credential_provider_test.cc"], + deps = [ + ":ec2_credential_provider", + ":test_utils", + "//tensorstore/internal:env", + "//tensorstore/internal/http", + "//tensorstore/util:result", + "//tensorstore/util:status_testutil", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings:cord", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tensorstore/kvstore/s3/credentials/aws_credentials.h b/tensorstore/kvstore/s3/credentials/aws_credentials.h new file mode 100644 index 000000000..fdd92607f --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/aws_credentials.h @@ -0,0 +1,63 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_AWS_CREDENTIALS_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_AWS_CREDENTIALS_H_ + +#include + +#include "absl/time/time.h" +#include "tensorstore/util/result.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +/// Holds AWS credentials +/// +/// Contains the access key, secret key, session token and expiry +/// An empty access key implies anonymous access, +/// while the presence of a session token implies the use of +/// short-lived STS credentials +/// https://docs.aws.amazon.com/STS/latest/APIReference/welcome.html +struct AwsCredentials { + /// AWS_ACCESS_KEY_ID + std::string access_key; + /// AWS_SECRET_KEY_ID + std::string secret_key; + /// AWS_SESSION_TOKEN + std::string session_token; + /// Expiration date + absl::Time expires_at = absl::InfinitePast(); + + /// Anonymous credentials that do not expire + static AwsCredentials Anonymous() { + return AwsCredentials{{}, {}, {}, absl::InfiniteFuture()}; + } + + bool IsAnonymous() const { return access_key.empty(); } +}; + +/// Base class for S3 Credential Providers +/// +/// Implementers should override GetCredentials. +class AwsCredentialProvider { + public: + virtual ~AwsCredentialProvider() = default; + virtual Result GetCredentials() = 0; +}; + +} // namespace internal_kvstore_s3 +} // namespace tensorstore + +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_AWS_CREDENTIALS_H_ diff --git a/tensorstore/kvstore/s3/credentials/default_credential_provider.cc b/tensorstore/kvstore/s3/credentials/default_credential_provider.cc new file mode 100644 index 000000000..6990425cb --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/default_credential_provider.cc @@ -0,0 +1,161 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/default_credential_provider.h" + +#include +#include +#include +#include +#include +#include + +#include "absl/functional/function_ref.h" +#include "absl/synchronization/mutex.h" +#include "absl/time/time.h" +#include "tensorstore/internal/http/http_transport.h" +#include "tensorstore/internal/no_destructor.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/kvstore/s3/credentials/ec2_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/environment_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/file_credential_provider.h" +#include "tensorstore/util/result.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { +namespace { + +/// Return a DefaultCredentialProvider that attempts to retrieve credentials +/// from +/// 1. AWS Environment Variables, e.g. AWS_ACCESS_KEY_ID +/// 2. Shared Credential File, e.g. $HOME/.aws/credentials +/// 3. EC2 Metadata Server +Result> GetDefaultAwsCredentialProvider( + std::string_view filename, std::string_view profile, + std::string_view metadata_endpoint, + std::shared_ptr transport) { + return std::make_unique( + DefaultAwsCredentialsProvider::Options{ + std::string{filename}, std::string{profile}, + std::string{metadata_endpoint}, transport}); +} + +struct AwsCredentialProviderRegistry { + std::vector> providers; + absl::Mutex mutex; +}; + +AwsCredentialProviderRegistry& GetAwsProviderRegistry() { + static internal::NoDestructor registry; + return *registry; +} + +} // namespace + +void RegisterAwsCredentialProviderProvider(AwsCredentialProviderFn provider, + int priority) { + auto& registry = GetAwsProviderRegistry(); + absl::WriterMutexLock lock(®istry.mutex); + registry.providers.emplace_back(priority, std::move(provider)); + std::sort(registry.providers.begin(), registry.providers.end(), + [](const auto& a, const auto& b) { return a.first < b.first; }); +} + +/// Obtain a credential provider from a series of registered and default +/// providers +/// +/// Providers are returned in the following order: +/// 1. Any registered providers that supply valid credentials +/// 2. Environment variable provider if valid credential can be obtained from +/// AWS_* environment variables +/// 3. File provider containing credentials from the $HOME/.aws/credentials +/// file. +/// The `profile` variable overrides the default profile in this file. +/// 4. EC2 Metadata server. The `transport` variable overrides the default +/// HttpTransport. +Result> GetAwsCredentialProvider( + std::string_view filename, std::string_view profile, + std::string_view metadata_endpoint, + std::shared_ptr transport) { + auto& registry = GetAwsProviderRegistry(); + absl::WriterMutexLock lock(®istry.mutex); + for (const auto& provider : registry.providers) { + auto credentials = provider.second(); + if (credentials.ok()) return credentials; + } + + return GetDefaultAwsCredentialProvider(filename, profile, metadata_endpoint, + transport); +} + +DefaultAwsCredentialsProvider::DefaultAwsCredentialsProvider( + Options options, absl::FunctionRef clock) + : options_(std::move(options)), + clock_(clock), + credentials_{{}, {}, {}, absl::InfinitePast()} {} + +Result DefaultAwsCredentialsProvider::GetCredentials() { + { + absl::ReaderMutexLock lock(&mutex_); + if (credentials_.expires_at > clock_()) { + return credentials_; + } + } + + absl::WriterMutexLock lock(&mutex_); + + // Refresh existing credentials + if (provider_) { + auto credentials_result = provider_->GetCredentials(); + if (credentials_result.ok()) { + credentials_ = credentials_result.value(); + return credentials_; + } + } + + // Return credentials in this order: + // 1. AWS Environment Variables, e.g. AWS_ACCESS_KEY_ID + provider_ = std::make_unique(); + auto credentials_result = provider_->GetCredentials(); + if (credentials_result.ok()) { + credentials_ = credentials_result.value(); + return credentials_; + } + + // 2. Shared Credential File, e.g. $HOME/.aws/credentials + provider_ = std::make_unique(options_.filename, + options_.profile); + credentials_result = provider_->GetCredentials(); + if (credentials_result.ok()) { + credentials_ = credentials_result.value(); + return credentials_; + } + + // 3. EC2 Metadata Server + provider_ = std::make_unique( + options_.endpoint, options_.transport); + credentials_result = provider_->GetCredentials(); + if (credentials_result.ok()) { + credentials_ = credentials_result.value(); + return credentials_; + } + + // 4. Anonymous credentials + provider_ = nullptr; + credentials_ = AwsCredentials::Anonymous(); + return credentials_; +} + +} // namespace internal_kvstore_s3 +} // namespace tensorstore diff --git a/tensorstore/kvstore/s3/credentials/default_credential_provider.h b/tensorstore/kvstore/s3/credentials/default_credential_provider.h new file mode 100644 index 000000000..ecd8c3bcd --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/default_credential_provider.h @@ -0,0 +1,88 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_DEFAULT_CREDENTIAL_PROVIDER_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_DEFAULT_CREDENTIAL_PROVIDER_H_ + +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/functional/function_ref.h" +#include "absl/synchronization/mutex.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "tensorstore/internal/http/curl_transport.h" +#include "tensorstore/internal/http/http_transport.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/util/result.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +/// A provider that implements a Default strategy for retrieving +/// and caching a set of AWS credentials from the following sources: +/// +/// 1. Environment variables such as AWS_ACCESS_KEY_ID +/// 2. Shared Credential Files such as ~/.aws/credentials +/// 3. The EC2 Metadata server +/// +/// The cached credentials are returned until they expire, +/// at which point the original source is queried again to +/// obtain fresher credentials +class DefaultAwsCredentialsProvider : public AwsCredentialProvider { + public: + /// Options to configure the provider. These include the: + /// + /// 1. Shared Credential Filename + /// 2. Shared Credential File Profile + /// 3. EC2 Metadata Server Endpoint + /// 3. Http Transport for querying the EC2 Metadata Server + struct Options { + std::string filename; + std::string profile; + std::string endpoint; + std::shared_ptr transport; + }; + + DefaultAwsCredentialsProvider( + Options options = {{}, {}, {}, internal_http::GetDefaultHttpTransport()}, + absl::FunctionRef clock = absl::Now); + Result GetCredentials() override; + + private: + Options options_; + absl::FunctionRef clock_; + absl::Mutex mutex_; + std::unique_ptr provider_ ABSL_GUARDED_BY(mutex_); + AwsCredentials credentials_ ABSL_GUARDED_BY(mutex_); +}; + +using AwsCredentialProviderFn = + std::function>()>; + +void RegisterAwsCredentialProviderProvider(AwsCredentialProviderFn provider, + int priority); + +Result> GetAwsCredentialProvider( + std::string_view filename, std::string_view profile, + std::string_view metadata_endpoint, + std::shared_ptr transport); + +} // namespace internal_kvstore_s3 +} // namespace tensorstore + +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_DEFAULT_CREDENTIAL_PROVIDER_H_ diff --git a/tensorstore/kvstore/s3/credentials/default_credential_provider_test.cc b/tensorstore/kvstore/s3/credentials/default_credential_provider_test.cc new file mode 100644 index 000000000..0eb533b18 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/default_credential_provider_test.cc @@ -0,0 +1,198 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/default_credential_provider.h" + +#include +#include +#include + +#include +#include "absl/container/flat_hash_map.h" +#include "absl/strings/cord.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "tensorstore/internal/env.h" +#include "tensorstore/internal/http/http_response.h" +#include "tensorstore/internal/path.h" +#include "tensorstore/internal/test_util.h" +#include "tensorstore/kvstore/s3/credentials/test_utils.h" +#include "tensorstore/util/status_testutil.h" + +namespace { + +using ::tensorstore::internal::JoinPath; +using ::tensorstore::internal::SetEnv; +using ::tensorstore::internal::UnsetEnv; +using ::tensorstore::internal_http::HttpResponse; +using ::tensorstore::internal_kvstore_s3::DefaultAwsCredentialsProvider; +using ::tensorstore::internal_kvstore_s3::DefaultEC2MetadataFlow; +using ::tensorstore::internal_kvstore_s3::EC2MetadataMockTransport; +using Options = + ::tensorstore::internal_kvstore_s3::DefaultAwsCredentialsProvider::Options; + +static constexpr char kEndpoint[] = "http://endpoint"; + +class CredentialFileFactory + : public tensorstore::internal::ScopedTemporaryDirectory { + public: + std::string WriteCredentialsFile() { + auto p = JoinPath(path(), "aws_config"); + std::ofstream ofs(p); + ofs << "[alice]\n" + "aws_access_key_id = AKIAIOSFODNN6EXAMPLE\n" + "aws_secret_access_key = " + "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY\n" + "aws_session_token = abcdef1234567890\n" + "\n"; + ofs.close(); + return p; + } +}; + +class DefaultCredentialProviderTest : public ::testing::Test { + protected: + void SetUp() override { + UnsetEnv("AWS_ACCESS_KEY_ID"); + UnsetEnv("AWS_SECRET_KEY_ID"); + UnsetEnv("AWS_SESSION_TOKEN"); + } +}; + +TEST_F(DefaultCredentialProviderTest, AnonymousCredentials) { + auto url_to_response = absl::flat_hash_map(); + auto mock_transport = + std::make_shared(url_to_response); + auto provider = std::make_unique( + Options{{}, {}, {}, mock_transport}); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, + provider->GetCredentials()); + EXPECT_TRUE(credentials.IsAnonymous()); + EXPECT_EQ(credentials.expires_at, absl::InfiniteFuture()); + + // Idempotent + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials2, + provider->GetCredentials()); + EXPECT_TRUE(credentials2.IsAnonymous()); + EXPECT_EQ(credentials2.expires_at, absl::InfiniteFuture()); +} + +TEST_F(DefaultCredentialProviderTest, EnvironmentCredentialIdempotency) { + SetEnv("AWS_ACCESS_KEY_ID", "access"); + SetEnv("AWS_SECRET_ACCESS_KEY", "secret"); + SetEnv("AWS_SESSION_TOKEN", "token"); + + auto provider = std::make_unique(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, + provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, "access"); + EXPECT_EQ(credentials.secret_key, "secret"); + EXPECT_EQ(credentials.session_token, "token"); + EXPECT_EQ(credentials.expires_at, absl::InfiniteFuture()); + + // Expect idempotency as environment credentials never expire + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials2, + provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, credentials2.access_key); + EXPECT_EQ(credentials.secret_key, credentials2.secret_key); + EXPECT_EQ(credentials.session_token, credentials2.session_token); + EXPECT_EQ(credentials.expires_at, credentials2.expires_at); +} + +/// Test configuration of FileCredentialProvider from +/// DefaultAwsCredentialsProvider::Options +TEST_F(DefaultCredentialProviderTest, ConfigureFileProviderFromOptions) { + auto factory = CredentialFileFactory{}; + auto credentials_file = factory.WriteCredentialsFile(); + + auto provider = std::make_unique( + Options{credentials_file, "alice"}); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, + provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); + EXPECT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); + EXPECT_EQ(credentials.session_token, "abcdef1234567890"); + EXPECT_EQ(credentials.expires_at, absl::InfiniteFuture()); + + // Expect idempotency as file credentials never expire + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials2, + provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, credentials2.access_key); + EXPECT_EQ(credentials.secret_key, credentials2.secret_key); + EXPECT_EQ(credentials.session_token, credentials2.session_token); + EXPECT_EQ(credentials.expires_at, credentials2.expires_at); +} + +/// Test configuration of EC2MetaDataProvider from +/// DefaultAwsCredentialsProvider::Options +TEST_F(DefaultCredentialProviderTest, ConfigureEC2ProviderFromOptions) { + auto now = absl::Now(); + auto stuck_clock = [&]() -> absl::Time { return now; }; + auto expiry = now + absl::Seconds(200); + auto url_to_response = DefaultEC2MetadataFlow( + kEndpoint, "1234", "ASIA1234567890", "1234567890abcdef", "token", expiry); + + auto mock_transport = + std::make_shared(url_to_response); + + auto provider = std::make_unique( + Options{{}, {}, kEndpoint, mock_transport}, stuck_clock); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, + provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, "ASIA1234567890"); + EXPECT_EQ(credentials.secret_key, "1234567890abcdef"); + EXPECT_EQ(credentials.session_token, "token"); + EXPECT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); + + /// Force failure on credential retrieval + url_to_response = absl::flat_hash_map{ + {"POST http://endpoint/latest/api/token", + HttpResponse{404, absl::Cord{""}}}, + }; + + // But we're not expired, so we get the original credentials + TENSORSTORE_ASSERT_OK_AND_ASSIGN(credentials, provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, "ASIA1234567890"); + EXPECT_EQ(credentials.secret_key, "1234567890abcdef"); + EXPECT_EQ(credentials.session_token, "token"); + EXPECT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); + + // Force expiry and retrieve new credentials + now += absl::Seconds(300); + url_to_response = DefaultEC2MetadataFlow(kEndpoint, "1234", "ASIA1234567890", + "1234567890abcdef", "TOKEN", expiry); + + // A new set of credentials is returned + TENSORSTORE_ASSERT_OK_AND_ASSIGN(credentials, provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, "ASIA1234567890"); + EXPECT_EQ(credentials.secret_key, "1234567890abcdef"); + EXPECT_EQ(credentials.session_token, "TOKEN"); + EXPECT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); + + /// Force failure on credential retrieval + url_to_response = absl::flat_hash_map{ + {"POST http://endpoint/latest/api/token", + HttpResponse{404, absl::Cord{""}}}, + }; + + // Anonymous credentials + TENSORSTORE_ASSERT_OK_AND_ASSIGN(credentials, provider->GetCredentials()); + EXPECT_EQ(credentials.access_key, ""); + EXPECT_EQ(credentials.secret_key, ""); + EXPECT_EQ(credentials.session_token, ""); + EXPECT_EQ(credentials.expires_at, absl::InfiniteFuture()); +} + +} // namespace diff --git a/tensorstore/kvstore/s3/aws_metadata_credential_provider.cc b/tensorstore/kvstore/s3/credentials/ec2_credential_provider.cc similarity index 83% rename from tensorstore/kvstore/s3/aws_metadata_credential_provider.cc rename to tensorstore/kvstore/s3/credentials/ec2_credential_provider.cc index daec0365a..230c8e57f 100644 --- a/tensorstore/kvstore/s3/aws_metadata_credential_provider.cc +++ b/tensorstore/kvstore/s3/credentials/ec2_credential_provider.cc @@ -12,9 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "tensorstore/kvstore/s3/aws_metadata_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/ec2_credential_provider.h" -#include #include #include #include @@ -25,9 +24,9 @@ #include "absl/strings/cord.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_split.h" -#include "absl/synchronization/mutex.h" #include "absl/time/clock.h" #include "absl/time/time.h" +#include #include "tensorstore/internal/env.h" #include "tensorstore/internal/http/http_request.h" #include "tensorstore/internal/http/http_response.h" @@ -35,7 +34,7 @@ #include "tensorstore/internal/json/json.h" #include "tensorstore/internal/json_binding/bindable.h" #include "tensorstore/internal/json_binding/json_binding.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" #include "tensorstore/util/result.h" #include "tensorstore/util/status.h" #include "tensorstore/util/str_cat.h" @@ -113,12 +112,12 @@ inline constexpr auto EC2CredentialsResponseBinder = jb::Object( jb::Projection(&EC2CredentialsResponse::expiration))); // Obtain a metadata token for communicating with the api server. -Result GetEC2ApiToken(internal_http::HttpTransport& transport) { +Result GetEC2ApiToken(std::string_view endpoint, + internal_http::HttpTransport& transport) { // Obtain Metadata server API tokens with a TTL of 21600 seconds auto token_request = HttpRequestBuilder("POST", - tensorstore::StrCat(GetEC2MetadataServiceEndpoint(), - "/latest/api/token")) + tensorstore::StrCat(endpoint, "/latest/api/token")) .AddHeader("x-aws-ec2-metadata-token-ttl-seconds: 21600") .BuildRequest(); @@ -135,11 +134,6 @@ Result GetEC2ApiToken(internal_http::HttpTransport& transport) { } // namespace -// Returns whether the EC2 Metadata Server is available. -bool IsEC2MetadataServiceAvailable(internal_http::HttpTransport& transport) { - return GetEC2ApiToken(transport).ok(); -} - /// Obtains AWS Credentials from the EC2Metadata. /// /// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html#instancedata-meta-data-retrieval-examples @@ -153,21 +147,18 @@ bool IsEC2MetadataServiceAvailable(internal_http::HttpTransport& transport) { /// 3. Obtain the associated credentials from path /// "/latest/meta-data/iam/security-credentials/". Result EC2MetadataCredentialProvider::GetCredentials() { - absl::MutexLock l(&mutex_); - if (absl::Now() + absl::Seconds(60) < timeout_) { - return credentials_; + if (endpoint_.empty()) { + endpoint_ = GetEC2MetadataServiceEndpoint(); } - auto default_timeout = absl::Now() + kDefaultTimeout; - // Obtain an API token for communicating with the EC2 Metadata server - TENSORSTORE_ASSIGN_OR_RETURN(auto api_token, GetEC2ApiToken(*transport_)); + TENSORSTORE_ASSIGN_OR_RETURN(auto api_token, + GetEC2ApiToken(endpoint_, *transport_)); auto token_header = tensorstore::StrCat(kMetadataTokenHeader, api_token); auto iam_role_request = HttpRequestBuilder("GET", - tensorstore::StrCat(GetEC2MetadataServiceEndpoint(), - kIamCredentialsPath)) + tensorstore::StrCat(endpoint_, kIamCredentialsPath)) .AddHeader(token_header) .BuildRequest(); @@ -184,8 +175,8 @@ Result EC2MetadataCredentialProvider::GetCredentials() { return absl::NotFoundError("Empty EC2 Role list"); } - auto iam_credentials_request_url = tensorstore::StrCat( - GetEC2MetadataServiceEndpoint(), kIamCredentialsPath, iam_roles[0]); + auto iam_credentials_request_url = + tensorstore::StrCat(endpoint_, kIamCredentialsPath, iam_roles[0]); auto iam_credentials_request = HttpRequestBuilder("GET", iam_credentials_request_url) @@ -213,11 +204,14 @@ Result EC2MetadataCredentialProvider::GetCredentials() { "] failed with ", json_sv)); } - timeout_ = iam_credentials.expiration.value_or(default_timeout); - credentials_ = AwsCredentials{iam_credentials.access_key_id.value_or(""), - iam_credentials.secret_access_key.value_or(""), - iam_credentials.token.value_or("")}; - return credentials_; + // Introduce a leeway of 60 seconds to avoid credential expiry conditions + auto default_timeout = absl::Now() + kDefaultTimeout; + auto expires_at = + iam_credentials.expiration.value_or(default_timeout) - absl::Seconds(60); + + return AwsCredentials{iam_credentials.access_key_id.value_or(""), + iam_credentials.secret_access_key.value_or(""), + iam_credentials.token.value_or(""), expires_at}; } } // namespace internal_kvstore_s3 diff --git a/tensorstore/kvstore/s3/aws_metadata_credential_provider.h b/tensorstore/kvstore/s3/credentials/ec2_credential_provider.h similarity index 63% rename from tensorstore/kvstore/s3/aws_metadata_credential_provider.h rename to tensorstore/kvstore/s3/credentials/ec2_credential_provider.h index 2e22992fc..d19aafbb4 100644 --- a/tensorstore/kvstore/s3/aws_metadata_credential_provider.h +++ b/tensorstore/kvstore/s3/credentials/ec2_credential_provider.h @@ -12,17 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -#ifndef TENSORSTORE_KVSTORE_S3_AWS_METADATA_CREDENTIAL_PROVIDER_H -#define TENSORSTORE_KVSTORE_S3_AWS_METADATA_CREDENTIAL_PROVIDER_H +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_EC2_CREDENTIAL_PROVIDER_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_EC2_CREDENTIAL_PROVIDER_H_ #include +#include +#include #include -#include "absl/base/thread_annotations.h" -#include "absl/synchronization/mutex.h" -#include "absl/time/time.h" #include "tensorstore/internal/http/http_transport.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" #include "tensorstore/util/result.h" namespace tensorstore { @@ -32,23 +31,19 @@ namespace internal_kvstore_s3 { class EC2MetadataCredentialProvider : public AwsCredentialProvider { public: EC2MetadataCredentialProvider( + std::string_view endpoint, std::shared_ptr transport) - : transport_(std::move(transport)), timeout_(absl::InfinitePast()) {} + : endpoint_(endpoint), transport_(std::move(transport)) {} Result GetCredentials() override; + inline const std::string& GetEndpoint() const { return endpoint_; } private: + std::string endpoint_; std::shared_ptr transport_; - - absl::Mutex mutex_; - absl::Time timeout_ ABSL_GUARDED_BY(mutex_); - AwsCredentials credentials_ ABSL_GUARDED_BY(mutex_); }; -// Returns whether the EC2 Metadata Server is available. -bool IsEC2MetadataServiceAvailable(internal_http::HttpTransport& transport); - } // namespace internal_kvstore_s3 } // namespace tensorstore -#endif // TENSORSTORE_KVSTORE_S3_AWS_METADATA_CREDENTIAL_PROVIDER_H +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_EC2_CREDENTIAL_PROVIDER_H_ diff --git a/tensorstore/kvstore/s3/credentials/ec2_credential_provider_test.cc b/tensorstore/kvstore/s3/credentials/ec2_credential_provider_test.cc new file mode 100644 index 000000000..f2707efe7 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/ec2_credential_provider_test.cc @@ -0,0 +1,166 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/ec2_credential_provider.h" + +#include +#include + +#include +#include +#include "absl/container/flat_hash_map.h" +#include "absl/status/status.h" +#include "absl/strings/cord.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "tensorstore/internal/env.h" +#include "tensorstore/internal/http/http_response.h" +#include "tensorstore/kvstore/s3/credentials/test_utils.h" +#include "tensorstore/util/result.h" +#include "tensorstore/util/status_testutil.h" + +using ::tensorstore::internal::SetEnv; +using ::tensorstore::internal::UnsetEnv; + +namespace { + +using ::tensorstore::MatchesStatus; +using ::tensorstore::internal_http::HttpResponse; +using ::tensorstore::internal_kvstore_s3::DefaultEC2MetadataFlow; +using ::tensorstore::internal_kvstore_s3::EC2MetadataCredentialProvider; +using ::tensorstore::internal_kvstore_s3::EC2MetadataMockTransport; + +static constexpr char kDefaultEndpoint[] = "http://169.254.169.254"; +static constexpr char kCustomEndpoint[] = "http://custom.endpoint"; +static constexpr char kApiToken[] = "1234567890"; +static constexpr char kAccessKey[] = "ASIA1234567890"; +static constexpr char kSecretKey[] = "1234567890abcdef"; +static constexpr char kSessionToken[] = "abcdef123456790"; + +class EC2MetadataCredentialProviderTest : public ::testing::Test { + protected: + void SetUp() override { UnsetEnv("AWS_EC2_METADATA_SERVICE_ENDPOINT"); } +}; + +TEST_F(EC2MetadataCredentialProviderTest, CredentialRetrievalFlow) { + auto expiry = absl::Now() + absl::Seconds(200); + auto url_to_response = + DefaultEC2MetadataFlow(kDefaultEndpoint, kApiToken, kAccessKey, + kSecretKey, kSessionToken, expiry); + + auto mock_transport = + std::make_shared(url_to_response); + auto provider = + std::make_shared("", mock_transport); + TENSORSTORE_CHECK_OK_AND_ASSIGN(auto credentials, provider->GetCredentials()); + ASSERT_EQ(provider->GetEndpoint(), kDefaultEndpoint); + ASSERT_EQ(credentials.access_key, kAccessKey); + ASSERT_EQ(credentials.secret_key, kSecretKey); + ASSERT_EQ(credentials.session_token, kSessionToken); + // expiry less the 60s leeway + ASSERT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); +} + +TEST_F(EC2MetadataCredentialProviderTest, EnvironmentVariableMetadataServer) { + SetEnv("AWS_EC2_METADATA_SERVICE_ENDPOINT", kCustomEndpoint); + auto expiry = absl::Now() + absl::Seconds(200); + auto url_to_response = + DefaultEC2MetadataFlow(kCustomEndpoint, kApiToken, kAccessKey, kSecretKey, + kSessionToken, expiry); + + auto mock_transport = + std::make_shared(url_to_response); + auto provider = + std::make_shared("", mock_transport); + TENSORSTORE_CHECK_OK_AND_ASSIGN(auto credentials, provider->GetCredentials()); + ASSERT_EQ(provider->GetEndpoint(), kCustomEndpoint); + ASSERT_EQ(credentials.access_key, kAccessKey); + ASSERT_EQ(credentials.secret_key, kSecretKey); + ASSERT_EQ(credentials.session_token, kSessionToken); + // expiry less the 60s leeway + ASSERT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); +} + +TEST_F(EC2MetadataCredentialProviderTest, InjectedMetadataServer) { + auto expiry = absl::Now() + absl::Seconds(200); + auto url_to_response = + DefaultEC2MetadataFlow(kCustomEndpoint, kApiToken, kAccessKey, kSecretKey, + kSessionToken, expiry); + + auto mock_transport = + std::make_shared(url_to_response); + auto provider = std::make_shared( + kCustomEndpoint, mock_transport); + TENSORSTORE_CHECK_OK_AND_ASSIGN(auto credentials, provider->GetCredentials()); + ASSERT_EQ(provider->GetEndpoint(), kCustomEndpoint); + ASSERT_EQ(credentials.access_key, kAccessKey); + ASSERT_EQ(credentials.secret_key, kSecretKey); + ASSERT_EQ(credentials.session_token, kSessionToken); + // expiry less the 60s leeway + ASSERT_EQ(credentials.expires_at, expiry - absl::Seconds(60)); +} + +TEST_F(EC2MetadataCredentialProviderTest, NoIamRolesInSecurityCredentials) { + auto url_to_response = absl::flat_hash_map{ + {"POST http://169.254.169.254/latest/api/token", + HttpResponse{200, absl::Cord{kApiToken}}}, + {"GET http://169.254.169.254/latest/meta-data/iam/security-credentials/", + HttpResponse{ + 200, absl::Cord{""}, {{"x-aws-ec2-metadata-token", kApiToken}}}}, + }; + + auto mock_transport = + std::make_shared(url_to_response); + auto provider = + std::make_shared("", mock_transport); + ASSERT_FALSE(provider->GetCredentials()); + ASSERT_EQ(provider->GetEndpoint(), kDefaultEndpoint); + EXPECT_THAT(provider->GetCredentials().status().ToString(), + ::testing::HasSubstr("Empty EC2 Role list")); +} + +TEST_F(EC2MetadataCredentialProviderTest, UnsuccessfulJsonResponse) { + // Test that "Code" != "Success" parsing succeeds + auto url_to_response = absl::flat_hash_map{ + {"POST http://169.254.169.254/latest/api/token", + HttpResponse{200, absl::Cord{kApiToken}}}, + {"GET http://169.254.169.254/latest/meta-data/iam/", + HttpResponse{ + 200, absl::Cord{"info"}, {{"x-aws-ec2-metadata-token", kApiToken}}}}, + {"GET http://169.254.169.254/latest/meta-data/iam/security-credentials/", + HttpResponse{200, + absl::Cord{"mock-iam-role"}, + {{"x-aws-ec2-metadata-token", kApiToken}}}}, + {"GET " + "http://169.254.169.254/latest/meta-data/iam/security-credentials/" + "mock-iam-role", + HttpResponse{200, + absl::Cord(R"({"Code": "EntirelyUnsuccessful"})"), + {{"x-aws-ec2-metadata-token", kApiToken}}}}}; + + auto mock_transport = + std::make_shared(url_to_response); + auto provider = + std::make_shared("", mock_transport); + auto credentials = provider->GetCredentials(); + + EXPECT_THAT(credentials.status(), MatchesStatus(absl::StatusCode::kNotFound)); + EXPECT_THAT(credentials.status().ToString(), + ::testing::AllOf( + ::testing::HasSubstr("EC2Metadata request"), + ::testing::HasSubstr( + "failed with {\"Code\": \"EntirelyUnsuccessful\"}"))); +} + +} // namespace diff --git a/tensorstore/kvstore/s3/credentials/environment_credential_provider.cc b/tensorstore/kvstore/s3/credentials/environment_credential_provider.cc new file mode 100644 index 000000000..bac1a6326 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/environment_credential_provider.cc @@ -0,0 +1,60 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/environment_credential_provider.h" + +#include "absl/log/absl_log.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "absl/time/time.h" +#include "tensorstore/internal/env.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/util/result.h" + +using ::tensorstore::internal::GetEnv; + +namespace tensorstore { +namespace internal_kvstore_s3 { + +namespace { + +// AWS user identifier +static constexpr char kEnvAwsAccessKeyId[] = "AWS_ACCESS_KEY_ID"; +// AWS user password +static constexpr char kEnvAwsSecretAccessKey[] = "AWS_SECRET_ACCESS_KEY"; +// AWS session token +static constexpr char kEnvAwsSessionToken[] = "AWS_SESSION_TOKEN"; + +} // namespace + +Result EnvironmentCredentialProvider::GetCredentials() { + auto access_key = GetEnv(kEnvAwsAccessKeyId); + if (!access_key) { + return absl::NotFoundError(absl::StrCat(kEnvAwsAccessKeyId, " not set")); + } + ABSL_LOG_FIRST_N(INFO, 1) + << "Using Environment Variable " << kEnvAwsAccessKeyId; + auto credentials = AwsCredentials{*access_key}; + if (auto secret_key = GetEnv(kEnvAwsSecretAccessKey); secret_key) { + credentials.secret_key = *secret_key; + } + if (auto session_token = GetEnv(kEnvAwsSessionToken); session_token) { + credentials.session_token = *session_token; + } + credentials.expires_at = absl::InfiniteFuture(); + return credentials; +} + +} // namespace internal_kvstore_s3 +} // namespace tensorstore diff --git a/tensorstore/kvstore/s3/credentials/environment_credential_provider.h b/tensorstore/kvstore/s3/credentials/environment_credential_provider.h new file mode 100644 index 000000000..e73202fd3 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/environment_credential_provider.h @@ -0,0 +1,34 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_ENVIRONMENT_CREDENTIAL_PROVIDER_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_ENVIRONMENT_CREDENTIAL_PROVIDER_H_ + +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/util/result.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +/// Provides credentials from the following environment variables: +/// AWS_ACCESS_KEY_ID, AWS_SECRET_KEY_ID, AWS_SESSION_TOKEN +class EnvironmentCredentialProvider : public AwsCredentialProvider { + public: + Result GetCredentials() override; +}; + +} // namespace internal_kvstore_s3 +} // namespace tensorstore + +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_ENVIRONMENT_CREDENTIAL_PROVIDER_H_ diff --git a/tensorstore/kvstore/s3/credentials/environment_credential_provider_test.cc b/tensorstore/kvstore/s3/credentials/environment_credential_provider_test.cc new file mode 100644 index 000000000..2771ff0a8 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/environment_credential_provider_test.cc @@ -0,0 +1,60 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/environment_credential_provider.h" + +#include +#include "tensorstore/internal/env.h" +#include "tensorstore/util/status_testutil.h" + +namespace { + +using ::tensorstore::internal::SetEnv; +using ::tensorstore::internal::UnsetEnv; +using ::tensorstore::internal_kvstore_s3::EnvironmentCredentialProvider; + +class EnvironmentCredentialProviderTest : public ::testing::Test { + protected: + void SetUp() override { + // Make sure that env vars are not set. + for (const char* var : + {"AWS_SHARED_CREDENTIALS_FILE", "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "AWS_PROFILE"}) { + UnsetEnv(var); + } + } +}; + +TEST_F(EnvironmentCredentialProviderTest, ProviderNoCredentials) { + auto provider = EnvironmentCredentialProvider(); + ASSERT_FALSE(provider.GetCredentials().ok()); + SetEnv("AWS_ACCESS_KEY_ID", "foo"); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, provider.GetCredentials()); + ASSERT_EQ(credentials.access_key, "foo"); + ASSERT_TRUE(credentials.secret_key.empty()); + ASSERT_TRUE(credentials.session_token.empty()); +} + +TEST_F(EnvironmentCredentialProviderTest, ProviderAwsCredentialsFromEnv) { + SetEnv("AWS_ACCESS_KEY_ID", "foo"); + SetEnv("AWS_SECRET_ACCESS_KEY", "bar"); + SetEnv("AWS_SESSION_TOKEN", "qux"); + auto provider = EnvironmentCredentialProvider(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, provider.GetCredentials()); + ASSERT_EQ(credentials.access_key, "foo"); + ASSERT_EQ(credentials.secret_key, "bar"); + ASSERT_EQ(credentials.session_token, "qux"); +} + +} // namespace diff --git a/tensorstore/kvstore/s3/credentials/file_credential_provider.cc b/tensorstore/kvstore/s3/credentials/file_credential_provider.cc new file mode 100644 index 000000000..27770df6f --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/file_credential_provider.cc @@ -0,0 +1,133 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/file_credential_provider.h" + +#include +#include + +#include "absl/log/absl_log.h" +#include "absl/status/status.h" +#include "absl/strings/ascii.h" +#include "absl/strings/str_cat.h" +#include "absl/time/time.h" +#include "tensorstore/internal/env.h" +#include "tensorstore/internal/path.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/util/result.h" + +using ::tensorstore::internal::GetEnv; +using ::tensorstore::internal::JoinPath; + +namespace tensorstore { +namespace internal_kvstore_s3 { + +namespace { + +// Credentials file environment variable +static constexpr char kEnvAwsCredentialsFile[] = "AWS_SHARED_CREDENTIALS_FILE"; +// Default path to the AWS credentials file, relative to the home folder +static constexpr char kDefaultAwsCredentialsFilePath[] = ".aws/credentials"; +// AWS user identifier +static constexpr char kCfgAwsAccessKeyId[] = "aws_access_key_id"; +// AWS user password +static constexpr char kCfgAwsSecretAccessKeyId[] = "aws_secret_access_key"; +// AWS session token +static constexpr char kCfgAwsSessionToken[] = "aws_session_token"; +// Discover AWS profile in environment variables +static constexpr char kEnvAwsProfile[] = "AWS_PROFILE"; +// Default profile +static constexpr char kDefaultProfile[] = "default"; + +Result GetAwsCredentialsFileName() { + auto credentials_file = GetEnv(kEnvAwsCredentialsFile); + if (!credentials_file) { + auto home_dir = GetEnv("HOME"); + if (!home_dir) { + return absl::NotFoundError("Could not read $HOME"); + } + return JoinPath(*home_dir, kDefaultAwsCredentialsFilePath); + } + return *credentials_file; +} + +} // namespace + +/// https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html#cli-configure-files-format +Result FileCredentialProvider::GetCredentials() { + if (filename_.empty()) { + TENSORSTORE_ASSIGN_OR_RETURN(filename_, GetAwsCredentialsFileName()); + } + + if (profile_.empty()) { + profile_ = GetEnv(kEnvAwsProfile).value_or(kDefaultProfile); + } + + std::ifstream ifs(filename_.c_str()); + if (!ifs) { + return absl::NotFoundError( + absl::StrCat("Could not open credentials file [", filename_, "]")); + } + + AwsCredentials credentials{}; + std::string section_name; + std::string line; + bool profile_found = false; + + while (std::getline(ifs, line)) { + auto sline = absl::StripAsciiWhitespace(line); + // Ignore empty and commented out lines + if (sline.empty() || sline[0] == '#') continue; + + // A configuration section name has been encountered + if (sline[0] == '[' && sline[sline.size() - 1] == ']') { + section_name = + absl::StripAsciiWhitespace(sline.substr(1, sline.size() - 2)); + continue; + } + + // Look for key=value pairs if we're in the appropriate profile + if (section_name == profile_) { + profile_found = true; + if (auto pos = sline.find('='); pos != std::string::npos) { + auto key = absl::StripAsciiWhitespace(sline.substr(0, pos)); + auto value = absl::StripAsciiWhitespace(sline.substr(pos + 1)); + + if (key == kCfgAwsAccessKeyId) { + credentials.access_key = value; + } else if (key == kCfgAwsSecretAccessKeyId) { + credentials.secret_key = value; + } else if (key == kCfgAwsSessionToken) { + credentials.session_token = value; + } + } + } + } + + if (!profile_found) { + return absl::NotFoundError(absl::StrCat("Profile [", profile_, + "] not found " + "in credentials file [", + filename_, "]")); + } + + ABSL_LOG_FIRST_N(INFO, 1) + << "Using profile [" << profile_ << "] in file [" << filename_ << "]"; + + credentials.expires_at = absl::InfiniteFuture(); + return credentials; +} + +} // namespace internal_kvstore_s3 +} // namespace tensorstore diff --git a/tensorstore/kvstore/s3/credentials/file_credential_provider.h b/tensorstore/kvstore/s3/credentials/file_credential_provider.h new file mode 100644 index 000000000..254446119 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/file_credential_provider.h @@ -0,0 +1,54 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_FILE_CREDENTIAL_PROVIDER_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_FILE_CREDENTIAL_PROVIDER_H_ + +#include +#include + +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/util/result.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +/// Obtains S3 credentials from a profile in a file, usually +/// `~/.aws/credentials` or a file specified in AWS_SHARED_CREDENTIALS_FILE. +/// A filename or desired profile may be specified in the constructor: +/// These values should be derived from the s3 json spec. +/// However, if filename is passed as an empty string, the filename is +/// obtained from AWS_SHARED_CREDENTIAL_FILE before defaulting to +/// `.aws/credentials`. +/// Likewise, if profile is passed as an empty string, +/// the profile is obtained from AWS_DEFAULT_PROFILE, AWS_PROFILE before +/// finally defaulting to "default". +class FileCredentialProvider : public AwsCredentialProvider { + private: + std::string filename_; + std::string profile_; + + public: + FileCredentialProvider(std::string_view filename, std::string_view profile) + : filename_(filename), profile_(profile) {} + + Result GetCredentials() override; + inline const std::string& GetFileName() const { return filename_; } + inline const std::string& GetProfile() const { return profile_; } +}; + +} // namespace internal_kvstore_s3 +} // namespace tensorstore + +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_FILE_CREDENTIAL_PROVIDER_H_ diff --git a/tensorstore/kvstore/s3/credentials/file_credential_provider_test.cc b/tensorstore/kvstore/s3/credentials/file_credential_provider_test.cc new file mode 100644 index 000000000..4408cac8b --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/file_credential_provider_test.cc @@ -0,0 +1,150 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/file_credential_provider.h" + +#include +#include +#include + +#include +#include "absl/time/time.h" +#include "tensorstore/internal/env.h" +#include "tensorstore/internal/path.h" +#include "tensorstore/internal/test_util.h" +#include "tensorstore/util/status_testutil.h" + +namespace { + +using ::tensorstore::internal::JoinPath; +using ::tensorstore::internal::SetEnv; +using ::tensorstore::internal::UnsetEnv; +using ::tensorstore::internal_kvstore_s3::FileCredentialProvider; + +class TestData : public tensorstore::internal::ScopedTemporaryDirectory { + public: + std::string WriteCredentialsFile() { + auto p = JoinPath(path(), "aws_config"); + std::ofstream ofs(p); + ofs << "discarded_value = 500\n" + "\n" + "[default]\n" + "aws_access_key_id =AKIAIOSFODNN7EXAMPLE\n" + "aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\n" + "aws_session_token= abcdef1234567890 \n" + "\n" + "[alice]\n" + "aws_access_key_id = AKIAIOSFODNN6EXAMPLE\n" + "aws_secret_access_key = " + "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY\n" + "\n"; + ofs.close(); + return p; + } +}; + +class FileCredentialProviderTest : public ::testing::Test { + protected: + void SetUp() override { + UnsetEnv("AWS_SHARED_CREDENTIALS_FILE"); + UnsetEnv("AWS_PROFILE"); + } +}; + +TEST_F(FileCredentialProviderTest, ProviderAwsCredentialsFromFileDefault) { + TestData test_data; + std::string credentials_filename = test_data.WriteCredentialsFile(); + + SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); + auto provider = FileCredentialProvider("", ""); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, provider.GetCredentials()); + ASSERT_EQ(provider.GetFileName(), credentials_filename); + ASSERT_EQ(provider.GetProfile(), "default"); + ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN7EXAMPLE"); + ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"); + ASSERT_EQ(credentials.session_token, "abcdef1234567890"); + ASSERT_EQ(credentials.expires_at, absl::InfiniteFuture()); +} + +TEST_F(FileCredentialProviderTest, + ProviderAwsCredentialsFromFileProfileOverride) { + TestData test_data; + auto credentials_filename = test_data.WriteCredentialsFile(); + + SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); + auto provider = FileCredentialProvider("", "alice"); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, provider.GetCredentials()); + ASSERT_EQ(provider.GetFileName(), credentials_filename); + ASSERT_EQ(provider.GetProfile(), "alice"); + ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); + ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); + ASSERT_EQ(credentials.session_token, ""); + ASSERT_EQ(credentials.expires_at, absl::InfiniteFuture()); +} + +TEST_F(FileCredentialProviderTest, ProviderAwsCredentialsFromFileProfileEnv) { + TestData test_data; + auto credentials_filename = test_data.WriteCredentialsFile(); + + SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); + SetEnv("AWS_PROFILE", "alice"); + auto provider = FileCredentialProvider("", ""); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, provider.GetCredentials()); + ASSERT_EQ(provider.GetFileName(), credentials_filename); + ASSERT_EQ(provider.GetProfile(), "alice"); + ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); + ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); + ASSERT_EQ(credentials.session_token, ""); + ASSERT_EQ(credentials.expires_at, absl::InfiniteFuture()); +} + +TEST_F(FileCredentialProviderTest, + ProviderAwsCredentialsFromFileInvalidProfileEnv) { + TestData test_data; + auto credentials_filename = test_data.WriteCredentialsFile(); + + SetEnv("AWS_SHARED_CREDENTIALS_FILE", credentials_filename.c_str()); + SetEnv("AWS_PROFILE", "bob"); + auto provider = FileCredentialProvider("", ""); + ASSERT_FALSE(provider.GetCredentials().ok()); + ASSERT_EQ(provider.GetFileName(), credentials_filename); + ASSERT_EQ(provider.GetProfile(), "bob"); +} + +TEST_F(FileCredentialProviderTest, ProviderAwsCredentialsFromFileOverride) { + TestData test_data; + auto credentials_filename = test_data.WriteCredentialsFile(); + auto provider = + std::make_unique(credentials_filename, ""); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto credentials, + provider->GetCredentials()); + ASSERT_EQ(provider->GetFileName(), credentials_filename); + ASSERT_EQ(provider->GetProfile(), "default"); + ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN7EXAMPLE"); + ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"); + ASSERT_EQ(credentials.session_token, "abcdef1234567890"); + ASSERT_EQ(credentials.expires_at, absl::InfiniteFuture()); + + provider = + std::make_unique(credentials_filename, "alice"); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(credentials, provider->GetCredentials()); + ASSERT_EQ(provider->GetFileName(), credentials_filename); + ASSERT_EQ(provider->GetProfile(), "alice"); + ASSERT_EQ(credentials.access_key, "AKIAIOSFODNN6EXAMPLE"); + ASSERT_EQ(credentials.secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCZEXAMPLEKEY"); + ASSERT_EQ(credentials.session_token, ""); + ASSERT_EQ(credentials.expires_at, absl::InfiniteFuture()); +} + +} // namespace diff --git a/tensorstore/kvstore/s3/credentials/test_utils.cc b/tensorstore/kvstore/s3/credentials/test_utils.cc new file mode 100644 index 000000000..1a84a9c26 --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/test_utils.cc @@ -0,0 +1,83 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tensorstore/kvstore/s3/credentials/test_utils.h" + +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/log/absl_log.h" +#include "absl/strings/cord.h" +#include "absl/strings/str_format.h" +#include "absl/time/time.h" +#include "tensorstore/internal/http/http_request.h" +#include "tensorstore/internal/http/http_response.h" +#include "tensorstore/util/future.h" +#include "tensorstore/util/str_cat.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +Future EC2MetadataMockTransport::IssueRequest( + const internal_http::HttpRequest& request, absl::Cord payload, + absl::Duration request_timeout, absl::Duration connect_timeout) { + ABSL_LOG(INFO) << request; + if (auto it = url_to_response_.find(StrCat(request.method, " ", request.url)); + it != url_to_response_.end()) { + return it->second; + } + + return internal_http::HttpResponse{404, absl::Cord(), {}}; +} + +absl::flat_hash_map +DefaultEC2MetadataFlow(const std::string& endpoint, + const std::string& api_token, + const std::string& access_key, + const std::string& secret_key, + const std::string& session_token, + const absl::Time& expires_at) { + return absl::flat_hash_map{ + {absl::StrFormat("POST %s/latest/api/token", endpoint), + internal_http::HttpResponse{200, absl::Cord{api_token}}}, + {absl::StrFormat("GET %s/latest/meta-data/iam/", endpoint), + internal_http::HttpResponse{ + 200, absl::Cord{"info"}, {{"x-aws-ec2-metadata-token", api_token}}}}, + {absl::StrFormat("GET %s/latest/meta-data/iam/security-credentials/", + endpoint), + internal_http::HttpResponse{200, + absl::Cord{"mock-iam-role"}, + {{"x-aws-ec2-metadata-token", api_token}}}}, + {absl::StrFormat( + "GET %s/latest/meta-data/iam/security-credentials/mock-iam-role", + endpoint), + internal_http::HttpResponse{ + 200, + absl::Cord(absl::StrFormat( + R"({ + "Code": "Success", + "AccessKeyId": "%s", + "SecretAccessKey": "%s", + "Token": "%s", + "Expiration": "%s" + })", + access_key, secret_key, session_token, + absl::FormatTime("%Y-%m-%d%ET%H:%M:%E*S%Ez", expires_at, + absl::UTCTimeZone()))), + {{"x-aws-ec2-metadata-token", api_token}}}}}; +} + +} // namespace internal_kvstore_s3 +} // namespace tensorstore diff --git a/tensorstore/kvstore/s3/credentials/test_utils.h b/tensorstore/kvstore/s3/credentials/test_utils.h new file mode 100644 index 000000000..ceea48a6a --- /dev/null +++ b/tensorstore/kvstore/s3/credentials/test_utils.h @@ -0,0 +1,62 @@ +// Copyright 2023 The TensorStore Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENSORSTORE_KVSTORE_S3_CREDENTIALS_TEST_UTILS_H_ +#define TENSORSTORE_KVSTORE_S3_CREDENTIALS_TEST_UTILS_H_ + +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/strings/cord.h" +#include "absl/time/time.h" +#include "tensorstore/internal/http/http_request.h" +#include "tensorstore/internal/http/http_response.h" +#include "tensorstore/internal/http/http_transport.h" +#include "tensorstore/util/future.h" + +namespace tensorstore { +namespace internal_kvstore_s3 { + +/// Mocks an HttpTransport by overriding the IssueRequest method to +/// respond with a predefined set of request-response pairs supplied +/// to the constructor +class EC2MetadataMockTransport : public internal_http::HttpTransport { + public: + EC2MetadataMockTransport( + const absl::flat_hash_map& + url_to_response) + : url_to_response_(url_to_response) {} + + Future IssueRequest( + const internal_http::HttpRequest& request, absl::Cord payload, + absl::Duration request_timeout, absl::Duration connect_timeout) override; + + const absl::flat_hash_map& + url_to_response_; +}; + +/// Return a Default EC2 Metadata Credential Retrieval Flow, suitable +/// for passing to EC2MetadataMockTransport +absl::flat_hash_map +DefaultEC2MetadataFlow(const std::string& endpoint, + const std::string& api_token, + const std::string& access_key, + const std::string& secret_key, + const std::string& session_token, + const absl::Time& expires_at); + +} // namespace internal_kvstore_s3 +} // namespace tensorstore + +#endif // TENSORSTORE_KVSTORE_S3_CREDENTIALS_TEST_UTILS_H_ diff --git a/tensorstore/kvstore/s3/localstack_test.cc b/tensorstore/kvstore/s3/localstack_test.cc index b6cc868f1..1d5d8869f 100644 --- a/tensorstore/kvstore/s3/localstack_test.cc +++ b/tensorstore/kvstore/s3/localstack_test.cc @@ -12,14 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include -#include -#include #include #include #include +#include "absl/container/flat_hash_map.h" #include "absl/flags/flag.h" #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" @@ -39,8 +39,9 @@ #include "tensorstore/internal/os/subprocess.h" #include "tensorstore/json_serialization_options_base.h" #include "tensorstore/kvstore/kvstore.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" #include "tensorstore/kvstore/s3/s3_request_builder.h" +#include "tensorstore/kvstore/spec.h" #include "tensorstore/kvstore/test_util.h" #include "tensorstore/util/future.h" #include "tensorstore/util/result.h" diff --git a/tensorstore/kvstore/s3/s3_endpoint_test.cc b/tensorstore/kvstore/s3/s3_endpoint_test.cc index 46515c56b..f6c1ba327 100644 --- a/tensorstore/kvstore/s3/s3_endpoint_test.cc +++ b/tensorstore/kvstore/s3/s3_endpoint_test.cc @@ -16,6 +16,7 @@ #include #include +#include #include #include diff --git a/tensorstore/kvstore/s3/s3_key_value_store.cc b/tensorstore/kvstore/s3/s3_key_value_store.cc index bdfde5337..5a9026a47 100644 --- a/tensorstore/kvstore/s3/s3_key_value_store.cc +++ b/tensorstore/kvstore/s3/s3_key_value_store.cc @@ -53,7 +53,6 @@ #include "tensorstore/internal/source_location.h" #include "tensorstore/internal/uri_utils.h" #include "tensorstore/kvstore/byte_range.h" -#include "tensorstore/kvstore/driver.h" #include "tensorstore/kvstore/gcs/validate.h" #include "tensorstore/kvstore/gcs_http/rate_limiter.h" #include "tensorstore/kvstore/generation.h" @@ -61,7 +60,8 @@ #include "tensorstore/kvstore/operations.h" #include "tensorstore/kvstore/read_result.h" #include "tensorstore/kvstore/registry.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" +#include "tensorstore/kvstore/s3/credentials/default_credential_provider.h" #include "tensorstore/kvstore/s3/s3_endpoint.h" #include "tensorstore/kvstore/s3/s3_metadata.h" #include "tensorstore/kvstore/s3/s3_request_builder.h" @@ -207,8 +207,11 @@ struct AwsCredentialsResource struct Spec { std::string profile; + std::string filename; + std::string metadata_endpoint; + constexpr static auto ApplyMembers = [](auto&& x, auto f) { - return f(x.profile); + return f(x.profile, x.filename, x.metadata_endpoint); }; }; @@ -222,16 +225,18 @@ struct AwsCredentialsResource static Spec Default() { return Spec{}; } static constexpr auto JsonBinder() { - return jb::Object( - jb::Member("profile", jb::Projection<&Spec::profile>()) /**/ - ); + return jb::Object(jb::Member("profile", jb::Projection<&Spec::profile>()), + jb::Member("filename", jb::Projection<&Spec::filename>()), + jb::Member("metadata_endpoint", + jb::Projection<&Spec::metadata_endpoint>())); } Result Create( const Spec& spec, internal::ContextResourceCreationContext context) const { auto result = GetAwsCredentialProvider( - spec.profile, internal_http::GetDefaultHttpTransport()); + spec.profile, spec.filename, spec.metadata_endpoint, + internal_http::GetDefaultHttpTransport()); if (!result.ok() && absl::IsNotFound(result.status())) { return Resource{spec, nullptr}; } diff --git a/tensorstore/kvstore/s3/s3_key_value_store_test.cc b/tensorstore/kvstore/s3/s3_key_value_store_test.cc index a3b97f4d4..25a6249e9 100644 --- a/tensorstore/kvstore/s3/s3_key_value_store_test.cc +++ b/tensorstore/kvstore/s3/s3_key_value_store_test.cc @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include @@ -23,7 +23,6 @@ #include "absl/status/status.h" #include "absl/strings/cord.h" #include "absl/time/time.h" -#include #include "tensorstore/context.h" #include "tensorstore/internal/http/curl_transport.h" #include "tensorstore/internal/http/http_request.h" @@ -37,6 +36,7 @@ #include "tensorstore/kvstore/read_result_testutil.h" #include "tensorstore/kvstore/test_util.h" #include "tensorstore/util/future.h" +#include "tensorstore/util/result.h" #include "tensorstore/util/status_testutil.h" #include "tensorstore/util/str_cat.h" diff --git a/tensorstore/kvstore/s3/s3_metadata.cc b/tensorstore/kvstore/s3/s3_metadata.cc index 549200a37..e762da29e 100644 --- a/tensorstore/kvstore/s3/s3_metadata.cc +++ b/tensorstore/kvstore/s3/s3_metadata.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include "absl/status/status.h" #include "absl/strings/str_cat.h" diff --git a/tensorstore/kvstore/s3/s3_metadata_test.cc b/tensorstore/kvstore/s3/s3_metadata_test.cc index 4fb4f782e..9de7fec12 100644 --- a/tensorstore/kvstore/s3/s3_metadata_test.cc +++ b/tensorstore/kvstore/s3/s3_metadata_test.cc @@ -14,8 +14,6 @@ #include "tensorstore/kvstore/s3/s3_metadata.h" -#include - #include #include #include diff --git a/tensorstore/kvstore/s3/s3_request_builder.cc b/tensorstore/kvstore/s3/s3_request_builder.cc index 9bb9f1f5b..442ddcd5d 100644 --- a/tensorstore/kvstore/s3/s3_request_builder.cc +++ b/tensorstore/kvstore/s3/s3_request_builder.cc @@ -37,13 +37,14 @@ #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/time/time.h" +#include #include // IWYU pragma: keep #include #include "tensorstore/internal/digest/sha256.h" #include "tensorstore/internal/http/http_request.h" #include "tensorstore/internal/log/verbose_flag.h" #include "tensorstore/internal/uri_utils.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" #include "tensorstore/kvstore/s3/s3_uri_utils.h" using ::tensorstore::internal::ParseGenericUri; diff --git a/tensorstore/kvstore/s3/s3_request_builder.h b/tensorstore/kvstore/s3/s3_request_builder.h index 9247e7b4b..89548fd3f 100644 --- a/tensorstore/kvstore/s3/s3_request_builder.h +++ b/tensorstore/kvstore/s3/s3_request_builder.h @@ -26,7 +26,7 @@ #include "absl/time/time.h" #include "tensorstore/internal/http/http_request.h" #include "tensorstore/kvstore/byte_range.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" #include "tensorstore/kvstore/s3/s3_uri_utils.h" namespace tensorstore { diff --git a/tensorstore/kvstore/s3/s3_request_builder_test.cc b/tensorstore/kvstore/s3/s3_request_builder_test.cc index 56565556b..a7905f6e9 100644 --- a/tensorstore/kvstore/s3/s3_request_builder_test.cc +++ b/tensorstore/kvstore/s3/s3_request_builder_test.cc @@ -15,7 +15,6 @@ #include "tensorstore/kvstore/s3/s3_request_builder.h" #include -#include #include #include @@ -25,7 +24,7 @@ #include "absl/time/civil_time.h" #include "absl/time/time.h" #include "tensorstore/internal/http/http_request.h" -#include "tensorstore/kvstore/s3/aws_credential_provider.h" +#include "tensorstore/kvstore/s3/credentials/aws_credentials.h" using ::tensorstore::internal_kvstore_s3::AwsCredentials; using ::tensorstore::internal_kvstore_s3::S3RequestBuilder; diff --git a/tensorstore/kvstore/s3/s3_resource.cc b/tensorstore/kvstore/s3/s3_resource.cc index 637cbbf2b..d171c6298 100644 --- a/tensorstore/kvstore/s3/s3_resource.cc +++ b/tensorstore/kvstore/s3/s3_resource.cc @@ -14,7 +14,8 @@ #include "tensorstore/kvstore/s3/s3_resource.h" -#include +#include + #include #include #include diff --git a/tensorstore/kvstore/s3/schema.yml b/tensorstore/kvstore/s3/schema.yml index 9a7bb3dc5..aaa02ab50 100644 --- a/tensorstore/kvstore/s3/schema.yml +++ b/tensorstore/kvstore/s3/schema.yml @@ -116,6 +116,16 @@ definitions: description: |- The profile name in the :file:`~/.aws/credentials` file, when used. Overrides the :envvar:`AWS_PROFILE` environment variables. + filename: + type: string + description: |- + The filename containing credentials. + Overrides the :envvar:`AWS_SHARED_CREDENTIALS_FILE` environment variable. + metadata_endpoint: + type: string + description: |- + The endpoint of the metadata server. + Overrides the :envvar:`AWS_EC2_METADATA_SERVICE_ENDPOINT` environment variable. experimental_s3_rate_limiter: $id: Context.experimental_s3_rate_limiter description: |-