feat: Add support for session tokens to S3 client

Signed-off-by: Raymond Chow <rchow@atlassian.com>
argoproj · Jan 3, 2024 · 50a23dd · 50a23dd
1 parent 4dd35c0
commit 50a23dd
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/s3/s3.go b/s3/s3.go
@@ -84,6 +84,7 @@ type S3ClientOpts struct {
 	Transport       http.RoundTripper
 	AccessKey       string
 	SecretKey       string
+	SessionToken    string
 	Trace           bool
 	RoleARN         string
 	RoleSessionName string
@@ -136,8 +137,13 @@ func GetAssumeRoleCredentials(opts S3ClientOpts) (*credentials.Credentials, erro
 
 func GetCredentials(opts S3ClientOpts) (*credentials.Credentials, error) {
 	if opts.AccessKey != "" && opts.SecretKey != "" {
-		log.WithField("endpoint", opts.Endpoint).Info("Creating minio client using static credentials")
-		return credentials.NewStaticV4(opts.AccessKey, opts.SecretKey, ""), nil
+		if opts.SessionToken != "" {
+			log.WithField("endpoint", opts.Endpoint).Info("Creating minio client using ephemeral credentials")
+			return credentials.NewStaticV4(opts.AccessKey, opts.SecretKey, opts.SessionToken), nil
+		} else {
+			log.WithField("endpoint", opts.Endpoint).Info("Creating minio client using static credentials")
+			return credentials.NewStaticV4(opts.AccessKey, opts.SecretKey, ""), nil
+		}
 	} else if opts.RoleARN != "" {
 		log.WithField("roleArn", opts.RoleARN).Info("Creating minio client using assumed-role credentials")
 		return GetAssumeRoleCredentials(opts)

diff --git a/s3/s3_test.go b/s3/s3_test.go
@@ -8,7 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-// TestNewS3Client tests the s3 construtor
+// TestNewS3Client tests the s3 constructor
 func TestNewS3Client(t *testing.T) {
 	opts := S3ClientOpts{
 		Endpoint:        "foo.com",
@@ -17,6 +17,7 @@ func TestNewS3Client(t *testing.T) {
 		Transport:       http.DefaultTransport,
 		AccessKey:       "key",
 		SecretKey:       "secret",
+		SessionToken:    "",
 		Trace:           true,
 		RoleARN:         "",
 		RoleSessionName: "",
@@ -31,14 +32,46 @@ func TestNewS3Client(t *testing.T) {
 	assert.Equal(t, opts.Secure, s3cli.Secure)
 	assert.Equal(t, opts.Transport, s3cli.Transport)
 	assert.Equal(t, opts.AccessKey, s3cli.AccessKey)
+	assert.Equal(t, opts.SessionToken, s3cli.SessionToken)
 	assert.Equal(t, opts.Trace, s3cli.Trace)
 	assert.Equal(t, opts.EncryptOpts, s3cli.EncryptOpts)
 	assert.Equal(t, opts.AddressingStyle, s3cli.AddressingStyle)
 	// s3cli.minioClient.
 	// 	s3client.minioClient
 }
 
-// TestNewS3Client tests the s3 construtor
+// TestNewS3Client tests the S3 constructor using ephemeral credentials
+func TestNewS3ClientEphemeral(t *testing.T) {
+	opts := S3ClientOpts{
+		Endpoint:        "foo.com",
+		Region:          "us-south-3",
+		Secure:          false,
+		Transport:       http.DefaultTransport,
+		AccessKey:       "key",
+		SecretKey:       "secret",
+		SessionToken:    "sessionToken",
+		Trace:           true,
+		RoleARN:         "",
+		RoleSessionName: "",
+		UseSDKCreds:     false,
+		EncryptOpts:     EncryptOpts{Enabled: true, ServerSideCustomerKey: "", KmsKeyId: "", KmsEncryptionContext: ""},
+	}
+	s3If, err := NewS3Client(context.Background(), opts)
+	assert.NoError(t, err)
+	s3cli := s3If.(*s3client)
+	assert.Equal(t, opts.Endpoint, s3cli.Endpoint)
+	assert.Equal(t, opts.Region, s3cli.Region)
+	assert.Equal(t, opts.Secure, s3cli.Secure)
+	assert.Equal(t, opts.Transport, s3cli.Transport)
+	assert.Equal(t, opts.AccessKey, s3cli.AccessKey)
+	assert.Equal(t, opts.SecretKey, s3cli.SecretKey)
+	assert.Equal(t, opts.SessionToken, s3cli.SessionToken)
+	assert.Equal(t, opts.Trace, s3cli.Trace)
+	assert.Equal(t, opts.EncryptOpts, s3cli.EncryptOpts)
+	assert.Equal(t, opts.AddressingStyle, s3cli.AddressingStyle)
+}
+
+// TestNewS3Client tests the s3 constructor
 func TestNewS3ClientWithDiff(t *testing.T) {
 	t.Run("IAMRole", func(t *testing.T) {
 		opts := S3ClientOpts{