Function AssumeRoleIfNeeded

backend/plugin/db/util/aws.go:22–58

AssumeRoleIfNeeded checks if role assumption is configured and updates the AWS config with assumed role credentials. This is a shared utility used by multiple AWS service drivers (Elasticsearch, RDS, etc.). Returns an error if role assumption fails.

(ctx context.Context, awsCfg *aws.Config, connectionCtx db.ConnectionContext, awsCredential *storepb.DataSource_AWSCredential)

// This is a shared utility used by multiple AWS service drivers (Elasticsearch, RDS, etc.).
// Returns an error if role assumption fails.
func AssumeRoleIfNeeded(ctx context.Context, awsCfg *aws.Config, connectionCtx db.ConnectionContext, awsCredential *storepb.DataSource_AWSCredential) error {
	// If no role ARN is provided, no assumption needed
	if awsCredential == nil || awsCredential.RoleArn == "" {
		return nil
	}

	roleArn := awsCredential.RoleArn

	// Create STS client with base credentials
	stsClient := sts.NewFromConfig(*awsCfg)

	// Generate descriptive session name for CloudTrail auditing
	sessionName := generateSessionName(connectionCtx.InstanceID)

	// Configure assume role provider
	assumeRoleProvider := stscreds.NewAssumeRoleProvider(stsClient, roleArn,
		func(o *stscreds.AssumeRoleOptions) {
			o.RoleSessionName = sessionName
			o.Duration = 1 * time.Hour // Temporary credentials valid for 1 hour

			// Add external ID if provided for additional security
			if externalID := awsCredential.ExternalId; externalID != "" {
				o.ExternalID = &externalID
			}
		})

	// Update config with assumed role credentials
	awsCfg.Credentials = assumeRoleProvider

	// Test credentials retrieval and provide context-specific error messages
	_, err := awsCfg.Credentials.Retrieve(ctx)
	if err != nil {
		return handleAssumeRoleError(err, roleArn, awsCredential.ExternalId)
	}

	return nil
}

// generateSessionName creates a descriptive session name for AWS CloudTrail auditing.
// Format: bytebase-{instance-id}-{timestamp}