feat: optionally assume a role (#13)

Author: cwinters8

docs/data-sources/ssm.md

@@ -44,6 +44,7 @@ provider "postgresql" {
 
 - `ssm_profile` (String) AWS profile name as set in credentials files. Can also be set using either the environment variables `AWS_PROFILE` or `AWS_DEFAULT_PROFILE`.
 - `ssm_region` (String) AWS Region where the instance is located. The Region must be set. Can also be set using either the environment variables `AWS_REGION` or `AWS_DEFAULT_REGION`.
+- `ssm_role_arn` (String) ARN of an IAM role to assume.
 
 ### Read-Only
 

docs/ephemeral-resources/ssm.md

@@ -44,6 +44,7 @@ provider "postgresql" {
 
 - `ssm_profile` (String) AWS profile name as set in credentials files. Can also be set using either the environment variables `AWS_PROFILE` or `AWS_DEFAULT_PROFILE`.
 - `ssm_region` (String) AWS Region where the instance is located. The Region must be set. Can also be set using either the environment variables `AWS_REGION` or `AWS_DEFAULT_REGION`.
+- `ssm_role_arn` (String) ARN of an IAM role to assume.
 
 ### Read-Only
 

internal/provider/data_source_ssm.go

@@ -28,6 +28,7 @@ type SSMDataSourceModel struct {
 	LocalPort   types.Int64  `tfsdk:"local_port"`
 	SSMInstance types.String `tfsdk:"ssm_instance"`
 	SSMProfile  types.String `tfsdk:"ssm_profile"`
+	SSMRoleARN  types.String `tfsdk:"ssm_role_arn"`
 	SSMRegion   types.String `tfsdk:"ssm_region"`
 	TargetHost  types.String `tfsdk:"target_host"`
 	TargetPort  types.Int64  `tfsdk:"target_port"`
@@ -60,6 +61,11 @@ func (d *SSMDataSource) Schema(ctx context.Context, req datasource.SchemaRequest
 				Optional:            true,
 				Computed:            true,
 			},
+			"ssm_role_arn": schema.StringAttribute{
+				MarkdownDescription: "ARN of an IAM role to assume.",
+				Optional:            true,
+				Computed:            true,
+			},
 			"ssm_region": schema.StringAttribute{
 				MarkdownDescription: "AWS Region where the instance is located. The Region must be set. Can also be set using either the environment variables `AWS_REGION` or `AWS_DEFAULT_REGION`.",
 				Optional:            true,
@@ -105,6 +111,7 @@ func (d *SSMDataSource) Read(ctx context.Context, req datasource.ReadRequest, re
 		LocalPort:   strconv.Itoa(localPort),
 		SSMInstance: data.SSMInstance.ValueString(),
 		SSMProfile:  data.SSMProfile.ValueString(),
+		SSMRoleARN:  data.SSMRoleARN.ValueString(),
 		SSMRegion:   data.SSMRegion.ValueString(),
 		TargetHost:  data.TargetHost.ValueString(),
 		TargetPort:  strconv.Itoa(int(data.TargetPort.ValueInt64())),
@@ -119,8 +126,14 @@ func (d *SSMDataSource) Read(ctx context.Context, req datasource.ReadRequest, re
 	tunnelCfg.SSMRegion = awsCfg.Region
 	tunnelCfg.SSMProfile = ssm.GetSDKConfigProfile(awsCfg)
 
+	// Only update SSMRoleARN if it wasn't explicitly provided
+	if tunnelCfg.SSMRoleARN == "" {
+		tunnelCfg.SSMRoleARN = ssm.GetSDKConfigRole(awsCfg)
+	}
+
 	data.SSMRegion = types.StringValue(tunnelCfg.SSMRegion)
 	data.SSMProfile = types.StringValue(tunnelCfg.SSMProfile)
+	data.SSMRoleARN = types.StringValue(tunnelCfg.SSMRoleARN)
 
 	_, err = ssm.ForkRemoteTunnel(ctx, awsCfg, tunnelCfg)
 	if err != nil {

internal/provider/ephemeral_ssm.go

@@ -28,6 +28,7 @@ type SSMEphemeralModel struct {
 	LocalPort   types.Int64  `tfsdk:"local_port"`
 	SSMInstance types.String `tfsdk:"ssm_instance"`
 	SSMProfile  types.String `tfsdk:"ssm_profile"`
+	SSMRoleARN  types.String `tfsdk:"ssm_role_arn"`
 	SSMRegion   types.String `tfsdk:"ssm_region"`
 	TargetHost  types.String `tfsdk:"target_host"`
 	TargetPort  types.Int64  `tfsdk:"target_port"`
@@ -60,6 +61,11 @@ func (d *SSMEphemeral) Schema(ctx context.Context, req ephemeral.SchemaRequest,
 				Optional:            true,
 				Computed:            true,
 			},
+			"ssm_role_arn": schema.StringAttribute{
+				MarkdownDescription: "ARN of an IAM role to assume.",
+				Optional:            true,
+				Computed:            true,
+			},
 			"ssm_region": schema.StringAttribute{
 				MarkdownDescription: "AWS Region where the instance is located. The Region must be set. Can also be set using either the environment variables `AWS_REGION` or `AWS_DEFAULT_REGION`.",
 				Optional:            true,
@@ -105,6 +111,7 @@ func (d *SSMEphemeral) Open(ctx context.Context, req ephemeral.OpenRequest, resp
 		LocalPort:   strconv.Itoa(localPort),
 		SSMInstance: data.SSMInstance.ValueString(),
 		SSMProfile:  data.SSMProfile.ValueString(),
+		SSMRoleARN:  data.SSMRoleARN.ValueString(),
 		SSMRegion:   data.SSMRegion.ValueString(),
 		TargetHost:  data.TargetHost.ValueString(),
 		TargetPort:  strconv.Itoa(int(data.TargetPort.ValueInt64())),
@@ -119,8 +126,14 @@ func (d *SSMEphemeral) Open(ctx context.Context, req ephemeral.OpenRequest, resp
 	tunnelCfg.SSMRegion = awsCfg.Region
 	tunnelCfg.SSMProfile = ssm.GetSDKConfigProfile(awsCfg)
 
+	// Only update SSMRoleARN if it wasn't explicitly provided
+	if tunnelCfg.SSMRoleARN == "" {
+		tunnelCfg.SSMRoleARN = ssm.GetSDKConfigRole(awsCfg)
+	}
+
 	data.SSMRegion = types.StringValue(tunnelCfg.SSMRegion)
 	data.SSMProfile = types.StringValue(tunnelCfg.SSMProfile)
+	data.SSMRoleARN = types.StringValue(tunnelCfg.SSMRoleARN)
 
 	cmd, err := ssm.ForkRemoteTunnel(ctx, awsCfg, tunnelCfg)
 	if err != nil {

internal/ssm/session.go

@@ -5,7 +5,9 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
 	"github.com/aws/aws-sdk-go-v2/service/ssm"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
 )
 
 const DEFAULT_SSM_ENV_NAME = "AWS_SSM_START_SESSION_RESPONSE"
@@ -14,6 +16,7 @@ type TunnelConfig struct {
 	LocalPort   string
 	SSMInstance string
 	SSMProfile  string
+	SSMRoleARN  string
 	SSMRegion   string
 	TargetHost  string
 	TargetPort  string
@@ -34,7 +37,20 @@ func GetNewSDKConfig(ctx context.Context, cfg TunnelConfig) (aws.Config, error)
 		loadOptions = append(loadOptions, config.WithSharedConfigProfile(cfg.SSMProfile))
 	}
 
-	return config.LoadDefaultConfig(ctx, loadOptions...)
+	// Load base config first
+	awsCfg, err := config.LoadDefaultConfig(ctx, loadOptions...)
+	if err != nil {
+		return aws.Config{}, err
+	}
+
+	// If role assumption is required, create STS client and configure assume role
+	if cfg.SSMRoleARN != "" {
+		stsClient := sts.NewFromConfig(awsCfg)
+		assumeRoleProvider := stscreds.NewAssumeRoleProvider(stsClient, cfg.SSMRoleARN)
+		awsCfg.Credentials = aws.NewCredentialsCache(assumeRoleProvider)
+	}
+
+	return awsCfg, nil
 }
 
 func GetSDKConfigProfile(awsCfg aws.Config) string {
@@ -46,6 +62,15 @@ func GetSDKConfigProfile(awsCfg aws.Config) string {
 	return ""
 }
 
+func GetSDKConfigRole(awsCfg aws.Config) string {
+	for _, cfg := range awsCfg.ConfigSources {
+		if p, ok := cfg.(config.SharedConfig); ok {
+			return p.RoleARN
+		}
+	}
+	return ""
+}
+
 func CreateSessionInput(cfg TunnelConfig) ssm.StartSessionInput {
 	reqParams := make(map[string][]string)
 	reqParams["portNumber"] = []string{cfg.TargetPort}