feat: Populate homeAZ on nodeInfo CRD (#4009)

* feat: Populate HomeAZ on nodeInfo CRD

* PR comments

* Lint

Author: carlotaarvela

.pipelines/build/dockerfiles/cns.Dockerfile

@@ -11,11 +11,11 @@ ENTRYPOINT ["azure-cns.exe"]
 EXPOSE 10090
 
 # mcr.microsoft.com/azurelinux/base/core:3.0
-FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/base/core@sha256:d472a34802cd535b24ed5fbb7869456e6d8ab2c087faedb9cb32a0efe5b67a15 AS build-helper
+FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/base/core@sha256:833693619d523c23b1fe4d9c1f64a6c697e2a82f7a6ee26e1564897c3fe3fa02 AS build-helper
 RUN tdnf install -y iptables
 
 # mcr.microsoft.com/azurelinux/distroless/minimal:3.0
-FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/distroless/minimal@sha256:77854f8f49c481de03b8c98a5cfba5066616ca5a0213e2f7d443eb542d0f64c4 AS linux
+FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/distroless/minimal@sha256:d784c8233e87e8bce2e902ff59a91262635e4cabc25ec55ac0a718344514db3c AS linux
 ARG ARTIFACT_DIR .
 
 COPY --from=build-helper /usr/sbin/*tables* /usr/sbin/

cni/Dockerfile

@@ -6,10 +6,10 @@ ARG OS_VERSION
 ARG OS
 
 # mcr.microsoft.com/oss/go/microsoft/golang:1.24-azurelinux3.0
-FROM --platform=linux/${ARCH} mcr.microsoft.com/oss/go/microsoft/golang@sha256:6623b87c05c87cf3a213d67f8e917a820227b7fe64582d16deaa8b486a37ef8d AS go
+FROM --platform=linux/${ARCH} mcr.microsoft.com/oss/go/microsoft/golang@sha256:281d086598c336073a59d32cd6fc614a892e90c0c0b881e5051d014859739f0e AS go
 
 # mcr.microsoft.com/azurelinux/base/core:3.0
-FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/base/core@sha256:d472a34802cd535b24ed5fbb7869456e6d8ab2c087faedb9cb32a0efe5b67a15 AS mariner-core
+FROM --platform=linux/${ARCH} mcr.microsoft.com/azurelinux/base/core@sha256:833693619d523c23b1fe4d9c1f64a6c697e2a82f7a6ee26e1564897c3fe3fa02 AS mariner-core
 
 FROM go AS azure-vnet
 ARG OS

cns/Dockerfile

@@ -5,13 +5,13 @@ ARG OS_VERSION
 ARG OS
 
 # mcr.microsoft.com/oss/go/microsoft/golang:1.24-azurelinux3.0
-FROM --platform=linux/${ARCH} mcr.microsoft.com/oss/go/microsoft/golang@sha256:6623b87c05c87cf3a213d67f8e917a820227b7fe64582d16deaa8b486a37ef8d AS go
+FROM --platform=linux/${ARCH} mcr.microsoft.com/oss/go/microsoft/golang@sha256:281d086598c336073a59d32cd6fc614a892e90c0c0b881e5051d014859739f0e AS go
 
 # mcr.microsoft.com/azurelinux/base/core:3.0
-FROM mcr.microsoft.com/azurelinux/base/core@sha256:d472a34802cd535b24ed5fbb7869456e6d8ab2c087faedb9cb32a0efe5b67a15 AS mariner-core
+FROM mcr.microsoft.com/azurelinux/base/core@sha256:833693619d523c23b1fe4d9c1f64a6c697e2a82f7a6ee26e1564897c3fe3fa02 AS mariner-core
 
 # mcr.microsoft.com/azurelinux/distroless/minimal:3.0
-FROM mcr.microsoft.com/azurelinux/distroless/minimal@sha256:77854f8f49c481de03b8c98a5cfba5066616ca5a0213e2f7d443eb542d0f64c4 AS mariner-distroless
+FROM mcr.microsoft.com/azurelinux/distroless/minimal@sha256:d784c8233e87e8bce2e902ff59a91262635e4cabc25ec55ac0a718344514db3c AS mariner-distroless
 
 FROM --platform=linux/${ARCH} go AS builder
 ARG OS

cns/service/main.go

@@ -118,6 +118,9 @@ const (
 	initialIBNICCount                = 0
 )
 
+// ErrHomeAzNotAvailable indicates that HomeAZ information is not available from CNS
+var ErrHomeAzNotAvailable = errors.New("home AZ not available from CNS")
+
 type cniConflistScenario string
 
 const (
@@ -1691,15 +1694,31 @@ func getPodInfoByIPProvider(
 	return podInfoByIPProvider, nil
 }
 
-// createOrUpdateNodeInfoCRD polls imds to learn the VM Unique ID and then creates or updates the NodeInfo CRD
-// with that vm unique ID
+// createOrUpdateNodeInfoCRD polls IMDS to learn the VM Unique ID and CNS to get the HomeAZ,
+// then creates or updates the NodeInfo CRD with that information
 func createOrUpdateNodeInfoCRD(ctx context.Context, restConfig *rest.Config, node *corev1.Node) error {
 	imdsCli := imds.NewClient()
 	vmUniqueID, err := imdsCli.GetVMUniqueID(ctx)
 	if err != nil {
 		return errors.Wrap(err, "error getting vm unique ID from imds")
 	}
 
+	cnsClient, err := cnsclient.New("", cnsReqTimeout)
+	if err != nil {
+		return errors.Wrap(err, "error creating CNS client")
+	}
+	var homeAZ string
+	homeAzResponse, err := cnsClient.GetHomeAz(ctx)
+	if err != nil {
+		return errors.Wrap(err, "error getting home AZ from CNS")
+	}
+	if homeAzResponse.Response.ReturnCode == cnstypes.Success && homeAzResponse.HomeAzResponse.IsSupported {
+		homeAZ = fmt.Sprintf("AZ%02d", homeAzResponse.HomeAzResponse.HomeAz)
+	} else {
+		return errors.Wrapf(ErrHomeAzNotAvailable, "ReturnCode=%d (expected=%d), IsSupported=%t",
+			homeAzResponse.Response.ReturnCode, cnstypes.Success, homeAzResponse.HomeAzResponse.IsSupported)
+	}
+
 	directcli, err := client.New(restConfig, client.Options{Scheme: multitenancy.Scheme})
 	if err != nil {
 		return errors.Wrap(err, "failed to create ctrl client")
@@ -1715,6 +1734,7 @@ func createOrUpdateNodeInfoCRD(ctx context.Context, restConfig *rest.Config, nod
 		},
 		Spec: mtv1alpha1.NodeInfoSpec{
 			VMUniqueID: vmUniqueID,
+			HomeAZ:     homeAZ,
 		},
 	}
 

cns/service/main_test.go

@@ -3,14 +3,24 @@ package main
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/Azure/azure-container-networking/cns"
 	"github.com/Azure/azure-container-networking/cns/fakes"
 	"github.com/Azure/azure-container-networking/cns/logger"
+	"github.com/Azure/azure-container-networking/crd/multitenancy/api/v1alpha1"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/rest"
 )
 
 // MockHTTPClient is a mock implementation of HTTPClient
@@ -69,3 +79,184 @@ func TestSendRegisterNodeRequest_StatusAccepted(t *testing.T) {
 
 	assert.Error(t, sendRegisterNodeRequest(ctx, mockClient, httpServiceFake, nodeRegisterReq, url))
 }
+
+func TestCreateOrUpdateNodeInfoCRD_PopulatesHomeAZ(t *testing.T) {
+	vmID := "test-vm-unique-id-12345"
+	homeAZ := uint(2)
+	HomeAZStr := fmt.Sprintf("AZ0%d", homeAZ)
+
+	// Create mock IMDS server
+	mockIMDSServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.Contains(r.URL.Path, "/metadata/instance/compute") {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			response := map[string]interface{}{
+				"vmId":              vmID,
+				"name":              "test-vm",
+				"resourceGroupName": "test-rg",
+			}
+			_ = json.NewEncoder(w).Encode(response)
+			return
+		}
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer mockIMDSServer.Close()
+
+	// Create mock CNS server
+	mockCNSServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.Contains(r.URL.Path, "/homeaz") || strings.Contains(r.URL.Path, "homeaz") {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			response := map[string]interface{}{
+				"ReturnCode": 0,
+				"Message":    "",
+				"HomeAzResponse": map[string]interface{}{
+					"IsSupported": true,
+					"HomeAz":      homeAZ,
+				},
+			}
+			_ = json.NewEncoder(w).Encode(response)
+			return
+		}
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer mockCNSServer.Close()
+
+	// Set up HTTP transport to mock IMDS and CNS
+	originalTransport := http.DefaultTransport
+	defer func() { http.DefaultTransport = originalTransport }()
+
+	http.DefaultTransport = &mockTransport{
+		imdsServer: mockIMDSServer,
+		cnsServer:  mockCNSServer,
+		original:   originalTransport,
+	}
+
+	// Create a mock Kubernetes server that captures the NodeInfo being created
+	var capturedNodeInfo *v1alpha1.NodeInfo
+
+	mockK8sServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Handle specific API group discovery - multitenancy.acn.azure.com
+		if r.URL.Path == "/apis/multitenancy.acn.azure.com/v1alpha1" && r.Method == "GET" {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusOK)
+			_ = json.NewEncoder(w).Encode(map[string]interface{}{
+				"kind":         "APIResourceList",
+				"groupVersion": "multitenancy.acn.azure.com/v1alpha1",
+				"resources": []map[string]interface{}{
+					{
+						"name":         "nodeinfos",
+						"singularName": "nodeinfo",
+						"namespaced":   false,
+						"kind":         "NodeInfo",
+						"verbs":        []string{"create", "delete", "get", "list", "patch", "update", "watch"},
+					},
+				},
+			})
+			return
+		}
+
+		// Handle NodeInfo resource requests
+		if strings.Contains(r.URL.Path, "nodeinfos") || strings.Contains(r.URL.Path, "multitenancy") {
+			if r.Method == "POST" || r.Method == "PATCH" || r.Method == "PUT" {
+				body, _ := io.ReadAll(r.Body)
+
+				// Try to parse the NodeInfo from the request
+				var nodeInfo v1alpha1.NodeInfo
+				if err := json.Unmarshal(body, &nodeInfo); err == nil {
+					capturedNodeInfo = &nodeInfo
+				}
+
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				// Return the created NodeInfo
+				_ = json.NewEncoder(w).Encode(map[string]interface{}{
+					"apiVersion": "multitenancy.acn.azure.com/v1alpha1",
+					"kind":       "NodeInfo",
+					"metadata": map[string]interface{}{
+						"name": "test-node",
+					},
+					"spec": map[string]interface{}{
+						"vmUniqueID": vmID,
+						"homeAZ":     HomeAZStr,
+					},
+				})
+				return
+			}
+		}
+
+		// Default success response for any other API calls
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"kind":   "Status",
+			"status": "Success",
+		})
+	}))
+	defer mockK8sServer.Close()
+
+	// Test the function with mocked dependencies
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// Point to our mock Kubernetes server
+	restConfig := &rest.Config{
+		Host: mockK8sServer.URL,
+	}
+
+	node := &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: "test-node"},
+	}
+
+	// Call the createOrUpdateNodeInfoCRD function
+	err := createOrUpdateNodeInfoCRD(ctx, restConfig, node)
+
+	// Verify the function succeeded
+	require.NoError(t, err, "Function should succeed with mocked dependencies")
+
+	// Verify the captured values
+	assert.NotNil(t, capturedNodeInfo, "NodeInfo should have been captured from K8s API call")
+	if capturedNodeInfo != nil {
+		assert.Equal(t, vmID, capturedNodeInfo.Spec.VMUniqueID, "VMUniqueID should be from IMDS")
+		assert.Equal(t, HomeAZStr, capturedNodeInfo.Spec.HomeAZ, "HomeAZ should be formatted from CNS response")
+	}
+}
+
+// mockTransport redirects HTTP requests to mock servers for testing.
+// It intercepts requests to IMDS and CNS endpoints and routes them to local test servers.
+type mockTransport struct {
+	imdsServer *httptest.Server
+	cnsServer  *httptest.Server
+	original   http.RoundTripper
+}
+
+func (m *mockTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	// Redirect IMDS calls to mock IMDS server
+	if req.URL.Host == "169.254.169.254" {
+		req.URL.Scheme = "http"
+		req.URL.Host = strings.TrimPrefix(m.imdsServer.URL, "http://")
+		resp, err := m.original.RoundTrip(req)
+		if err != nil {
+			return nil, fmt.Errorf("IMDS mock transport failed: %w", err)
+		}
+		return resp, nil
+	}
+
+	// Redirect CNS calls to mock CNS server
+	if req.URL.Host == "localhost:10090" || strings.Contains(req.URL.Host, "10090") {
+		req.URL.Scheme = "http"
+		req.URL.Host = strings.TrimPrefix(m.cnsServer.URL, "http://")
+		resp, err := m.original.RoundTrip(req)
+		if err != nil {
+			return nil, fmt.Errorf("CNS mock transport failed: %w", err)
+		}
+		return resp, nil
+	}
+
+	// All other calls go through original transport
+	resp, err := m.original.RoundTrip(req)
+	if err != nil {
+		return nil, fmt.Errorf("mock transport failed: %w", err)
+	}
+	return resp, nil
+}