Merge pull request #54347 from ClickHouse/add-sso-for-s3

Add support for SSO credentials in S3

src/IO/S3/Credentials.cpp

@@ -11,6 +11,9 @@
 #    include <aws/core/utils/UUID.h>
 #    include <aws/core/http/HttpClientFactory.h>
 
+#    include <aws/core/utils/HashingUtils.h>
+#    include <aws/core/platform/FileSystem.h>
+
 #    include <Common/logger_useful.h>
 
 #    include <IO/S3/PocoHTTPClient.h>
@@ -43,6 +46,8 @@ bool areCredentialsEmptyOrExpired(const Aws::Auth::AWSCredentials & credentials,
     return now >= credentials.GetExpiration() - std::chrono::seconds(expiration_window_seconds);
 }
 
+const char SSO_CREDENTIALS_PROVIDER_LOG_TAG[] = "SSOCredentialsProvider";
+
 }
 
 AWSEC2MetadataClient::AWSEC2MetadataClient(const Aws::Client::ClientConfiguration & client_configuration, const char * endpoint_)
@@ -449,6 +454,139 @@ void AwsAuthSTSAssumeRoleWebIdentityCredentialsProvider::refreshIfExpired()
     Reload();
 }
 
+
+SSOCredentialsProvider::SSOCredentialsProvider(DB::S3::PocoHTTPClientConfiguration aws_client_configuration_, uint64_t expiration_window_seconds_)
+    : profile_to_use(Aws::Auth::GetConfigProfileName())
+    , aws_client_configuration(std::move(aws_client_configuration_))
+    , expiration_window_seconds(expiration_window_seconds_)
+    , logger(&Poco::Logger::get(SSO_CREDENTIALS_PROVIDER_LOG_TAG))
+{
+    LOG_INFO(logger, "Setting sso credentials provider to read config from {}", profile_to_use);
+}
+
+Aws::Auth::AWSCredentials SSOCredentialsProvider::GetAWSCredentials()
+{
+    refreshIfExpired();
+    Aws::Utils::Threading::ReaderLockGuard guard(m_reloadLock);
+    return credentials;
+}
+
+void SSOCredentialsProvider::Reload()
+{
+    auto profile = Aws::Config::GetCachedConfigProfile(profile_to_use);
+    const auto access_token = [&]
+    {
+        // If we have an SSO Session set, use the refreshed token.
+        if (profile.IsSsoSessionSet())
+        {
+            sso_region = profile.GetSsoSession().GetSsoRegion();
+            auto token = bearer_token_provider.GetAWSBearerToken();
+            expires_at = token.GetExpiration();
+            return token.GetToken();
+        }
+
+        Aws::String hashed_start_url = Aws::Utils::HashingUtils::HexEncode(Aws::Utils::HashingUtils::CalculateSHA1(profile.GetSsoStartUrl()));
+        auto profile_directory = Aws::Auth::ProfileConfigFileAWSCredentialsProvider::GetProfileDirectory();
+        Aws::StringStream ss_token;
+        ss_token << profile_directory;
+        ss_token << Aws::FileSystem::PATH_DELIM << "sso"  << Aws::FileSystem::PATH_DELIM << "cache" << Aws::FileSystem::PATH_DELIM << hashed_start_url << ".json";
+        auto sso_token_path = ss_token.str();
+        LOG_INFO(logger, "Loading token from: {}", sso_token_path);
+        sso_region = profile.GetSsoRegion();
+        return loadAccessTokenFile(sso_token_path);
+    }();
+
+    if (access_token.empty())
+    {
+        LOG_TRACE(logger, "Access token for SSO not available");
+        return;
+    }
+    if (expires_at < Aws::Utils::DateTime::Now())
+    {
+        LOG_TRACE(logger, "Cached Token expired at {}", expires_at.ToGmtString(Aws::Utils::DateFormat::ISO_8601));
+        return;
+    }
+
+    Aws::Internal::SSOCredentialsClient::SSOGetRoleCredentialsRequest request;
+    request.m_ssoAccountId = profile.GetSsoAccountId();
+    request.m_ssoRoleName = profile.GetSsoRoleName();
+    request.m_accessToken = access_token;
+
+    aws_client_configuration.scheme = Aws::Http::Scheme::HTTPS;
+    aws_client_configuration.region = sso_region;
+    LOG_TRACE(logger, "Passing config to client for region: {}", sso_region);
+
+    Aws::Vector<Aws::String> retryable_errors;
+    retryable_errors.push_back("TooManyRequestsException");
+
+    aws_client_configuration.retryStrategy = Aws::MakeShared<Aws::Client::SpecifiedRetryableErrorsRetryStrategy>(SSO_CREDENTIALS_PROVIDER_LOG_TAG, retryable_errors, /*maxRetries=*/3);
+    client = Aws::MakeUnique<Aws::Internal::SSOCredentialsClient>(SSO_CREDENTIALS_PROVIDER_LOG_TAG, aws_client_configuration);
+
+    LOG_TRACE(logger, "Requesting credentials with AWS_ACCESS_KEY: {}", sso_account_id);
+    auto result = client->GetSSOCredentials(request);
+    LOG_TRACE(logger, "Successfully retrieved credentials with AWS_ACCESS_KEY: {}", result.creds.GetAWSAccessKeyId());
+
+    credentials = result.creds;
+}
+
+void SSOCredentialsProvider::refreshIfExpired()
+{
+    Aws::Utils::Threading::ReaderLockGuard guard(m_reloadLock);
+    if (!areCredentialsEmptyOrExpired(credentials, expiration_window_seconds))
+        return;
+
+    guard.UpgradeToWriterLock();
+
+    if (!areCredentialsEmptyOrExpired(credentials, expiration_window_seconds)) // double-checked lock to avoid refreshing twice
+        return;
+
+    Reload();
+}
+
+Aws::String SSOCredentialsProvider::loadAccessTokenFile(const Aws::String & sso_access_token_path)
+{
+    LOG_TRACE(logger, "Preparing to load token from: {}", sso_access_token_path);
+
+    Aws::IFStream input_file(sso_access_token_path.c_str());
+
+    if (input_file)
+    {
+        LOG_TRACE(logger, "Reading content from token file: {}", sso_access_token_path);
+
+        Aws::Utils::Json::JsonValue token_doc(input_file);
+        if (!token_doc.WasParseSuccessful())
+        {
+            LOG_TRACE(logger, "Failed to parse token file: {}", sso_access_token_path);
+            return "";
+        }
+        Aws::Utils::Json::JsonView token_view(token_doc);
+        Aws::String tmp_access_token, expiration_str;
+        tmp_access_token = token_view.GetString("accessToken");
+        expiration_str = token_view.GetString("expiresAt");
+        Aws::Utils::DateTime expiration(expiration_str, Aws::Utils::DateFormat::ISO_8601);
+
+        LOG_TRACE(logger, "Token cache file contains accessToken [{}], expiration [{}]", tmp_access_token, expiration_str);
+
+        if (tmp_access_token.empty() || !expiration.WasParseSuccessful())
+        {
+            LOG_TRACE(logger, R"(The SSO session associated with this profile has expired or is otherwise invalid. To refresh this SSO session run aws sso login with the corresponding profile.)");
+            LOG_TRACE(
+                logger,
+                "Token cache file failed because {}{}",
+                (tmp_access_token.empty() ? "AccessToken was empty " : ""),
+                (!expiration.WasParseSuccessful() ? "failed to parse expiration" : ""));
+            return "";
+        }
+        expires_at = expiration;
+        return tmp_access_token;
+    }
+    else
+    {
+        LOG_TRACE(logger, "Unable to open token file on path: {}", sso_access_token_path);
+        return "";
+    }
+}
+
 S3CredentialsProviderChain::S3CredentialsProviderChain(
         const DB::S3::PocoHTTPClientConfiguration & configuration,
         const Aws::Auth::AWSCredentials & credentials,
@@ -494,6 +632,18 @@ S3CredentialsProviderChain::S3CredentialsProviderChain(
 
         AddProvider(std::make_shared<Aws::Auth::EnvironmentAWSCredentialsProvider>());
 
+        {
+            DB::S3::PocoHTTPClientConfiguration aws_client_configuration = DB::S3::ClientFactory::instance().createClientConfiguration(
+                configuration.region,
+                configuration.remote_host_filter,
+                configuration.s3_max_redirects,
+                configuration.enable_s3_requests_logging,
+                configuration.for_disk_s3,
+                configuration.get_request_throttler,
+                configuration.put_request_throttler);
+            AddProvider(std::make_shared<SSOCredentialsProvider>(
+                std::move(aws_client_configuration), credentials_configuration.expiration_window_seconds));
+        }
 
         /// ECS TaskRole Credentials only available when ENVIRONMENT VARIABLE is set.
         const auto relative_uri = Aws::Environment::GetEnv(AWS_ECS_CONTAINER_CREDENTIALS_RELATIVE_URI);

src/IO/S3/Credentials.h

@@ -8,6 +8,7 @@
 #    include <aws/core/internal/AWSHttpResourceClient.h>
 #    include <aws/core/config/AWSProfileConfigLoader.h>
 #    include <aws/core/auth/AWSCredentialsProviderChain.h>
+#    include <aws/core/auth/bearer-token-provider/SSOBearerTokenProvider.h>
 
 #    include <IO/S3/PocoHTTPClient.h>
 
@@ -124,6 +125,39 @@ private:
     uint64_t expiration_window_seconds;
 };
 
+class SSOCredentialsProvider : public Aws::Auth::AWSCredentialsProvider
+{
+public:
+    SSOCredentialsProvider(DB::S3::PocoHTTPClientConfiguration aws_client_configuration_, uint64_t expiration_window_seconds_);
+
+    Aws::Auth::AWSCredentials GetAWSCredentials() override;
+
+private:
+    Aws::UniquePtr<Aws::Internal::SSOCredentialsClient> client;
+    Aws::Auth::AWSCredentials credentials;
+
+    // Profile description variables
+    Aws::String profile_to_use;
+
+    // The AWS account ID that temporary AWS credentials are resolved for.
+    Aws::String sso_account_id;
+    // The AWS region where the SSO directory for the given sso_start_url is hosted.
+    // This is independent of the general region configuration and MUST NOT be conflated.
+    Aws::String sso_region;
+    // The expiration time of the accessToken.
+    Aws::Utils::DateTime expires_at;
+    // The SSO Token Provider
+    Aws::Auth::SSOBearerTokenProvider bearer_token_provider;
+
+    DB::S3::PocoHTTPClientConfiguration aws_client_configuration;
+    uint64_t expiration_window_seconds;
+    Poco::Logger * logger;
+
+    void Reload() override;
+    void refreshIfExpired();
+    Aws::String loadAccessTokenFile(const Aws::String & sso_access_token_path);
+};
+
 struct CredentialsConfiguration
 {
     bool use_environment_credentials = false;