support assume role (bug fix included) (ClickHouse#449)

binmahone · kyligence-git · commit 7d41c19d0fd2 · 2023-12-27T23:02:45.000Z
(cherry picked from commit e7061da)
diff --git a/contrib/aws-cmake/CMakeLists.txt b/contrib/aws-cmake/CMakeLists.txt
@@ -53,6 +53,9 @@ endif()
 # Directories.
 SET(AWS_SDK_DIR "${ClickHouse_SOURCE_DIR}/contrib/aws")
 SET(AWS_SDK_CORE_DIR "${AWS_SDK_DIR}/src/aws-cpp-sdk-core")
+SET(AWS_SDK_IM_DIR "${AWS_SDK_DIR}/src/aws-cpp-sdk-identity-management")
+SET(AWS_SDK_STS_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-sts")
+SET(AWS_SDK_COGNITO_IDENTITY_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-cognito-identity")
 SET(AWS_SDK_S3_DIR "${AWS_SDK_DIR}/generated/src/aws-cpp-sdk-s3")
 
 SET(AWS_AUTH_DIR "${ClickHouse_SOURCE_DIR}/contrib/aws-c-auth")
@@ -103,6 +106,9 @@ file(GLOB AWS_SDK_CORE_SRC
     "${AWS_SDK_CORE_DIR}/source/utils/xml/*.cpp"
 )
 
+file(GLOB AWS_SDK_IM_SRC
+        "${AWS_SDK_IM_DIR}/source/auth/*.cpp")
+
 if(OS_LINUX OR OS_DARWIN)
     file(GLOB AWS_SDK_CORE_NET_SRC "${AWS_SDK_CORE_DIR}/source/net/linux-shared/*.cpp")
     file(GLOB AWS_SDK_CORE_PLATFORM_SRC "${AWS_SDK_CORE_DIR}/source/platform/linux-shared/*.cpp")
@@ -119,13 +125,23 @@ list(APPEND AWS_PUBLIC_COMPILE_DEFS "-DAWS_SDK_VERSION_MAJOR=1")
 list(APPEND AWS_PUBLIC_COMPILE_DEFS "-DAWS_SDK_VERSION_MINOR=10")
 list(APPEND AWS_PUBLIC_COMPILE_DEFS "-DAWS_SDK_VERSION_PATCH=36")
 
-list(APPEND AWS_SOURCES ${AWS_SDK_CORE_SRC} ${AWS_SDK_CORE_NET_SRC} ${AWS_SDK_CORE_PLATFORM_SRC})
+list(APPEND AWS_SOURCES ${AWS_SDK_CORE_SRC} ${AWS_SDK_IM_SRC} ${AWS_SDK_CORE_NET_SRC} ${AWS_SDK_CORE_PLATFORM_SRC})
 
 list(APPEND AWS_PUBLIC_INCLUDES
     "${AWS_SDK_CORE_DIR}/include/"
+    "${AWS_SDK_IM_DIR}/include/"
     "${CMAKE_CURRENT_BINARY_DIR}/include"
 )
 
+# aws-cpp-sdk-sts
+file(GLOB AWS_SDK_STS_SRC
+        "${AWS_SDK_STS_DIR}/source/*.cpp"
+        "${AWS_SDK_STS_DIR}/source/model/*.cpp"
+        )
+
+list(APPEND AWS_SOURCES ${AWS_SDK_STS_SRC})
+list(APPEND AWS_PUBLIC_INCLUDES "${AWS_SDK_STS_DIR}/include/")
+list(APPEND AWS_PUBLIC_INCLUDES "${AWS_SDK_COGNITO_IDENTITY_DIR}/include/")
 
 # aws-cpp-sdk-s3
 file(GLOB AWS_SDK_S3_SRC
diff --git a/src/IO/S3/Client.cpp b/src/IO/S3/Client.cpp
@@ -12,6 +12,8 @@
 #include <aws/core/endpoint/EndpointParameter.h>
 #include <aws/core/utils/HashingUtils.h>
 #include <aws/core/utils/logging/ErrorMacros.h>
+#include <aws/sts/STSClient.h>
+#include <aws/identity-management/auth/STSAssumeRoleCredentialsProvider.h>
 
 #include <Poco/Net/NetException.h>
 
@@ -885,11 +887,25 @@ std::unique_ptr<S3::Client> ClientFactory::create( // NOLINT
     client_configuration.extra_headers = std::move(headers);
 
     Aws::Auth::AWSCredentials credentials(access_key_id, secret_access_key, session_token);
-    auto credentials_provider = std::make_shared<S3CredentialsProviderChain>(
+    std::shared_ptr<Aws::Auth::AWSCredentialsProvider> credentials_provider = std::make_shared<S3CredentialsProviderChain>(
             client_configuration,
             std::move(credentials),
             credentials_configuration);
 
+    if (!credentials_configuration.role_arn.empty())
+    {
+        // why set to empty? because client_configuration's endpointOverride is pointed to a s3 endpoint, whereas we
+        // expect are going to visit a sts endpoint.
+        client_configuration.endpointOverride = "";
+        const auto x = std::make_shared<Aws::STS::STSClient>(credentials_provider, client_configuration);
+        credentials_provider = std::make_shared<Aws::Auth::STSAssumeRoleCredentialsProvider>(
+            credentials_configuration.role_arn,
+            credentials_configuration.session_name,
+            credentials_configuration.external_id,
+            Aws::Auth::DEFAULT_CREDS_LOAD_FREQ_SECONDS,
+            x);
+    }
+
     client_configuration.retryStrategy = std::make_shared<Client::RetryStrategy>(client_configuration.s3_retry_attempts);
 
     return Client::create(
diff --git a/src/IO/S3/Credentials.h b/src/IO/S3/Credentials.h
@@ -176,6 +176,11 @@ struct CredentialsConfiguration
     bool use_insecure_imds_request = false;
     uint64_t expiration_window_seconds = DEFAULT_EXPIRATION_WINDOW_SECONDS;
     bool no_sign_request = false;
+
+    // STS Assume Role related
+    std::string role_arn = "";
+    std::string session_name = "";
+    std::string external_id = "";
 };
 
 class S3CredentialsProviderChain : public Aws::Auth::AWSCredentialsProviderChain
