Adds webIndentity token for aws (#1099)

https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1080

lib/promscrape/discovery/ec2/api.go

@@ -18,13 +18,17 @@ import (
 const (
 	awsAccessKeyEnv = "AWS_ACCESS_KEY_ID"
 	awsSecretKeyEnv = "AWS_SECRET_ACCESS_KEY"
+	awsRegionEnv    = "AWS_REGION"
+	awsRoleARNEnv   = "AWS_ROLE_ARN"
+	awsWITPath      = "AWS_WEB_IDENTITY_TOKEN_FILE"
 )
 
 type apiConfig struct {
-	region  string
+	region       string
-	roleARN string
+	roleARN      string
-	filters string
+	webTokenPath string
-	port    int
+	filters      string
+	port         int
 
 	ec2Endpoint string
 	stsEndpoint string
@@ -79,6 +83,14 @@ func newAPIConfig(sdc *SDConfig) (*apiConfig, error) {
 	cfg.ec2Endpoint = buildAPIEndpoint(sdc.Endpoint, region, "ec2")
 	cfg.stsEndpoint = buildAPIEndpoint(sdc.Endpoint, region, "sts")
 
+	envARN := os.Getenv(awsRoleARNEnv)
+	if envARN != "" {
+		cfg.roleARN = envARN
+	}
+	cfg.webTokenPath = os.Getenv(awsWITPath)
+	if cfg.webTokenPath != "" && cfg.roleARN == "" {
+		return nil, fmt.Errorf("roleARN is missing for %q, set it with cfg or env var %q", awsWITPath, awsRoleARNEnv)
+	}
 	// explicitly set credentials has priority over env variables
 	cfg.defaultAccessKey = os.Getenv(awsAccessKeyEnv)
 	cfg.defaultSecretKey = os.Getenv(awsSecretKeyEnv)
@@ -108,6 +120,11 @@ func getFiltersQueryString(filters []Filter) string {
 }
 
 func getDefaultRegion() (string, error) {
+	envRegion := os.Getenv(awsRegionEnv)
+	if envRegion != "" {
+		return envRegion, nil
+	}
+
 	data, err := getMetadataByPath("dynamic/instance-identity/document")
 	if err != nil {
 		return "", err
@@ -156,6 +173,13 @@ func getAPICredentials(cfg *apiConfig) (*apiCredentials, error) {
 		AccessKeyID:     cfg.defaultAccessKey,
 		SecretAccessKey: cfg.defaultSecretKey,
 	}
+	if len(cfg.webTokenPath) > 0 {
+		token, err := ioutil.ReadFile(cfg.webTokenPath)
+		if err != nil {
+			return nil, fmt.Errorf("cannot read webToken from path: %q, err: %w", cfg.webTokenPath, err)
+		}
+		return getRoleWebIdentityCredentials(cfg.stsEndpoint, cfg.roleARN, string(token))
+	}
 
 	// we need instance credentials if dont have access keys
 	if len(acNew.AccessKeyID) == 0 && len(acNew.SecretAccessKey) == 0 {
@@ -263,42 +287,73 @@ func getMetadataByPath(apiPath string) ([]byte, error) {
 	return readResponseBody(resp, apiURL)
 }
 
-// getRoleARNCredentials obtains credentials fo the given roleARN.
+// getRoleWebIdentityCredentials obtains credentials fo the given roleARN with webToken.
-func getRoleARNCredentials(region, stsEndpoint, roleARN string, creds *apiCredentials) (*apiCredentials, error) {
+// https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
-	data, err := getSTSAPIResponse(region, stsEndpoint, roleARN, creds)
+// aws IRSA for kubernetes.
+// https://aws.amazon.com/blogs/opensource/introducing-fine-grained-iam-roles-service-accounts/
+func getRoleWebIdentityCredentials(stsEndpoint, roleARN string, token string) (*apiCredentials, error) {
+	data, err := getSTSAPIResponse("AssumeRoleWithWebIdentity", stsEndpoint, roleARN, func(apiURL string) (*http.Request, error) {
+		apiURL += fmt.Sprintf("&WebIdentityToken=%s", token)
+		return http.NewRequest("GET", apiURL, nil)
+	})
 	if err != nil {
 		return nil, err
 	}
-	return parseARNCredentials(data)
+	return parseARNCredentials(data, "AssumeRoleWithWebIdentity")
+}
+
+// getRoleARNCredentials obtains credentials fo the given roleARN.
+func getRoleARNCredentials(region, stsEndpoint, roleARN string, creds *apiCredentials) (*apiCredentials, error) {
+	data, err := getSTSAPIResponse("AssumeRole", stsEndpoint, roleARN, func(apiURL string) (*http.Request, error) {
+		return newSignedRequest(apiURL, "sts", region, creds)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return parseARNCredentials(data, "AssumeRole")
 }
 
 // parseARNCredentials parses apiCredentials from AssumeRole response.
 //
 // See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
-func parseARNCredentials(data []byte) (*apiCredentials, error) {
+func parseARNCredentials(data []byte, role string) (*apiCredentials, error) {
 	var arr AssumeRoleResponse
 	if err := xml.Unmarshal(data, &arr); err != nil {
 		return nil, fmt.Errorf("cannot parse AssumeRoleResponse response from %q: %w", data, err)
 	}
+	var cred assumeCredentials
+	switch role {
+	case "AssumeRole":
+		cred = arr.AssumeRoleResult.Credentials
+	case "AssumeRoleWithWebIdentity":
+		cred = arr.AssumeRoleWithWebIdentityResult.Credentials
+	default:
+		return nil, fmt.Errorf("bug, unexpected role: %q", role)
+	}
 	return &apiCredentials{
-		AccessKeyID:     arr.AssumeRoleResult.Credentials.AccessKeyID,
+		AccessKeyID:     cred.AccessKeyID,
-		SecretAccessKey: arr.AssumeRoleResult.Credentials.SecretAccessKey,
+		SecretAccessKey: cred.SecretAccessKey,
-		Token:           arr.AssumeRoleResult.Credentials.SessionToken,
+		Token:           cred.SessionToken,
-		Expiration:      arr.AssumeRoleResult.Credentials.Expiration,
+		Expiration:      cred.Expiration,
 	}, nil
 }
 
+type assumeCredentials struct {
+	AccessKeyID     string    `xml:"AccessKeyId"`
+	SecretAccessKey string    `xml:"SecretAccessKey"`
+	SessionToken    string    `xml:"SessionToken"`
+	Expiration      time.Time `xml:"Expiration"`
+}
+
 // AssumeRoleResponse represents AssumeRole response
 //
 // See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
 type AssumeRoleResponse struct {
 	AssumeRoleResult struct {
-		Credentials struct {
+		Credentials assumeCredentials
-			AccessKeyID     string    `xml:"AccessKeyId"`
+	}
-			SecretAccessKey string    `xml:"SecretAccessKey"`
+	AssumeRoleWithWebIdentityResult struct {
-			SessionToken    string    `xml:"SessionToken"`
+		Credentials assumeCredentials
-			Expiration      time.Time `xml:"Expiration"`
-		}
 	}
 }
 
@@ -323,14 +378,14 @@ func buildAPIEndpoint(customEndpoint, region, service string) string {
 // and returns temporary credentials with expiration time
 //
 // See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
-func getSTSAPIResponse(region, stsEndpoint, roleARN string, creds *apiCredentials) ([]byte, error) {
+func getSTSAPIResponse(action, stsEndpoint, roleARN string, reqBuilder func(apiURL string) (*http.Request, error)) ([]byte, error) {
 	// See https://docs.aws.amazon.com/AWSEC2/latest/APIReference/Query-Requests.html
-	apiURL := fmt.Sprintf("%s?Action=%s", stsEndpoint, "AssumeRole")
+	apiURL := fmt.Sprintf("%s?Action=%s", stsEndpoint, action)
 	apiURL += "&Version=2011-06-15"
 	apiURL += fmt.Sprintf("&RoleArn=%s", roleARN)
 	// we have to provide unique session name for cloudtrail audit
 	apiURL += "&RoleSessionName=vmagent-ec2-discovery"
-	req, err := newSignedRequest(apiURL, "sts", region, creds)
+	req, err := reqBuilder(apiURL)
 	if err != nil {
 		return nil, fmt.Errorf("cannot create signed request: %w", err)
 	}

lib/promscrape/discovery/ec2/api_test.go

@@ -51,7 +51,7 @@ func TestParseMetadataSecurityCredentialsSuccess(t *testing.T) {
 func TestParseARNCredentialsFailure(t *testing.T) {
 	f := func(s string) {
 		t.Helper()
-		creds, err := parseARNCredentials([]byte(s))
+		creds, err := parseARNCredentials([]byte(s), "")
 		if err == nil {
 			t.Fatalf("expecting non-nil error")
 		}
@@ -64,6 +64,19 @@ func TestParseARNCredentialsFailure(t *testing.T) {
 }
 
 func TestParseARNCredentialsSuccess(t *testing.T) {
+
+	f := func(data, role string, credsExpected *apiCredentials) {
+		t.Helper()
+		creds, err := parseARNCredentials([]byte(data), role)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		if !reflect.DeepEqual(creds, credsExpected) {
+			t.Fatalf("unexpected creds;\ngot\n%+v\nwant\n%+v", creds, credsExpected)
+		}
+
+	}
 	// See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html
 	s := `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
   <AssumeRoleResult>
@@ -90,10 +103,6 @@ func TestParseARNCredentialsSuccess(t *testing.T) {
   </ResponseMetadata>
 </AssumeRoleResponse>
 `
-	creds, err := parseARNCredentials([]byte(s))
-	if err != nil {
-		t.Fatalf("unexpected error: %s", err)
-	}
 	credsExpected := &apiCredentials{
 		AccessKeyID:     "ASIAIOSFODNN7EXAMPLE",
 		SecretAccessKey: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY",
@@ -106,9 +115,35 @@ func TestParseARNCredentialsSuccess(t *testing.T) {
       `,
 		Expiration: mustParseRFC3339("2019-11-09T13:34:41Z"),
 	}
-	if !reflect.DeepEqual(creds, credsExpected) {
+	s2 := `<AssumeRoleWithWebIdentityResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
-		t.Fatalf("unexpected creds;\ngot\n%+v\nwant\n%+v", creds, credsExpected)
+  <AssumeRoleWithWebIdentityResult>
+    <Audience>sts.amazonaws.com</Audience>
+    <AssumedRoleUser>
+      <AssumedRoleId>AROA2X6NOXN27E3OGMK3T:vmagent-ec2-discovery</AssumedRoleId>
+      <Arn>arn:aws:sts::111111111:assumed-role/eks-role-9N0EFKEDJ1X/vmagent-ec2-discovery</Arn>
+    </AssumedRoleUser>
+    <Provider>arn:aws:iam::111111111:oidc-provider/oidc.eks.eu-west-1.amazonaws.com/id/111111111</Provider>
+    <Credentials>
+      <AccessKeyId>ASIABYASSDASF</AccessKeyId>
+      <SecretAccessKey>asffasfasf/RvxIQpCid4iRMGm56nnRs2oKgV</SecretAccessKey>
+      <SessionToken>asfafsassssssssss/MlyKUPOYAiEAq5HgS19Mf8SJ3kIKU3NCztDeZW5EUW4NrPrPyXQ8om0q/AQIjv//////////</SessionToken>
+      <Expiration>2021-03-01T13:38:15Z</Expiration>
+    </Credentials>
+    <SubjectFromWebIdentityToken>system:serviceaccount:default:vmagent</SubjectFromWebIdentityToken>
+  </AssumeRoleWithWebIdentityResult>
+  <ResponseMetadata>
+    <RequestId>1214124-7bb0-4673-ad6d-af9e67fc1141</RequestId>
+  </ResponseMetadata>
+</AssumeRoleWithWebIdentityResponse>`
+	credsExpected2 := &apiCredentials{
+		AccessKeyID:     "ASIABYASSDASF",
+		SecretAccessKey: "asffasfasf/RvxIQpCid4iRMGm56nnRs2oKgV",
+		Token:           "asfafsassssssssss/MlyKUPOYAiEAq5HgS19Mf8SJ3kIKU3NCztDeZW5EUW4NrPrPyXQ8om0q/AQIjv//////////",
+		Expiration:      mustParseRFC3339("2021-03-01T13:38:15Z"),
 	}
+
+	f(s, "AssumeRole", credsExpected)
+	f(s2, "AssumeRoleWithWebIdentity", credsExpected2)
 }
 
 func mustParseRFC3339(s string) time.Time {