fix(provider): add session token for credential and use different dev…

…ice name by AMI
cnrancher · Jun 21, 2023 · fa84c48 · fa84c48
1 parent 4864363
commit fa84c48
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/pkg/providers/aws/aws.go b/pkg/providers/aws/aws.go
@@ -350,7 +350,7 @@ func (p *Amazon) generateInstance(ssh *types.SSH) (*types.Cluster, error) {
 func (p *Amazon) newClient() {
 	config := aws.NewConfig()
 	config = config.WithRegion(p.Region)
-	config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKey, p.SecretKey, ""))
+	config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKey, p.SecretKey, p.SessionToken))
 	sess := session.Must(session.NewSession(config))
 	p.client = ec2.New(sess)
 }
@@ -360,8 +360,19 @@ func (p *Amazon) runInstances(num int, master bool, ssh *types.SSH) error {
 	if err != nil {
 		return fmt.Errorf("[%s] --root-size is invalid %v, must be integer: %v", p.GetProviderName(), p.RootSize, err)
 	}
+	img, err := p.client.DescribeImages(&ec2.DescribeImagesInput{
+		ImageIds: aws.StringSlice([]string{p.AMI}),
+	})
+	if err != nil {
+		return fmt.Errorf("[%s] AMI %s is invalid: %v", p.GetProviderName(), p.AMI, err)
+	}
+	// get root device name from AMI
+	deviceName := aws.String(defaultDeviceName)
+	if len(img.Images) > 0 {
+		deviceName = img.Images[0].RootDeviceName
+	}
 	bdm := &ec2.BlockDeviceMapping{
-		DeviceName: aws.String(defaultDeviceName),
+		DeviceName: deviceName,
 		Ebs: &ec2.EbsBlockDevice{
 			VolumeSize:          aws.Int64(rootSize),
 			VolumeType:          aws.String(p.VolumeType),

diff --git a/pkg/providers/aws/flag.go b/pkg/providers/aws/flag.go
@@ -140,6 +140,14 @@ func (p *Amazon) GetCredentialFlags() []types.Flag {
 			Required: true,
 			EnvVar:   "AWS_SECRET_ACCESS_KEY",
 		},
+		{
+			Name:     "session-token",
+			P:        &p.SessionToken,
+			V:        p.SessionToken,
+			Usage:    "AWS session token",
+			Required: false,
+			EnvVar:   "AWS_SESSION_TOKEN",
+		},
 	}
 
 	return fs
@@ -157,8 +165,9 @@ func (p *Amazon) GetSSHConfig() *types.SSH {
 // BindCredential bind aws credential.
 func (p *Amazon) BindCredential() error {
 	secretMap := map[string]string{
-		"access-key": p.AccessKey,
-		"secret-key": p.SecretKey,
+		"access-key":    p.AccessKey,
+		"secret-key":    p.SecretKey,
+		"session-token": p.SessionToken,
 	}
 	return p.SaveCredential(secretMap)
 }

diff --git a/pkg/types/aws/aws.go b/pkg/types/aws/aws.go
@@ -4,6 +4,7 @@ package aws
 type Options struct {
 	AccessKey                    string   `json:"access-key,omitempty" yaml:"access-key,omitempty"`
 	SecretKey                    string   `json:"secret-key,omitempty" yaml:"secret-key,omitempty"`
+	SessionToken                 string   `json:"session-token,omitempty" yaml:"session-token,omitempty"`
 	Region                       string   `json:"region,omitempty" yaml:"region,omitempty"`
 	AMI                          string   `json:"ami,omitempty" yaml:"ami,omitempty"`
 	KeypairName                  string   `json:"keypair-name,omitempty" yaml:"keypair-name,omitempty"`