From f97eb963554a8883f2d0e19f326f4a99f8b64a60 Mon Sep 17 00:00:00 2001
From: Matthew Long <61910737+thatmattlong@users.noreply.github.com>
Date: Thu, 8 Feb 2024 08:32:36 -0800
Subject: [PATCH] fix: build imds compute url correctly and give a default
 retryAttempts (#2566)

* fix: build imds compute url correctly

* fix: remove httputil debug stuff
---
 cns/imds/client.go      | 19 +++++++++++--------
 cns/imds/client_test.go |  6 ++++++
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/cns/imds/client.go b/cns/imds/client.go
index 5ab43e76ba..6e210e4705 100644
--- a/cns/imds/client.go
+++ b/cns/imds/client.go
@@ -44,12 +44,14 @@ func RetryAttempts(attempts uint) ClientOption {
 }
 
 const (
-	vmUniqueIDProperty   = "vmId"
-	imdsComputePath      = "/metadata/instance/compute?api-version=2021-01-01&format=json"
-	metadataHeaderKey    = "Metadata"
-	metadataHeaderValue  = "true"
-	defaultRetryAttempts = 10
-	defaultIMDSEndpoint  = "http://169.254.169.254"
+	vmUniqueIDProperty    = "vmId"
+	imdsComputePath       = "/metadata/instance/compute"
+	imdsComputeAPIVersion = "api-version=2021-01-01"
+	imdsFormatJSON        = "format=json"
+	metadataHeaderKey     = "Metadata"
+	metadataHeaderValue   = "true"
+	defaultRetryAttempts  = 3
+	defaultIMDSEndpoint   = "http://169.254.169.254"
 )
 
 var (
@@ -60,7 +62,8 @@ var (
 // NewClient creates a new imds client
 func NewClient(opts ...ClientOption) *Client {
 	config := clientConfig{
-		endpoint: defaultIMDSEndpoint,
+		endpoint:      defaultIMDSEndpoint,
+		retryAttempts: defaultRetryAttempts,
 	}
 
 	for _, o := range opts {
@@ -104,6 +107,7 @@ func (c *Client) getInstanceComputeMetadata(ctx context.Context) (map[string]any
 	if err != nil {
 		return nil, errors.Wrap(err, "unable to build path to IMDS compute metadata")
 	}
+	imdsComputeURL = imdsComputeURL + "?" + imdsComputeAPIVersion + "&" + imdsFormatJSON
 
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imdsComputeURL, http.NoBody)
 	if err != nil {
@@ -112,7 +116,6 @@ func (c *Client) getInstanceComputeMetadata(ctx context.Context) (map[string]any
 
 	// IMDS requires the "Metadata: true" header
 	req.Header.Add(metadataHeaderKey, metadataHeaderValue)
-
 	resp, err := c.cli.Do(req)
 	if err != nil {
 		return nil, errors.Wrap(err, "error querying IMDS")
diff --git a/cns/imds/client_test.go b/cns/imds/client_test.go
index c0c03b40a3..fe8324666f 100644
--- a/cns/imds/client_test.go
+++ b/cns/imds/client_test.go
@@ -23,6 +23,12 @@ func TestGetVMUniqueID(t *testing.T) {
 		// request header "Metadata: true" must be present
 		metadataHeader := r.Header.Get("Metadata")
 		assert.Equal(t, "true", metadataHeader)
+
+		// query params should include apiversion and json format
+		apiVersion := r.URL.Query().Get("api-version")
+		assert.Equal(t, "2021-01-01", apiVersion)
+		format := r.URL.Query().Get("format")
+		assert.Equal(t, "json", format)
 		w.WriteHeader(http.StatusOK)
 		_, writeErr := w.Write(computeMetadata)
 		require.NoError(t, writeErr, "error writing response")