Update plugin.go

Author: Ompragash

plugin.go

@@ -29,6 +29,7 @@ type Plugin struct {
 	AssumeRoleSessionName string
 	Bucket                string
 	UserRoleArn           string
+	UserRoleExternalID    string
 
 	// if not "", enable server-side encryption
 	// valid values are:
@@ -99,7 +100,7 @@ type Plugin struct {
 	// set externalID for assume role
 	ExternalID string
 
-	// set OIDC ID Token to retrieve temporary credentials
+	// set OIDC ID Token to retrieve temporary credentials 
 	IdToken string
 }
 
@@ -280,7 +281,8 @@ func matchExtension(match string, stringMap map[string]string) string {
 	return ""
 }
 
-func assumeRole(sess *session.Session, roleArn, roleSessionName, externalID string) *credentials.Credentials {
+func assumeRole(roleArn, roleSessionName, externalID string) *credentials.Credentials {
+	sess, _ := session.NewSession()
 	client := sts.New(sess)
 	duration := time.Hour * 1
 	stsProvider := &stscreds.AssumeRoleProvider{
@@ -433,95 +435,69 @@ func (p *Plugin) downloadS3Objects(client *s3.S3, sourceDir string) error {
 
 // createS3Client creates and returns an S3 client based on the plugin configuration
 func (p *Plugin) createS3Client() *s3.S3 {
-	conf := &aws.Config{
-		Region:           aws.String(p.Region),
-		Endpoint:         &p.Endpoint,
-		DisableSSL:       aws.Bool(strings.HasPrefix(p.Endpoint, "http://")),
-		S3ForcePathStyle: aws.Bool(p.PathStyle),
-	}
-
-	var creds *credentials.Credentials
-	var err error
-
-	// Set credentials before creating the session
-	if p.Key != "" && p.Secret != "" {
-		// Use static credentials
-		creds = credentials.NewStaticCredentials(p.Key, p.Secret, "")
-		conf.Credentials = creds
-	} else if p.IdToken != "" && p.AssumeRole != "" {
-		// Assume role with web identity
-		tempSess, err := session.NewSession(conf)
-		if err != nil {
-			log.Fatalf("failed to create temporary AWS session: %v", err)
-		}
-
-		creds, err = assumeRoleWithWebIdentity(tempSess, p.AssumeRole, p.AssumeRoleSessionName, p.IdToken)
-		if err != nil {
-			log.Fatalf("failed to assume role with web identity: %v", err)
-		}
-
-		// Update the credentials in the config
-		conf.Credentials = creds
-	} else if p.AssumeRole != "" {
-		// Standard AssumeRole
-		tempSess, err := session.NewSession(conf)
-		if err != nil {
-			log.Fatalf("failed to create temporary AWS session: %v", err)
-		}
-
-		creds = assumeRole(tempSess, p.AssumeRole, p.AssumeRoleSessionName, p.ExternalID)
-
-		// Update the credentials in the config
-		conf.Credentials = creds
-	} else {
-		log.Warn("AWS Key and/or Secret not provided (falling back to EC2 instance profile or environment variables)")
-	}
-
-	// Now create the main session with the credentials
-	sess, err := session.NewSession(conf)
-	if err != nil {
-		log.Fatalf("failed to create AWS session: %v", err)
-	}
-
-	// Create the S3 client using the session
-	client := s3.New(sess)
-
-	// Only attempt to assume UserRoleArn if it's provided
-	if len(p.UserRoleArn) > 0 {
-		log.WithFields(log.Fields{
-			"UserRoleArn": p.UserRoleArn,
-		}).Info("Assuming user role ARN")
-
-		// Create new credentials by assuming the UserRoleArn
-		creds = stscreds.NewCredentials(sess, p.UserRoleArn)
-
-		// Create a new session with the new credentials
-		confWithUserRole := &aws.Config{
-			Region:      aws.String(p.Region),
-			Credentials: creds,
-		}
-
-		sessWithUserRole, err := session.NewSession(confWithUserRole)
-		if err != nil {
-			log.Fatalf("failed to create AWS session with user role: %v", err)
-		}
-
-		client = s3.New(sessWithUserRole)
-	}
-
-	return client
+    conf := &aws.Config{
+        Region:           aws.String(p.Region),
+        Endpoint:         &p.Endpoint,
+        DisableSSL:       aws.Bool(strings.HasPrefix(p.Endpoint, "http://")),
+        S3ForcePathStyle: aws.Bool(p.PathStyle),
+    }
+
+    sess, err := session.NewSession(conf)
+    if err != nil {
+        log.Fatalf("failed to create AWS session: %v", err)
+    }
+
+    if p.Key != "" && p.Secret != "" {
+        conf.Credentials = credentials.NewStaticCredentials(p.Key, p.Secret, "")
+    } else if p.IdToken != "" && p.AssumeRole != "" {
+        creds, err := assumeRoleWithWebIdentity(sess, p.AssumeRole, p.AssumeRoleSessionName, p.IdToken)
+        if err != nil {
+            log.Fatalf("failed to assume role with web identity: %v", err)
+        }
+        conf.Credentials = creds
+    } else if p.AssumeRole != "" {
+        conf.Credentials = assumeRole(p.AssumeRole, p.AssumeRoleSessionName, p.ExternalID)
+    } else {
+        log.Warn("AWS Key and/or Secret not provided (falling back to ec2 instance profile)")
+    }
+
+    client := s3.New(sess, conf)
+
+    if len(p.UserRoleArn) > 0 {
+	// Create new credentials by assuming the UserRoleArn (with ExternalID when provided)
+        creds := stscreds.NewCredentials(sess, p.UserRoleArn, func(provider *stscreds.AssumeRoleProvider) {
+            if p.UserRoleExternalID != "" {
+                provider.ExternalID = aws.String(p.UserRoleExternalID)
+            }
+        })
+
+        // Create a new session with the new credentials
+        confWithUserRole := &aws.Config{
+            Region:      aws.String(p.Region),
+            Credentials: creds,
+        }
+
+        sessWithUserRole, err := session.NewSession(confWithUserRole)
+        if err != nil {
+            log.Fatalf("failed to create AWS session with user role: %v", err)
+        }
+
+        client = s3.New(sessWithUserRole)
+    }
+
+    return client
 }
 
 func assumeRoleWithWebIdentity(sess *session.Session, roleArn, roleSessionName, idToken string) (*credentials.Credentials, error) {
-	svc := sts.New(sess)
-	input := &sts.AssumeRoleWithWebIdentityInput{
-		RoleArn:          aws.String(roleArn),
-		RoleSessionName:  aws.String(roleSessionName),
-		WebIdentityToken: aws.String(idToken),
-	}
-	result, err := svc.AssumeRoleWithWebIdentity(input)
-	if err != nil {
-		log.Fatalf("failed to assume role with web identity: %v", err)
-	}
-	return credentials.NewStaticCredentials(*result.Credentials.AccessKeyId, *result.Credentials.SecretAccessKey, *result.Credentials.SessionToken), nil
+    svc := sts.New(sess)
+    input := &sts.AssumeRoleWithWebIdentityInput{
+        RoleArn:          aws.String(roleArn),
+        RoleSessionName:  aws.String(roleSessionName),
+        WebIdentityToken: aws.String(idToken),
+    }
+    result, err := svc.AssumeRoleWithWebIdentity(input)
+    if err != nil {
+        log.Fatalf("failed to assume role with web identity: %v", err)
+    }
+    return credentials.NewStaticCredentials(*result.Credentials.AccessKeyId, *result.Credentials.SecretAccessKey, *result.Credentials.SessionToken), nil
 }