From ad50b11fdcdb7016948e9dfaf51b6af361a35432 Mon Sep 17 00:00:00 2001
From: Matthew Long <61910737+thatmattlong@users.noreply.github.com>
Date: Wed, 7 Feb 2024 12:31:16 -0800
Subject: [PATCH 1/2] fix: build imds compute url correctly

---
 cns/imds/client.go      | 22 +++++++++++++++-------
 cns/imds/client_test.go |  6 ++++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/cns/imds/client.go b/cns/imds/client.go
index 5ab43e76ba..a0cc678bb3 100644
--- a/cns/imds/client.go
+++ b/cns/imds/client.go
@@ -6,7 +6,9 @@ package imds
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
+	"net/http/httputil"
 	"net/url"
 
 	"github.com/avast/retry-go/v4"
@@ -44,12 +46,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 +64,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 +109,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 {
@@ -113,6 +119,8 @@ func (c *Client) getInstanceComputeMetadata(ctx context.Context) (map[string]any
 	// IMDS requires the "Metadata: true" header
 	req.Header.Add(metadataHeaderKey, metadataHeaderValue)
 
+	reqDump, _ := httputil.DumpRequestOut(req, false)
+	fmt.Printf("REQUEST:\n%s", string(reqDump))
 	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")

From 135d47b36161bb8a3c1801796c36529827699aef Mon Sep 17 00:00:00 2001
From: Matthew Long <61910737+thatmattlong@users.noreply.github.com>
Date: Wed, 7 Feb 2024 12:33:58 -0800
Subject: [PATCH 2/2] fix: remove httputil debug stuff

---
 cns/imds/client.go | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/cns/imds/client.go b/cns/imds/client.go
index a0cc678bb3..6e210e4705 100644
--- a/cns/imds/client.go
+++ b/cns/imds/client.go
@@ -6,9 +6,7 @@ package imds
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"net/http"
-	"net/http/httputil"
 	"net/url"
 
 	"github.com/avast/retry-go/v4"
@@ -118,9 +116,6 @@ func (c *Client) getInstanceComputeMetadata(ctx context.Context) (map[string]any
 
 	// IMDS requires the "Metadata: true" header
 	req.Header.Add(metadataHeaderKey, metadataHeaderValue)
-
-	reqDump, _ := httputil.DumpRequestOut(req, false)
-	fmt.Printf("REQUEST:\n%s", string(reqDump))
 	resp, err := c.cli.Do(req)
 	if err != nil {
 		return nil, errors.Wrap(err, "error querying IMDS")