From 51073607eaef1fabd486609dbad29683d15e7b0f Mon Sep 17 00:00:00 2001
From: Dean De Leo <hello@whatsthecraic.net>
Date: Thu, 13 Mar 2025 20:02:09 +0700
Subject: [PATCH] S3: opt-in the STSProfileCredentialsProvider

The STSProfileCredentialsProviders allows to assume a specific IAM role
when accessing an S3 repository. Sometimes this is needed to obtain the
permissions to operate on the bucket.
---
 packaging/dependencies.nix            |  1 +
 src/libstore/meson.build              |  3 +++
 src/libstore/s3-binary-cache-store.cc | 31 +++++++++++++++++++++------
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/packaging/dependencies.nix b/packaging/dependencies.nix
index ed05843c7..16d1f1376 100644
--- a/packaging/dependencies.nix
+++ b/packaging/dependencies.nix
@@ -38,6 +38,7 @@ scope: {
   aws-sdk-cpp =
     (pkgs.aws-sdk-cpp.override {
       apis = [
+        "identity-management"
         "s3"
         "transfer"
       ];
diff --git a/src/libstore/meson.build b/src/libstore/meson.build
index ae50686b6..255f83f74 100644
--- a/src/libstore/meson.build
+++ b/src/libstore/meson.build
@@ -135,6 +135,9 @@ if aws_s3.found()
       '-L' + aws_s3.get_variable('libdir'),
       '-laws-cpp-sdk-transfer',
       '-laws-cpp-sdk-s3',
+      '-laws-cpp-sdk-identity-management',
+      '-laws-cpp-sdk-cognito-identity',
+      '-laws-cpp-sdk-sts',
       '-laws-cpp-sdk-core',
       '-laws-crt-cpp',
     ],
diff --git a/src/libstore/s3-binary-cache-store.cc b/src/libstore/s3-binary-cache-store.cc
index f9e583307..26c21de44 100644
--- a/src/libstore/s3-binary-cache-store.cc
+++ b/src/libstore/s3-binary-cache-store.cc
@@ -21,6 +21,7 @@
 #include <aws/core/utils/logging/FormattedLogSystem.h>
 #include <aws/core/utils/logging/LogMacros.h>
 #include <aws/core/utils/threading/Executor.h>
+#include <aws/identity-management/auth/STSProfileCredentialsProvider.h>
 #include <aws/s3/S3Client.h>
 #include <aws/s3/model/GetObjectRequest.h>
 #include <aws/s3/model/HeadObjectRequest.h>
@@ -71,6 +72,29 @@ class AwsLogger : public Aws::Utils::Logging::FormattedLogSystem
 #endif
 };
 
+/* Retrieve the credentials from the list of AWS default providers, with the addition of the STS creds provider. This
+   last can be used to acquire further permissions with a specific IAM role.
+   Roughly based on https://github.com/aws/aws-sdk-cpp/issues/150#issuecomment-538548438
+*/
+struct CustomAwsCredentialsProviderChain : public Aws::Auth::AWSCredentialsProviderChain
+{
+    CustomAwsCredentialsProviderChain(const std::string & profile)
+    {
+        if (profile.empty()) {
+            // Use all the default AWS providers, plus the possibility to acquire a IAM role directly via a profile.
+            Aws::Auth::DefaultAWSCredentialsProviderChain default_aws_chain;
+            for (auto provider : default_aws_chain.GetProviders())
+                AddProvider(provider);
+            AddProvider(std::make_shared<Aws::Auth::STSProfileCredentialsProvider>());
+        } else {
+            // Override the profile name to retrieve from the AWS config and credentials. I believe this option
+            // comes from the ?profile querystring in nix.conf.
+            AddProvider(std::make_shared<Aws::Auth::ProfileConfigFileAWSCredentialsProvider>(profile.c_str()));
+            AddProvider(std::make_shared<Aws::Auth::STSProfileCredentialsProvider>(profile));
+        }
+    }
+};
+
 static void initAWS()
 {
     static std::once_flag flag;
@@ -102,13 +126,8 @@ S3Helper::S3Helper(
     const std::string & endpoint)
     : config(makeConfig(region, scheme, endpoint))
     , client(make_ref<Aws::S3::S3Client>(
-            profile == ""
-            ? std::dynamic_pointer_cast<Aws::Auth::AWSCredentialsProvider>(
-                std::make_shared<Aws::Auth::DefaultAWSCredentialsProviderChain>())
-            : std::dynamic_pointer_cast<Aws::Auth::AWSCredentialsProvider>(
-                std::make_shared<Aws::Auth::ProfileConfigFileAWSCredentialsProvider>(profile.c_str())),
+            std::make_shared<CustomAwsCredentialsProviderChain>(profile),
             *config,
-            // FIXME: https://github.com/aws/aws-sdk-cpp/issues/759
 #if AWS_SDK_VERSION_MAJOR == 1 && AWS_SDK_VERSION_MINOR < 3
             false,
 #else