Merge pull request #627 from shogo82148/support-imds-v2

Support IMDSv2 for AWS EC2

spec/cloud.go

@@ -1,6 +1,7 @@
 package spec
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -14,9 +15,8 @@ import (
 
 	"github.com/Songmu/retry"
 	"github.com/mackerelio/golib/logging"
-	"github.com/mackerelio/mackerel-client-go"
-
 	"github.com/mackerelio/mackerel-agent/config"
+	"github.com/mackerelio/mackerel-client-go"
 )
 
 // This Generator collects metadata about cloud instances.
@@ -122,12 +122,12 @@ func (s *cloudGeneratorSuggester) Suggest(conf *config.Config) *CloudGenerator {
 var CloudGeneratorSuggester *cloudGeneratorSuggester
 
 func init() {
-	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
+	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest")
 	gceMetaURL, _ = url.Parse("http://metadata.google.internal./computeMetadata/v1")
 	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
 
 	CloudGeneratorSuggester = &cloudGeneratorSuggester{
-		ec2Generator:     &EC2Generator{ec2BaseURL},
+		ec2Generator:     &EC2Generator{baseURL: ec2BaseURL},
 		gceGenerator:     &GCEGenerator{gceMetaURL, gceMeta{}},
 		azureVMGenerator: &AzureVMGenerator{azureVMBaseURL},
 	}
@@ -148,12 +148,94 @@ func httpCli() *http.Client {
 // EC2Generator meta generator for EC2
 type EC2Generator struct {
 	baseURL *url.URL
+
+	mu       sync.Mutex
+	token    string
+	deadline time.Time
 }
 
 // IsEC2 checks current environment is EC2 or not
 func (g *EC2Generator) IsEC2(ctx context.Context) bool {
 	// implementation varies between OSs. see isec2_XXX.go
-	return isEC2(ctx)
+	return g.isEC2(ctx)
+}
+
+// check whether IMDS(Instance Metadata Service) is available.
+// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
+func (g *EC2Generator) hasMetadataService(ctx context.Context) (bool, error) {
+	cl := httpCli()
+
+	// services/partition is an AWS specific URL
+	req, err := http.NewRequest("GET", g.baseURL.String()+"/meta-data/services/partition", nil)
+	if err != nil {
+		return false, nil // something wrong. give up
+	}
+
+	// try to refresh api token for IMDSv2
+	// if it fails, fallback to IMDSv1.
+	if token := g.refreshToken(ctx); token != "" {
+		req.Header.Set("X-aws-ec2-metadata-token", token)
+	}
+
+	resp, err := cl.Do(req.WithContext(ctx))
+	if err != nil {
+		return false, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		return false, nil
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return false, err
+	}
+
+	// For standard AWS Regions, the partition is "aws".
+	// If you have resources in other partitions, the partition is "aws-partitionname".
+	// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-categories.html
+	return bytes.HasPrefix(body, []byte("aws")), nil
+}
+
+// refresh api token for IMDSv2
+func (g *EC2Generator) refreshToken(ctx context.Context) string {
+	cl := httpCli()
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	now := time.Now()
+	if !g.deadline.IsZero() && now.Before(g.deadline) {
+		return g.token // no need to refresh
+	}
+
+	req, err := http.NewRequest("PUT", g.baseURL.String()+"/api/token", nil)
+	if err != nil {
+		return ""
+	}
+	// TTL is 6 hours
+	req.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+	resp, err := cl.Do(req.WithContext(ctx))
+	if err != nil {
+		// IMDSv2 may be disabled? fallback to IMDSv1
+		g.token = ""
+		g.deadline = now.Add(30 * time.Minute)
+		return ""
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		// IMDSv2 may be disabled? fallback to IMDSv1
+		g.token = ""
+		g.deadline = now.Add(30 * time.Minute)
+		return ""
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return ""
+	}
+	g.deadline = now.Add(6*time.Hour - 30*time.Minute)
+	g.token = string(body)
+	return g.token
 }
 
 // Generate collects metadata from cloud platform.
@@ -177,7 +259,15 @@ func (g *EC2Generator) Generate() (*mackerel.Cloud, error) {
 	metadata := make(map[string]string)
 
 	for _, key := range metadataKeys {
-		resp, err := cl.Get(g.baseURL.String() + "/" + key)
+		req, err := http.NewRequest("GET", g.baseURL.String()+"/meta-data/"+key, nil)
+		if err != nil {
+			cloudLogger.Debugf("Unexpected error while requesting metadata: '%s'", err)
+			return nil, nil
+		}
+		if token := g.refreshToken(context.Background()); token != "" {
+			req.Header.Set("X-aws-ec2-metadata-token", token)
+		}
+		resp, err := cl.Do(req)
 		if err != nil {
 			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
 			return nil, nil
@@ -201,11 +291,17 @@ func (g *EC2Generator) Generate() (*mackerel.Cloud, error) {
 
 // SuggestCustomIdentifier suggests the identifier of the EC2 instance
 func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
+	cl := httpCli()
 	identifier := ""
 	err := retry.Retry(3, 2*time.Second, func() error {
-		cl := httpCli()
-		key := "instance-id"
-		resp, err := cl.Get(g.baseURL.String() + "/" + key)
+		req, err := http.NewRequest("GET", g.baseURL.String()+"/meta-data/instance-id", nil)
+		if err != nil {
+			return fmt.Errorf("error while retrieving instance-id: %s", err)
+		}
+		if token := g.refreshToken(context.Background()); token != "" {
+			req.Header.Set("X-aws-ec2-metadata-token", token)
+		}
+		resp, err := cl.Do(req)
 		if err != nil {
 			return fmt.Errorf("error while retrieving instance-id: %s", err)
 		}

spec/cloud_test.go

@@ -8,12 +8,12 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"reflect"
+	"strconv"
 	"testing"
 	"time"
 
-	"github.com/mackerelio/mackerel-client-go"
-
 	"github.com/mackerelio/mackerel-agent/config"
+	"github.com/mackerelio/mackerel-client-go"
 )
 
 type mockCloudMetaGenerator struct {
@@ -102,11 +102,11 @@ func TestCloudGenerator(t *testing.T) {
 	}
 }
 
-func TestEC2Generator(t *testing.T) {
+func TestEC2GeneratorIMDSv1(t *testing.T) {
 	handler := func(res http.ResponseWriter, req *http.Request) {
 		// The REAL path is /latest/meta-data/instance-id.
 		// This odd Path is due to current implementation.
-		if req.URL.Path == "/instance-id" {
+		if req.URL.Path == "/meta-data/instance-id" {
 			fmt.Fprint(res, "i-4f90d537")
 		} else {
 			http.Error(res, "not found", 404)
@@ -121,7 +121,74 @@ func TestEC2Generator(t *testing.T) {
 	if err != nil {
 		t.Errorf("should not raise error: %s", err)
 	}
-	g := &EC2Generator{u}
+	g := &EC2Generator{baseURL: u}
+
+	customIdentifier, err := g.SuggestCustomIdentifier()
+	if err != nil {
+		t.Errorf("should not raise error: %s", err)
+	}
+
+	if customIdentifier != "i-4f90d537.ec2.amazonaws.com" {
+		t.Errorf("Unexpected customIdentifier: %s", customIdentifier)
+	}
+
+	cloud, err := g.Generate()
+	if err != nil {
+		t.Errorf("should not raise error: %s", err)
+	}
+
+	if cloud == nil {
+		t.Error("cloud should not be nil")
+		return
+	}
+
+	metadata, typeOk := cloud.MetaData.(map[string]string)
+	if !typeOk {
+		t.Errorf("MetaData should be map. %+v", cloud.MetaData)
+	}
+
+	if metadata == nil || metadata["instance-id"] != "i-4f90d537" {
+		t.Errorf("Unexpected metadata: %+v", metadata)
+	}
+}
+
+func TestEC2GeneratorIMDSv2(t *testing.T) {
+	handler := func(res http.ResponseWriter, req *http.Request) {
+		const token = "very-secret"
+		switch req.Method {
+		case "PUT":
+			if _, err := strconv.Atoi(req.Header.Get("X-aws-ec2-metadata-token-ttl-seconds")); err != nil {
+				http.Error(res, "X-aws-ec2-metadata-token-ttl-seconds header is missing", 400)
+				return
+			}
+			if req.URL.Path == "/api/token" {
+				fmt.Fprint(res, token)
+				return
+			}
+		case "GET":
+			if req.Header.Get("X-aws-ec2-metadata-token") != token {
+				http.Error(res, "Unauthorized", 400)
+				return
+			}
+			// The REAL path is /latest/meta-data/instance-id.
+			// This odd Path is due to current implementation.
+			if req.URL.Path == "/meta-data/instance-id" {
+				fmt.Fprint(res, "i-4f90d537")
+				return
+			}
+		}
+		http.Error(res, "not found", 404)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
+		handler(res, req)
+	}))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Errorf("should not raise error: %s", err)
+	}
+	g := &EC2Generator{baseURL: u}
 
 	customIdentifier, err := g.SuggestCustomIdentifier()
 	if err != nil {
@@ -172,7 +239,7 @@ func TestEC2SuggestCustomIdentifier_ChangingHttpStatus(t *testing.T) {
 	if err != nil {
 		t.Errorf("should not raise error: %s", err)
 	}
-	g := &EC2Generator{u}
+	g := &EC2Generator{baseURL: u}
 
 	// 404, 404, 404 => give up
 	{

spec/isec2.go

@@ -4,30 +4,21 @@ package spec
 
 import (
 	"context"
-	"net/http"
 	"time"
 
 	"github.com/Songmu/retry"
 )
 
 // For instances other than Linux, retry only 1 times to shorten whole process
-func isEC2(ctx context.Context) bool {
-	isEC2 := false
+func (g *EC2Generator) isEC2(ctx context.Context) bool {
+	var res bool
 	err := retry.WithContext(ctx, 2, 2*time.Second, func() error {
-		cl := httpCli()
-		// '/ami-id` is probably an AWS specific URL
-		req, err := http.NewRequest("GET", ec2BaseURL.String()+"/ami-id", nil)
+		res0, err := g.hasMetadataService(ctx)
 		if err != nil {
 			return err
 		}
-		resp, err := cl.Do(req.WithContext(ctx))
-		if err != nil {
-			return err
-		}
-		defer resp.Body.Close()
-
-		isEC2 = resp.StatusCode == 200
+		res = res0
 		return nil
 	})
-	return err == nil && isEC2
+	return err == nil && res
 }

spec/isec2_linux.go

@@ -7,7 +7,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"io/ioutil"
-	"net/http"
 	"strings"
 	"time"
 
@@ -16,16 +15,16 @@ import (
 
 // If the OS is Linux, check /sys/hypervisor/uuid and /sys/devices/virtual/dmi/id/product_uuid files first. If UUID seems to be EC2-ish, call the metadata API (up to 3 times).
 // ref. https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html
-func isEC2(ctx context.Context) bool {
+func (g *EC2Generator) isEC2(ctx context.Context) bool {
 	var uuidFiles = []string{
 		"/sys/hypervisor/uuid",
 		"/sys/devices/virtual/dmi/id/product_uuid",
 	}
 
-	return isEC2WithSpecifiedUUIDFiles(ctx, uuidFiles)
+	return g.isEC2WithSpecifiedUUIDFiles(ctx, uuidFiles)
 }
 
-func isEC2WithSpecifiedUUIDFiles(ctx context.Context, uuidFiles []string) bool {
+func (g *EC2Generator) isEC2WithSpecifiedUUIDFiles(ctx context.Context, uuidFiles []string) bool {
 	looksLikeEC2 := false
 	for _, u := range uuidFiles {
 		data, err := ioutil.ReadFile(u)
@@ -48,29 +47,16 @@ func isEC2WithSpecifiedUUIDFiles(ctx context.Context, uuidFiles []string) bool {
 	default:
 	}
 
-	res := false
-	cl := httpCli()
+	var res bool
 	err := retry.WithContext(ctx, 3, 2*time.Second, func() error {
-		// '/ami-id` is probably an AWS specific URL
-		req, err := http.NewRequest("GET", ec2BaseURL.String()+"/ami-id", nil)
-		if err != nil {
-			return nil // something wrong. give up
-		}
-		resp, err := cl.Do(req.WithContext(ctx))
+		res0, err := g.hasMetadataService(ctx)
 		if err != nil {
 			return err
 		}
-		defer resp.Body.Close()
-
-		res = resp.StatusCode == 200
+		res = res0
 		return nil
 	})
-
-	if err == nil {
-		return res
-	}
-
-	return false
+	return err == nil && res
 }
 
 func isEC2UUID(uuid string) bool {