feat: add imds client (#2537)

thatmattlong · web-flow · commit ab5dfe0de94c · 2024-01-24T00:59:33.000Z
* feat: add imds client

* fix: lint
diff --git a/cns/imds/client.go b/cns/imds/client.go
@@ -0,0 +1,132 @@
+// Copyright 2024 Microsoft. All rights reserved.
+// MIT License
+
+package imds
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/url"
+
+	"github.com/avast/retry-go/v4"
+	"github.com/pkg/errors"
+)
+
+// see docs for IMDS here: https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service
+
+// Client returns metadata about the VM by querying IMDS
+type Client struct {
+	cli    *http.Client
+	config clientConfig
+}
+
+// clientConfig holds config options for a Client
+type clientConfig struct {
+	endpoint      string
+	retryAttempts uint
+}
+
+type ClientOption func(*clientConfig)
+
+// Endpoint overrides the default endpoint for a Client
+func Endpoint(endpoint string) ClientOption {
+	return func(c *clientConfig) {
+		c.endpoint = endpoint
+	}
+}
+
+// RetryAttempts overrides the default retry attempts for the client
+func RetryAttempts(attempts uint) ClientOption {
+	return func(c *clientConfig) {
+		c.retryAttempts = attempts
+	}
+}
+
+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"
+)
+
+var (
+	ErrVMUniqueIDNotFound   = errors.New("vm unique ID not found")
+	ErrUnexpectedStatusCode = errors.New("imds returned an unexpected status code")
+)
+
+// NewClient creates a new imds client
+func NewClient(opts ...ClientOption) *Client {
+	config := clientConfig{
+		endpoint: defaultIMDSEndpoint,
+	}
+
+	for _, o := range opts {
+		o(&config)
+	}
+
+	return &Client{
+		cli:    &http.Client{},
+		config: config,
+	}
+}
+
+func (c *Client) GetVMUniqueID(ctx context.Context) (string, error) {
+	var vmUniqueID string
+	err := retry.Do(func() error {
+		computeDoc, err := c.getInstanceComputeMetadata(ctx)
+		if err != nil {
+			return errors.Wrap(err, "error getting IMDS compute metadata")
+		}
+		vmUniqueIDUntyped := computeDoc[vmUniqueIDProperty]
+		var ok bool
+		vmUniqueID, ok = vmUniqueIDUntyped.(string)
+		if !ok {
+			return errors.New("unable to parse IMDS compute metadata, vmId property is not a string")
+		}
+		return nil
+	}, retry.Context(ctx), retry.Attempts(c.config.retryAttempts), retry.DelayType(retry.BackOffDelay))
+	if err != nil {
+		return "", errors.Wrap(err, "exhausted retries querying IMDS compute metadata")
+	}
+
+	if vmUniqueID == "" {
+		return "", ErrVMUniqueIDNotFound
+	}
+
+	return vmUniqueID, nil
+}
+
+func (c *Client) getInstanceComputeMetadata(ctx context.Context) (map[string]any, error) {
+	imdsComputeURL, err := url.JoinPath(c.config.endpoint, imdsComputePath)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to build path to IMDS compute metadata")
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imdsComputeURL, http.NoBody)
+	if err != nil {
+		return nil, errors.Wrap(err, "error building IMDS http request")
+	}
+
+	// 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")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, errors.Wrapf(ErrUnexpectedStatusCode, "unexpected status code %d", resp.StatusCode)
+	}
+
+	var m map[string]any
+	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
+		return nil, errors.Wrap(err, "error decoding IMDS response as json")
+	}
+
+	return m, nil
+}
diff --git a/cns/imds/client_test.go b/cns/imds/client_test.go
@@ -0,0 +1,70 @@
+// Copyright 2024 Microsoft. All rights reserved.
+// MIT License
+
+package imds_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/Azure/azure-container-networking/cns/imds"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetVMUniqueID(t *testing.T) {
+	computeMetadata, err := os.ReadFile("testdata/computeMetadata.json")
+	require.NoError(t, err, "error reading testdata compute metadata file")
+
+	mockIMDSServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// request header "Metadata: true" must be present
+		metadataHeader := r.Header.Get("Metadata")
+		assert.Equal(t, "true", metadataHeader)
+		w.WriteHeader(http.StatusOK)
+		_, writeErr := w.Write(computeMetadata)
+		require.NoError(t, writeErr, "error writing response")
+	}))
+	defer mockIMDSServer.Close()
+
+	imdsClient := imds.NewClient(imds.Endpoint(mockIMDSServer.URL))
+	vmUniqueID, err := imdsClient.GetVMUniqueID(context.Background())
+	require.NoError(t, err, "error querying testserver")
+
+	require.Equal(t, "55b8499d-9b42-4f85-843f-24ff69f4a643", vmUniqueID)
+}
+
+func TestGetVMUniqueIDInvalidEndpoint(t *testing.T) {
+	imdsClient := imds.NewClient(imds.Endpoint(string([]byte{0x7f})), imds.RetryAttempts(1))
+	_, err := imdsClient.GetVMUniqueID(context.Background())
+	require.Error(t, err, "expected invalid path")
+}
+
+func TestIMDSInternalServerError(t *testing.T) {
+	mockIMDSServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// request header "Metadata: true" must be present
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer mockIMDSServer.Close()
+
+	imdsClient := imds.NewClient(imds.Endpoint(mockIMDSServer.URL), imds.RetryAttempts(1))
+
+	_, err := imdsClient.GetVMUniqueID(context.Background())
+	require.ErrorIs(t, err, imds.ErrUnexpectedStatusCode, "expected internal server error")
+}
+
+func TestIMDSInvalidJSON(t *testing.T) {
+	mockIMDSServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, err := w.Write([]byte("not json"))
+		require.NoError(t, err)
+	}))
+	defer mockIMDSServer.Close()
+
+	imdsClient := imds.NewClient(imds.Endpoint(mockIMDSServer.URL), imds.RetryAttempts(1))
+
+	_, err := imdsClient.GetVMUniqueID(context.Background())
+	require.Error(t, err, "expected json decoding error")
+}
diff --git a/cns/imds/testdata/computeMetadata.json b/cns/imds/testdata/computeMetadata.json
@@ -0,0 +1,130 @@
+{
+    "azEnvironment": "AzurePublicCloud",
+    "customData": "",
+    "evictionPolicy": "",
+    "isHostCompatibilityLayerVm": "false",
+    "licenseType": "",
+    "location": "westus2",
+    "name": "aks-nodepool1-25781205-vmss_0",
+    "offer": "",
+    "osProfile": {
+        "adminUsername": "azureuser",
+        "computerName": "aks-nodepool1-25781205-vmss000000",
+        "disablePasswordAuthentication": "true"
+    },
+    "osType": "Linux",
+    "placementGroupId": "078e7cc3-c76f-46d2-91ff-64c4dbdd0d7c",
+    "plan": {
+        "name": "",
+        "product": "",
+        "publisher": ""
+    },
+    "platformFaultDomain": "0",
+    "platformUpdateDomain": "0",
+    "priority": "",
+    "provider": "Microsoft.Compute",
+    "publicKeys": [],
+    "publisher": "",
+    "resourceGroupName": "MC_matlong-test-imds_matlong-test-imds_westus2",
+    "resourceId": "/subscriptions/9b8218f9-902a-4d20-a65c-e98acec5362f/resourceGroups/MC_matlong-test-imds_matlong-test-imds_westus2/providers/Microsoft.Compute/virtualMachineScaleSets/aks-nodepool1-25781205-vmss/virtualMachines/0",
+    "securityProfile": {
+        "secureBootEnabled": "false",
+        "virtualTpmEnabled": "false"
+    },
+    "sku": "",
+    "storageProfile": {
+        "dataDisks": [],
+        "imageReference": {
+            "id": "/subscriptions/109a5e88-712a-48ae-9078-9ca8b3c81345/resourceGroups/AKS-Ubuntu/providers/Microsoft.Compute/galleries/AKSUbuntu/images/2204gen2containerd/versions/202401.09.0",
+            "offer": "",
+            "publisher": "",
+            "sku": "",
+            "version": ""
+        },
+        "osDisk": {
+            "caching": "ReadWrite",
+            "createOption": "FromImage",
+            "diffDiskSettings": {
+                "option": ""
+            },
+            "diskSizeGB": "128",
+            "encryptionSettings": {
+                "enabled": "false"
+            },
+            "image": {
+                "uri": ""
+            },
+            "managedDisk": {
+                "id": "/subscriptions/9b8218f9-902a-4d20-a65c-e98acec5362f/resourceGroups/MC_matlong-test-imds_matlong-test-imds_westus2/providers/Microsoft.Compute/disks/aks-nodepool1-257812aks-nodepool1-2578120disk1_7f5fbaaac4d7480b9f78b5da0206643c",
+                "storageAccountType": "Premium_LRS"
+            },
+            "name": "aks-nodepool1-257812aks-nodepool1-2578120disk1_7f5fbaaac4d7480b9f78b5da0206643c",
+            "osType": "Linux",
+            "vhd": {
+                "uri": ""
+            },
+            "writeAcceleratorEnabled": "false"
+        },
+        "resourceDisk": {
+            "size": "14336"
+        }
+    },
+    "subscriptionId": "9b8218f9-902a-4d20-a65c-e98acec5362f",
+    "tags": "aks-managed-azure-cni-overlay:true;aks-managed-consolidated-additional-properties:d9dfbd53-b974-11ee-bba9-b29ebea78e50;aks-managed-coordination:true;aks-managed-createOperationID:f9031ed6-4bd9-47e8-908f-90c6429c0790;aks-managed-creationSource:vmssclient-aks-nodepool1-25781205-vmss;aks-managed-kubeletIdentityClientID:58ada700-3bba-437e-b6a2-71df9a5148b5;aks-managed-orchestrator:Kubernetes:1.27.7;aks-managed-poolName:nodepool1;aks-managed-resourceNameSuffix:22973824;aks-managed-ssh-access:LocalUser;azsecpack:nonprod;platformsettings.host_environment.service.platform_optedin_for_rootcerts:true",
+    "tagsList": [
+        {
+            "name": "aks-managed-azure-cni-overlay",
+            "value": "true"
+        },
+        {
+            "name": "aks-managed-consolidated-additional-properties",
+            "value": "d9dfbd53-b974-11ee-bba9-b29ebea78e50"
+        },
+        {
+            "name": "aks-managed-coordination",
+            "value": "true"
+        },
+        {
+            "name": "aks-managed-createOperationID",
+            "value": "f9031ed6-4bd9-47e8-908f-90c6429c0790"
+        },
+        {
+            "name": "aks-managed-creationSource",
+            "value": "vmssclient-aks-nodepool1-25781205-vmss"
+        },
+        {
+            "name": "aks-managed-kubeletIdentityClientID",
+            "value": "58ada700-3bba-437e-b6a2-71df9a5148b5"
+        },
+        {
+            "name": "aks-managed-orchestrator",
+            "value": "Kubernetes:1.27.7"
+        },
+        {
+            "name": "aks-managed-poolName",
+            "value": "nodepool1"
+        },
+        {
+            "name": "aks-managed-resourceNameSuffix",
+            "value": "22973824"
+        },
+        {
+            "name": "aks-managed-ssh-access",
+            "value": "LocalUser"
+        },
+        {
+            "name": "azsecpack",
+            "value": "nonprod"
+        },
+        {
+            "name": "platformsettings.host_environment.service.platform_optedin_for_rootcerts",
+            "value": "true"
+        }
+    ],
+    "userData": "",
+    "version": "202401.09.0",
+    "vmId": "55b8499d-9b42-4f85-843f-24ff69f4a643",
+    "vmScaleSetName": "aks-nodepool1-25781205-vmss",
+    "vmSize": "Standard_DS2_v2",
+    "zone": ""
+}
