Add support for external-id and session-name when assuming a pods role

using STS

pkg/aws/sts/cache.go

@@ -36,7 +36,9 @@ type credentialsCache struct {
 }
 
 type CredentialsIdentity struct {
-	Role string
+	Role        string
+	SessionName string
+	ExternalID  string
 }
 
 type CachedCredentials struct {
@@ -128,7 +130,20 @@ func (c *credentialsCache) CredentialsForRole(ctx context.Context, identity *Cre
 
 	issue := func() (interface{}, error) {
 		arn := c.arnResolver.Resolve(identity.Role)
-		credentials, err := c.gateway.Issue(ctx, arn, c.sessionName, c.sessionDuration)
+		sessionName := c.sessionName
+
+		if identity.SessionName != "" {
+			sessionName = identity.SessionName
+		}
+
+		stsRequest := &IssueRequest{
+			RoleARN:         arn,
+			SessionName:     sessionName,
+			ExternalID:      identity.ExternalID,
+			SessionDuration: c.sessionDuration,
+		}
+
+		credentials, err := c.gateway.Issue(ctx, stsRequest)
 		if err != nil {
 			errorIssuing.Inc()
 			logger.Errorf("error requesting credentials: %s", err.Error())
@@ -157,5 +172,5 @@ func (c *credentialsCache) CredentialsForRole(ctx context.Context, identity *Cre
 }
 
 func (i *CredentialsIdentity) String() string {
-	return i.Role
+	return fmt.Sprintf("%s:%s:%s", i.Role, i.SessionName, i.ExternalID)
 }

pkg/aws/sts/credentials_cache_test.go

@@ -25,9 +25,9 @@ type stubGateway struct {
 	requestedRole string
 }
 
-func (s *stubGateway) Issue(ctx context.Context, roleARN, sessionName string, expiry time.Duration) (*Credentials, error) {
+func (s *stubGateway) Issue(ctx context.Context, request *IssueRequest) (*Credentials, error) {
 	s.issueCount = s.issueCount + 1
-	s.requestedRole = roleARN
+	s.requestedRole = request.RoleARN
 	return s.c, nil
 }
 

pkg/aws/sts/gateway.go

@@ -29,8 +29,15 @@ import (
 	"github.com/uswitch/kiam/pkg/statsd"
 )
 
+type IssueRequest struct {
+	RoleARN         string
+	SessionName     string
+	ExternalID      string
+	SessionDuration time.Duration
+}
+
 type STSGateway interface {
-	Issue(ctx context.Context, role, session string, expiry time.Duration) (*Credentials, error)
+	Issue(ctx context.Context, request *IssueRequest) (*Credentials, error)
 }
 
 type regionalResolver struct {
@@ -103,7 +110,7 @@ func DefaultGateway(assumeRoleArn, region string) (*DefaultSTSGateway, error) {
 	return &DefaultSTSGateway{session: session}, nil
 }
 
-func (g *DefaultSTSGateway) Issue(ctx context.Context, roleARN, sessionName string, expiry time.Duration) (*Credentials, error) {
+func (g *DefaultSTSGateway) Issue(ctx context.Context, request *IssueRequest) (*Credentials, error) {
 	timer := prometheus.NewTimer(assumeRole)
 	defer timer.ObserveDuration()
 	if statsd.Enabled {
@@ -115,10 +122,15 @@ func (g *DefaultSTSGateway) Issue(ctx context.Context, roleARN, sessionName stri
 
 	svc := sts.New(g.session)
 	in := &sts.AssumeRoleInput{
-		DurationSeconds: aws.Int64(int64(expiry.Seconds())),
-		RoleArn:         aws.String(roleARN),
-		RoleSessionName: aws.String(sessionName),
+		DurationSeconds: aws.Int64(int64(request.SessionDuration.Seconds())),
+		RoleArn:         aws.String(request.RoleARN),
+		RoleSessionName: aws.String(request.SessionName),
 	}
+
+	if request.ExternalID != "" {
+		in.ExternalId = aws.String(request.ExternalID)
+	}
+
 	resp, err := svc.AssumeRoleWithContext(ctx, in)
 	if err != nil {
 		return nil, err

pkg/k8s/pod_cache.go

@@ -19,7 +19,7 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/tools/cache"
 )
 
@@ -180,9 +180,25 @@ func PodRole(pod *v1.Pod) string {
 	return pod.ObjectMeta.Annotations[AnnotationIAMRoleKey]
 }
 
+// PodSessionName returns the IAM role session-name specified in the annotation for the Pod
+func PodSessionName(pod *v1.Pod) string {
+	return pod.ObjectMeta.Annotations[AnnotationIAMSessionNameKey]
+}
+
+// PodExternalID returns the IAM role external-id specified in the annotation for the Pod
+func PodExternalID(pod *v1.Pod) string {
+	return pod.ObjectMeta.Annotations[AnnotationIAMExternalIDKey]
+}
+
 // AnnotationIAMRoleKey is the key for the annotation specifying the IAM Role
 const AnnotationIAMRoleKey = "iam.amazonaws.com/role"
 
+// AnnotationIAMSessionNameKey is the key for the annotation specifying the session-name
+const AnnotationIAMSessionNameKey = "iam.amazonaws.com/session-name"
+
+// AnnotationIAMExternalIDKey is the key for the annotation specifying the external-id
+const AnnotationIAMExternalIDKey = "iam.amazonaws.com/external-id"
+
 type podHandler struct {
 	pods chan<- *v1.Pod
 }

pkg/prefetch/manager.go

@@ -38,8 +38,12 @@ func (m *CredentialManager) fetchCredentials(ctx context.Context, pod *v1.Pod) {
 		return
 	}
 
-	role := k8s.PodRole(pod)
-	identity := &sts.CredentialsIdentity{Role: role}
+	identity := &sts.CredentialsIdentity{
+		Role:        k8s.PodRole(pod),
+		SessionName: k8s.PodSessionName(pod),
+		ExternalID:  k8s.PodExternalID(pod),
+	}
+
 	issued, err := m.fetchCredentialsFromCache(ctx, identity)
 	if err != nil {
 		logger.Errorf("error warming credentials: %s", err.Error())

pkg/server/server.go

@@ -116,7 +116,13 @@ func (k *KiamServer) GetPodCredentials(ctx context.Context, req *pb.GetPodCreden
 		return nil, ErrPolicyForbidden
 	}
 
-	creds, err := k.credentialsProvider.CredentialsForRole(ctx, &sts.CredentialsIdentity{Role: req.Role})
+	identity := &sts.CredentialsIdentity{
+		Role:        req.Role,
+		SessionName: k8s.PodSessionName(pod),
+		ExternalID:  k8s.PodExternalID(pod),
+	}
+
+	creds, err := k.credentialsProvider.CredentialsForRole(ctx, identity)
 	if err != nil {
 		logger.Errorf("error retrieving credentials: %s", err.Error())
 		k.recordEvent(pod, v1.EventTypeWarning, "KiamCredentialError", fmt.Sprintf("failed retrieving credentials: %s", simplifyAWSErrorMessage(err)))