Merge pull request #167 from drone-plugins/CI-14214

feat: [CI-14214]: Add PLUGIN_USER_ROLE_EXTERNAL_ID to pass external ID for the secondary role when required
drone-plugins · Sep 25, 2024 · f132cba · f132cba
2 parents c9a496f + 8813c3d
commit f132cba
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 53 deletions.
diff --git a/README.md b/README.md
@@ -48,7 +48,7 @@ docker: Error response from daemon: Container command
 
 Execute from the working directory:
 
-* For upload
+* For Upload
 ```
 docker run --rm \
   -e PLUGIN_SOURCE=<source> \
@@ -61,7 +61,7 @@ docker run --rm \
   plugins/s3 --dry-run
 ```
 
-* For download
+* For Download
 ```
 docker run --rm \
   -e PLUGIN_SOURCE=<source directory to be downloaded from bucket> \

diff --git a/main.go b/main.go
@@ -52,6 +52,11 @@ func main() {
 			Usage:  "AWS user role",
 			EnvVar: "PLUGIN_USER_ROLE_ARN,AWS_USER_ROLE_ARN",
 		},
+		cli.StringFlag{
+			Name:   "user-role-external-id",
+			Usage:  "external ID to use when assuming secondary role",
+			EnvVar: "PLUGIN_USER_ROLE_EXTERNAL_ID",
+		},
 		cli.StringFlag{
 			Name:   "bucket",
 			Usage:  "aws bucket",
@@ -166,6 +171,7 @@ func run(c *cli.Context) error {
 		AssumeRoleSessionName: c.String("assume-role-session-name"),
 		Bucket:                c.String("bucket"),
 		UserRoleArn:           c.String("user-role-arn"),
+		UserRoleExternalID:    c.String("user-role-external-id"),
 		Region:                c.String("region"),
 		Access:                c.String("acl"),
 		Source:                c.String("source"),
@@ -181,7 +187,7 @@ func run(c *cli.Context) error {
 		PathStyle:             c.Bool("path-style"),
 		DryRun:                c.Bool("dry-run"),
 		ExternalID:            c.String("external-id"),
-		IdToken: 			   c.String("oidc-token-id"),
+		IdToken:               c.String("oidc-token-id"),
 	}
 
 	return plugin.Exec()

diff --git a/plugin.go b/plugin.go
@@ -29,6 +29,7 @@ type Plugin struct {
 	AssumeRoleSessionName string
 	Bucket                string
 	UserRoleArn           string
+	UserRoleExternalID    string // New field for UserRoleArn ExternalID
 
 	// 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
 }
 
@@ -434,60 +435,79 @@ 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),
-    }
-
-    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)")
-    }
+	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)")
+	}
 
 	sess, err = session.NewSession(conf)
-    if err != nil {
-        log.Fatalf("failed to create AWS session: %v", err)
-    }
+	if err != nil {
+		log.Fatalf("failed to create AWS session: %v", err)
+	}
+
+	client := s3.New(sess, conf)
+
+	if len(p.UserRoleArn) > 0 {
+		log.WithFields(log.Fields{
+			"UserRoleArn": p.UserRoleArn,
+		}).Info("Assuming user role ARN")
 
-    client := s3.New(sess, conf)
+		// Create new credentials by assuming the UserRoleArn with ExternalID
+		creds := stscreds.NewCredentials(sess, p.UserRoleArn, func(provider *stscreds.AssumeRoleProvider) {
+			if p.UserRoleExternalID != "" {
+				provider.ExternalID = aws.String(p.UserRoleExternalID)
+			}
+		})
 
-    if len(p.UserRoleArn) > 0 {
-        confRoleArn := aws.Config{
-            Region:      aws.String(p.Region),
-            Credentials: stscreds.NewCredentials(sess, p.UserRoleArn),
-        }
-        client = s3.New(sess, &confRoleArn)
-    }
+		// 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
 
-    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
+}