From 9ce63f12a7067515e9c9b2c56ee8e4530c210c37 Mon Sep 17 00:00:00 2001
From: "OP (oppenheimer)" <21008429+Ompragash@users.noreply.github.com>
Date: Mon, 13 May 2024 13:19:42 +0530
Subject: [PATCH] Add PLUGIN_OIDC_TOKEN_ID support to AWS S3 (#103)

* Add PLUGIN_OIDC_TOKEN_ID support to AWS S3

* Updated storage/backend/s3/s3.go

* Updated storage/backend/s3/s3.go

* Updated storage/backend/s3/s3.go

* Improved logging on S3 assumeRoleWithWebIdentity method

* Dummy commit
---
 main.go                      | 31 ++++++++++++++++----------
 storage/backend/s3/config.go |  1 +
 storage/backend/s3/s3.go     | 42 +++++++++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/main.go b/main.go
index 09952a29..5a722a34 100644
--- a/main.go
+++ b/main.go
@@ -357,6 +357,13 @@ func main() {
 			EnvVars: []string{"PLUGIN_FILESYSTEM_CACHE_ROOT", "FILESYSTEM_CACHE_ROOT"},
 		},
 
+		// OIDC
+		&cli.StringFlag{
+			Name:    "oidc-token-id",
+			Usage:   "OIDC token ID for assuming role with web identity",
+			EnvVars: []string{"PLUGIN_OIDC_TOKEN_ID"},
+		},
+
 		// S3 specific Config flags
 
 		&cli.StringFlag{
@@ -591,17 +598,19 @@ func run(c *cli.Context) error {
 			CacheRoot: c.String("filesystem.cache-root"),
 		},
 		S3: s3.Config{
-			ACL:           c.String("acl"),
-			Bucket:        c.String("bucket"),
-			Encryption:    c.String("encryption"),
-			Endpoint:      c.String("endpoint"),
-			Key:           c.String("access-key"),
-			PathStyle:     c.Bool("path-style"),
-			Region:        c.String("region"),
-			Secret:        c.String("secret-key"),
-			StsEndpoint:   c.String("sts-endpoint"),
-			AssumeRoleARN: c.String("assume-role-arn"),
-			UserRoleArn:   c.String("user-role-arn"),
+			ACL:                   c.String("acl"),
+			Bucket:                c.String("bucket"),
+			Encryption:            c.String("encryption"),
+			Endpoint:              c.String("endpoint"),
+			Key:                   c.String("access-key"),
+			PathStyle:             c.Bool("path-style"),
+			Region:                c.String("region"),
+			Secret:                c.String("secret-key"),
+			StsEndpoint:           c.String("sts-endpoint"),
+			AssumeRoleARN:         c.String("assume-role-arn"),
+			AssumeRoleSessionName: c.String("assume-role-session-name"),
+			UserRoleArn:           c.String("user-role-arn"),
+			OIDCTokenID:           c.String("oidc-token-id"),
 		},
 		Azure: azure.Config{
 			AccountName:    c.String("azure.account-name"),
diff --git a/storage/backend/s3/config.go b/storage/backend/s3/config.go
index d338e2ce..0570a99c 100644
--- a/storage/backend/s3/config.go
+++ b/storage/backend/s3/config.go
@@ -19,6 +19,7 @@ type Config struct {
 	AssumeRoleARN         string // if "", do not assume IAM role i.e. use the IAM user.
 	AssumeRoleSessionName string
 	UserRoleArn           string
+	OIDCTokenID           string
 
 	// us-east-1
 	// us-west-1
diff --git a/storage/backend/s3/s3.go b/storage/backend/s3/s3.go
index d39d43fc..c086b9f4 100644
--- a/storage/backend/s3/s3.go
+++ b/storage/backend/s3/s3.go
@@ -44,7 +44,17 @@ func New(l log.Logger, c Config, debug bool) (*Backend, error) {
 	if c.Key != "" && c.Secret != "" { // nolint:gocritic
 		conf.Credentials = credentials.NewStaticCredentials(c.Key, c.Secret, "")
 	} else if c.AssumeRoleARN != "" {
-		conf.Credentials = assumeRole(c.AssumeRoleARN, c.AssumeRoleSessionName)
+		if c.OIDCTokenID != "" {
+			// Assume role with OIDC
+			creds, err := assumeRoleWithWebIdentity(c.AssumeRoleARN, c.AssumeRoleSessionName, c.OIDCTokenID)
+			if err != nil {
+				level.Error(l).Log("msg", "failed to assume role with OIDC", "error", err)
+				return nil, err
+			}
+			conf.Credentials = creds
+		} else {
+			conf.Credentials = assumeRole(c.AssumeRoleARN, c.AssumeRoleSessionName)
+		}
 	} else {
 		level.Warn(l).Log("msg", "aws key and/or Secret not provided (falling back to anonymous credentials)")
 	}
@@ -195,3 +205,33 @@ func assumeRole(roleArn, roleSessionName string) *credentials.Credentials {
 
 	return credentials.NewCredentials(stsProvider)
 }
+
+func assumeRoleWithWebIdentity(roleArn, roleSessionName, webIdentityToken string) (*credentials.Credentials, error) {
+	// Create a new session
+	sess, err := session.NewSession()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create AWS session: %v", err)
+	}
+
+	// Create a new STS client
+	svc := sts.New(sess)
+
+	// Prepare the input parameters for the STS call
+	duration := int64(time.Hour / time.Second)
+	input := &sts.AssumeRoleWithWebIdentityInput{
+		RoleArn:          aws.String(roleArn),
+		RoleSessionName:  aws.String(roleSessionName),
+		WebIdentityToken: aws.String(webIdentityToken),
+		DurationSeconds:  aws.Int64(duration),
+	}
+
+	// Call the AssumeRoleWithWebIdentity function
+	result, err := svc.AssumeRoleWithWebIdentity(input)
+	if err != nil {
+		return nil, fmt.Errorf("failed to assume role with web identity: %v", err)
+	}
+
+	// Create credentials using the response from STS
+	newCreds := credentials.NewStaticCredentials(*result.Credentials.AccessKeyId, *result.Credentials.SecretAccessKey, *result.Credentials.SessionToken)
+	return newCreds, nil
+}