Add support for autodetecting EKS when IMDSv2 is in use

cmd/kubeletmein/generate.go

@@ -55,6 +55,12 @@ func Generate(c *config.Config) *cobra.Command {
 				if err != nil {
 					return err
 				}
+			case "eks-imdsv2":
+				logger.Info("EKS using IMDSv2 detected")
+				err := eks.Generate(c)
+				if err != nil {
+					return err
+				}
 			case "autodetect":
 				logger.Debug("autodetect? We should not have got here")
 			default:

pkg/autodetect/autodetect.go

@@ -12,8 +12,10 @@ type Providers map[string]Provider
 // Provider stores the details of each cloud provider
 type Provider struct {
 	Path               string
-	Header             map[string]string
+	RequestHeader      map[string]string
+	ResponseHeader     map[string]string
 	ExpectedStatusCode int
+	Method             string
 }
 
 // Client wraps http.Client so we can mock
@@ -30,17 +32,30 @@ var (
 	PublicCloudProviders = Providers{
 		"gke": Provider{
 			Path:               "/",
-			Header:             map[string]string{"Server": "Metadata Server for VM"},
+			Method:             "GET",
+			RequestHeader:      map[string]string{},
+			ResponseHeader:     map[string]string{"Server": "Metadata Server for VM"},
 			ExpectedStatusCode: http.StatusOK,
 		},
 		"do": Provider{
 			Path:               "/metadata/v1/id",
-			Header:             map[string]string{"Content-Type": "text/plain; charset=utf-8"},
+			Method:             "GET",
+			RequestHeader:      map[string]string{},
+			ResponseHeader:     map[string]string{"Content-Type": "text/plain; charset=utf-8"},
 			ExpectedStatusCode: http.StatusOK,
 		},
 		"eks": Provider{
 			Path:               "/",
-			Header:             map[string]string{"Server": "EC2ws"},
+			Method:             "GET",
+			RequestHeader:      map[string]string{},
+			ResponseHeader:     map[string]string{"Server": "EC2ws"},
+			ExpectedStatusCode: http.StatusOK,
+		},
+		"eks-imdsv2": Provider{
+			Path:               "/latest/api/token",
+			Method:             "PUT",
+			RequestHeader:      map[string]string{"X-aws-ec2-metadata-token-ttl-seconds": "21600"},
+			ResponseHeader:     map[string]string{"Server": "EC2ws"},
 			ExpectedStatusCode: http.StatusOK,
 		},
 	}
@@ -72,7 +87,17 @@ func (c *Client) GetProvider() string {
 }
 
 func checkProvider(hc *http.Client, provider Provider) bool {
-	rs, err := hc.Get(metadataServerURL + provider.Path)
+
+	rq, err := http.NewRequest(provider.Method, metadataServerURL+provider.Path, nil)
+	if err != nil {
+		return false
+	}
+
+	for k, v := range provider.RequestHeader {
+		rq.Header.Add(k, v)
+	}
+
+	rs, err := hc.Do(rq)
 	if err != nil {
 		return false
 	}
@@ -81,7 +106,7 @@ func checkProvider(hc *http.Client, provider Provider) bool {
 		return false
 	}
 
-	for k, v := range provider.Header {
+	for k, v := range provider.ResponseHeader {
 		header := rs.Header.Get(k)
 		if header != v {
 			logger.Debug("header %s does not match expected value %s, got %s", k, v, header)

pkg/autodetect/autodetect_test.go

@@ -26,6 +26,10 @@ func newMetadataServer(t *testing.T, testServer *testServer) *httptest.Server {
 	case "eks":
 		// AWS EKS
 		mux.HandleFunc("/", testServer.eksHandler)
+
+	case "eks-imdsv2":
+		// AWS EKS when using IMDSv2
+		mux.HandleFunc("/latest/api/token", testServer.eksImdsv2Handler)
 	}
 
 	return httptest.NewServer(mux)
@@ -80,6 +84,14 @@ func (s *testServer) eksHandler(w http.ResponseWriter, r *http.Request) {
 	w.Write([]byte(responseData))
 }
 
+func (s *testServer) eksImdsv2Handler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Server", "EC2ws")
+
+	responseData := `AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==`
+
+	w.Write([]byte(responseData))
+}
+
 func TestAutoDetectGKE(t *testing.T) {
 	ts := &testServer{
 		t:             t,
@@ -179,6 +191,28 @@ func TestAutoDetectEKS(t *testing.T) {
 
 }
 
+func TestAutoDetectEKSImdsV2(t *testing.T) {
+	ts := &testServer{
+		t:             t,
+		cloudProvider: "eks-imdsv2",
+	}
+
+	server := newMetadataServer(t, ts)
+	defer server.Close()
+
+	metadataServerURL = server.URL
+
+	pc, err := New(&http.Client{}, nil)
+	if err != nil {
+		t.Errorf("err: %v", err)
+	}
+
+	provider := pc.GetProvider()
+
+	assert.Equal(t, "eks-imdsv2", provider, "should be equal")
+
+}
+
 func TestCheckProviderEKS(t *testing.T) {
 	ts := &testServer{
 		t:             t,