Add support for oidc aws session tags

buildkite · Sep 24, 2024 · 243b1a8 · 243b1a8
1 parent 65ed6aa
commit 243b1a8
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 15 deletions.
diff --git a/api/oidc.go b/api/oidc.go
@@ -10,21 +10,24 @@ type OIDCToken struct {
 }
 
 type OIDCTokenRequest struct {
-	Job      string
-	Audience string
-	Lifetime int
-	Claims   []string
+	Job            string
+	Audience       string
+	Lifetime       int
+	Claims         []string
+	AwsSessionTags []string
 }
 
 func (c *Client) OIDCToken(ctx context.Context, methodReq *OIDCTokenRequest) (*OIDCToken, *Response, error) {
 	m := &struct {
-		Audience string   `json:"audience,omitempty"`
-		Lifetime int      `json:"lifetime,omitempty"`
-		Claims   []string `json:"claims,omitempty"`
+		Audience       string   `json:"audience,omitempty"`
+		Lifetime       int      `json:"lifetime,omitempty"`
+		Claims         []string `json:"claims,omitempty"`
+		AwsSessionTags []string `json:"aws_session_tags,omitempty"`
 	}{
-		Audience: methodReq.Audience,
-		Lifetime: methodReq.Lifetime,
-		Claims:   methodReq.Claims,
+		Audience:       methodReq.Audience,
+		Lifetime:       methodReq.Lifetime,
+		Claims:         methodReq.Claims,
+		AwsSessionTags: methodReq.AwsSessionTags,
 	}
 
 	u := fmt.Sprintf("jobs/%s/oidc/tokens", railsPathEscape(methodReq.Job))

diff --git a/api/oidc_test.go b/api/oidc_test.go
@@ -124,6 +124,24 @@ func TestOIDCToken(t *testing.T) {
 			ExpectedBody: []byte(fmt.Sprintf(`{"lifetime":%d}`+"\n", lifetime)),
 			OIDCToken:    &api.OIDCToken{Token: oidcToken},
 		},
+		{
+			AccessToken: accessToken,
+			OIDCTokenRequest: &api.OIDCTokenRequest{
+				Job:    jobID,
+				Claims: []string{"organization_id", "pipeline_id"},
+			},
+			ExpectedBody: []byte(`{"claims":["organization_id","pipeline_id"]}` + "\n"),
+			OIDCToken:    &api.OIDCToken{Token: oidcToken},
+		},
+		{
+			AccessToken: accessToken,
+			OIDCTokenRequest: &api.OIDCTokenRequest{
+				Job:            jobID,
+				AwsSessionTags: []string{"organization_id", "pipeline_id"},
+			},
+			ExpectedBody: []byte(`{"aws_session_tags":["organization_id","pipeline_id"]}` + "\n"),
+			OIDCToken:    &api.OIDCToken{Token: oidcToken},
+		},
 	}
 
 	for _, test := range tests {

diff --git a/clicommand/oidc_request_token.go b/clicommand/oidc_request_token.go
@@ -16,7 +16,8 @@ type OIDCTokenConfig struct {
 	Lifetime int    `cli:"lifetime"`
 	Job      string `cli:"job"      validate:"required"`
 	// TODO: enumerate possible values, perhaps by adding a link to the documentation
-	Claims []string `cli:"claim"    normalize:"list"`
+	Claims         []string `cli:"claim"           normalize:"list"`
+	AwsSessionTags []string `cli:"aws-session-tag" normalize:"list"`
 
 	// Global flags
 	Debug       bool     `cli:"debug"`
@@ -79,6 +80,13 @@ var OIDCRequestTokenCommand = cli.Command{
 			EnvVar: "BUILDKITE_OIDC_TOKEN_CLAIMS",
 		},
 
+		cli.StringSliceFlag{
+			Name:   "aws-session-tag",
+			Value:  &cli.StringSlice{},
+			Usage:  "Add claims as AWS Session Tags",
+			EnvVar: "BUILDKITE_OIDC_TOKEN_AWS_SESSION_TAGS",
+		},
+
 		// API Flags
 		AgentAccessTokenFlag,
 		EndpointFlag,
@@ -112,10 +120,11 @@ var OIDCRequestTokenCommand = cli.Command{
 		)
 		token, err := roko.DoFunc(ctx, r, func(r *roko.Retrier) (*api.OIDCToken, error) {
 			req := &api.OIDCTokenRequest{
-				Job:      cfg.Job,
-				Audience: cfg.Audience,
-				Lifetime: cfg.Lifetime,
-				Claims:   cfg.Claims,
+				Job:            cfg.Job,
+				Audience:       cfg.Audience,
+				Lifetime:       cfg.Lifetime,
+				Claims:         cfg.Claims,
+				AwsSessionTags: cfg.AwsSessionTags,
 			}
 
 			token, resp, err := client.OIDCToken(ctx, req)