From eda644cf0d56584bb358ab4b400b924f97fe30e0 Mon Sep 17 00:00:00 2001
From: nicolaferraro <ni.ferraro@gmail.com>
Date: Wed, 13 Nov 2024 18:05:16 +0100
Subject: [PATCH] secrets: allow AWS client to assume a different role

---
 secrets/aws.go | 28 ++++++++++++++++++++++++----
 secrets/go.mod |  6 +++---
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/secrets/aws.go b/secrets/aws.go
index 0edd3e4..ef825a4 100644
--- a/secrets/aws.go
+++ b/secrets/aws.go
@@ -6,9 +6,12 @@ import (
 	"fmt"
 	"log/slog"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
 	"github.com/aws/aws-sdk-go-v2/service/secretsmanager"
 	"github.com/aws/aws-sdk-go-v2/service/secretsmanager/types"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
 )
 
 type awsSecretsManager struct {
@@ -16,14 +19,14 @@ type awsSecretsManager struct {
 	logger *slog.Logger
 }
 
-func NewAWSSecretsManager(ctx context.Context, logger *slog.Logger, region string) (SecretAPI, error) {
-	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+func NewAWSSecretsManager(ctx context.Context, logger *slog.Logger, region string, roleARN string) (SecretAPI, error) {
+	cl, err := createAWSClient(ctx, region, roleARN)
 	if err != nil {
-		return nil, fmt.Errorf("failed to load AWS config: %w", err)
+		return nil, fmt.Errorf("failed to create secrets manager client: %w", err)
 	}
 
 	return &awsSecretsManager{
-		client: secretsmanager.NewFromConfig(cfg),
+		client: cl,
 		logger: logger,
 	}, nil
 }
@@ -49,3 +52,20 @@ func (a *awsSecretsManager) CheckSecretExists(ctx context.Context, key string) b
 	})
 	return err == nil
 }
+
+func createAWSClient(ctx context.Context, region string, roleARN string) (*secretsmanager.Client, error) {
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+	if err != nil {
+		return nil, fmt.Errorf("failed to load AWS config: %w", err)
+	}
+	if roleARN == "" {
+		return secretsmanager.NewFromConfig(cfg), nil
+	}
+
+	creds := aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(sts.NewFromConfig(cfg), roleARN))
+	secretsManagerClient := secretsmanager.New(secretsmanager.Options{
+		Credentials: creds,
+		Region:      region,
+	})
+	return secretsManagerClient, nil
+}
diff --git a/secrets/go.mod b/secrets/go.mod
index e955274..91619cf 100644
--- a/secrets/go.mod
+++ b/secrets/go.mod
@@ -6,8 +6,11 @@ require (
 	cloud.google.com/go/secretmanager v1.14.2
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.0
 	github.com/Azure/azure-sdk-for-go/sdk/keyvault/azsecrets v0.12.0
+	github.com/aws/aws-sdk-go-v2 v1.32.3
 	github.com/aws/aws-sdk-go-v2/config v1.28.1
+	github.com/aws/aws-sdk-go-v2/credentials v1.17.42
 	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.34.3
+	github.com/aws/aws-sdk-go-v2/service/sts v1.32.3
 	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/gjson v1.18.0
 	google.golang.org/grpc v1.67.1
@@ -22,8 +25,6 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/keyvault/internal v0.7.1 // indirect
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 // indirect
-	github.com/aws/aws-sdk-go-v2 v1.32.3 // indirect
-	github.com/aws/aws-sdk-go-v2/credentials v1.17.42 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.18 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.22 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.22 // indirect
@@ -32,7 +33,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.3 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.24.3 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.32.3 // indirect
 	github.com/aws/smithy-go v1.22.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect