fix: use cloud-specific ARM endpoint for IMDS token in ORAS login

aks-node-controller/parser/helper.go

@@ -569,6 +569,37 @@ func getTargetCloud(v *aksnodeconfigv1.Configuration) string {
 	return getTargetEnvironment(v)
 }
 
+// getArmResourceEndpoint returns the ARM resource endpoint to use as the IMDS
+// "resource" parameter when acquiring an AAD token.
+//   - For AKS custom clouds (Azure Stack): sourced from
+//     CustomEnvJsonContent.resourceManagerEndpoint, which is populated by AKS RP.
+//   - For public sovereign clouds (Fairfax / Mooncake): mapped by cloud name.
+//     These endpoints are public knowledge so hardcoding is acceptable.
+//   - For Azure public cloud (and any unknown): returns empty; cse_helpers.sh
+//     defaults to https://management.azure.com/.
+func getArmResourceEndpoint(v *aksnodeconfigv1.Configuration) string {
+	if getIsAksCustomCloud(v.GetCustomCloudConfig()) {
+		raw := v.GetCustomCloudConfig().GetCustomEnvJsonContent()
+		if raw == "" {
+			return ""
+		}
+		var env struct {
+			ResourceManagerEndpoint string `json:"resourceManagerEndpoint"`
+		}
+		if err := json.Unmarshal([]byte(raw), &env); err != nil {
+			return ""
+		}
+		return env.ResourceManagerEndpoint
+	}
+	switch getCloudTargetEnv(v) {
+	case "AzureUSGovernmentCloud":
+		return "https://management.usgovcloudapi.net/"
+	case "AzureChinaCloud":
+		return "https://management.chinacloudapi.cn/"
+	}
+	return ""
+}
+
 func getAzureEnvironmentFilepath(v *aksnodeconfigv1.Configuration) string {
 	if getIsAksCustomCloud(v.GetCustomCloudConfig()) {
 		return fmt.Sprintf("/etc/kubernetes/%s.json", getTargetEnvironment(v))

aks-node-controller/parser/helper_test.go

@@ -1185,6 +1185,77 @@ func Test_getTargetCloud(t *testing.T) {
 	}
 }
 
+func Test_getArmResourceEndpoint(t *testing.T) {
+	tests := []struct {
+		name string
+		v    *aksnodeconfigv1.Configuration
+		want string
+	}{
+		{
+			name: "Nil config returns empty (public cloud default)",
+			v:    &aksnodeconfigv1.Configuration{},
+			want: "",
+		},
+		{
+			name: "Public cloud location returns empty",
+			v: &aksnodeconfigv1.Configuration{
+				ClusterConfig: &aksnodeconfigv1.ClusterConfig{Location: "eastus"},
+			},
+			want: "",
+		},
+		{
+			name: "China cloud by location returns Mooncake endpoint",
+			v: &aksnodeconfigv1.Configuration{
+				ClusterConfig: &aksnodeconfigv1.ClusterConfig{Location: "chinaeast2"},
+			},
+			want: "https://management.chinacloudapi.cn/",
+		},
+		{
+			name: "US Gov by location returns Fairfax endpoint",
+			v: &aksnodeconfigv1.Configuration{
+				ClusterConfig: &aksnodeconfigv1.ClusterConfig{Location: "usgovvirginia"},
+			},
+			want: "https://management.usgovcloudapi.net/",
+		},
+		{
+			name: "AKS custom cloud with resourceManagerEndpoint in CustomEnvJsonContent",
+			v: &aksnodeconfigv1.Configuration{
+				CustomCloudConfig: &aksnodeconfigv1.CustomCloudConfig{
+					CustomCloudEnvName:   helpers.AksCustomCloudName,
+					CustomEnvJsonContent: `{"resourceManagerEndpoint":"https://management.azure.microsoft.fakecustomcloud/"}`,
+				},
+			},
+			want: "https://management.azure.microsoft.fakecustomcloud/",
+		},
+		{
+			name: "AKS custom cloud with empty CustomEnvJsonContent returns empty",
+			v: &aksnodeconfigv1.Configuration{
+				CustomCloudConfig: &aksnodeconfigv1.CustomCloudConfig{
+					CustomCloudEnvName: helpers.AksCustomCloudName,
+				},
+			},
+			want: "",
+		},
+		{
+			name: "AKS custom cloud with malformed JSON returns empty",
+			v: &aksnodeconfigv1.Configuration{
+				CustomCloudConfig: &aksnodeconfigv1.CustomCloudConfig{
+					CustomCloudEnvName:   helpers.AksCustomCloudName,
+					CustomEnvJsonContent: `{not-json`,
+				},
+			},
+			want: "",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := getArmResourceEndpoint(tt.v); got != tt.want {
+				t.Errorf("getArmResourceEndpoint() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func Test_getLinuxAdminUsername(t *testing.T) {
 	type args struct {
 		username string

aks-node-controller/parser/parser.go

@@ -111,6 +111,7 @@ func getCSEEnv(config *aksnodeconfigv1.Configuration) map[string]string {
 		"CONTAINERD_ULIMITS":                                   getUlimitContent(config.GetCustomLinuxOsConfig().GetUlimitConfig()),
 		"TARGET_CLOUD":                                         getTargetCloud(config),
 		"TARGET_ENVIRONMENT":                                   getTargetEnvironment(config),
+		"ARM_RESOURCE_ENDPOINT":                                getArmResourceEndpoint(config),
 		"CUSTOM_ENV_JSON":                                      config.GetCustomCloudConfig().GetCustomEnvJsonContent(),
 		"IS_CUSTOM_CLOUD":                                      fmt.Sprintf("%v", getIsAksCustomCloud(config.GetCustomCloudConfig())),
 		"AKS_CUSTOM_CLOUD_CONTAINER_REGISTRY_DNS_SUFFIX":       config.GetCustomCloudConfig().GetContainerRegistryDnsSuffix(),

aks-node-controller/parser/parser_test.go

@@ -130,6 +130,7 @@ oom_score = -999
 				assert.Equal(t, "AzureChinaCloud", vars["TARGET_ENVIRONMENT"])
 				assert.Equal(t, "AzureChinaCloud", vars["TARGET_CLOUD"])
 				assert.Equal(t, "false", vars["IS_CUSTOM_CLOUD"])
+				assert.Equal(t, "https://management.chinacloudapi.cn/", vars["ARM_RESOURCE_ENDPOINT"])
 			},
 		},
 		{
@@ -429,6 +430,7 @@ func TestAKSNodeConfigCompatibilityFromJsonToCSECommand(t *testing.T) {
 				assertHasKeyWithValue(t, vars, "NO_PROXY_URLS", "")
 				assertHasKeyWithValue(t, vars, "HTTP_PROXY_TRUSTED_CA", "")
 				assertHasKeyWithValue(t, vars, "TARGET_ENVIRONMENT", helpers.DefaultCloudName)
+				assertHasKeyWithValue(t, vars, "ARM_RESOURCE_ENDPOINT", "")
 				assertHasKeyWithValue(t, vars, "TLS_BOOTSTRAP_TOKEN", "")
 				assertHasKeyWithValue(t, vars, "ENABLE_SECURE_TLS_BOOTSTRAPPING", "false")
 				assertHasKeyWithValue(t, vars, "SECURE_TLS_BOOTSTRAPPING_AAD_RESOURCE", "")

parts/linux/cloud-init/artifacts/cse_cmd.sh

@@ -114,6 +114,7 @@ CONTAINERD_ULIMITS="{{GetContainerdUlimitString}}"
 {{/* see GetCustomEnvironmentJSON for more weirdness. */}}
 TARGET_CLOUD="{{- if IsAKSCustomCloud -}} AzureStackCloud {{- else -}} {{GetTargetEnvironment}} {{- end -}}"
 TARGET_ENVIRONMENT="{{GetTargetEnvironment}}"
+ARM_RESOURCE_ENDPOINT="{{GetArmResourceEndpoint}}"
 CUSTOM_ENV_JSON="{{GetBase64EncodedEnvironmentJSON}}"
 IS_CUSTOM_CLOUD="{{IsAKSCustomCloud}}"
 AKS_CUSTOM_CLOUD_CONTAINER_REGISTRY_DNS_SUFFIX="{{- if IsAKSCustomCloud}}{{AKSCustomCloudContainerRegistryDNSSuffix}}{{end}}"

parts/linux/cloud-init/artifacts/cse_helpers.sh

@@ -1282,7 +1282,8 @@ oras_login_with_kubelet_identity() {
     fi
 
     set +x
-    access_url="http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https://management.azure.com/&client_id=$client_id"
+    local arm_endpoint="${ARM_RESOURCE_ENDPOINT:-https://management.azure.com/}"
+    access_url="http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=${arm_endpoint}&client_id=$client_id"
     raw_access_token=$(retrycmd_get_aad_access_token 5 15 $access_url)
     ret_code=$?
     if [ "$ret_code" -ne 0 ]; then

parts/windows/kuberneteswindowssetup.ps1

@@ -47,6 +47,7 @@ $Location="{{ GetVariable "location" }}"
 $UserAssignedClientID="{{ GetVariable "userAssignedIdentityID" }}"
 {{ end }}
 $TargetEnvironment="{{ GetTargetEnvironment }}"
+$ArmResourceEndpoint="{{ GetArmResourceEndpoint }}"
 $AADClientId="{{ GetParameter "servicePrincipalClientId" }}"
 $NetworkAPIVersion="2018-08-01"
 

pkg/agent/baker.go

@@ -1016,6 +1016,22 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration
 			}
 			return GetCloudTargetEnv(cs.Location)
 		},
+		"GetArmResourceEndpoint": func() string {
+			// Custom clouds (Azure Stack): RP populates CustomCloudEnv.ResourceManagerEndpoint.
+			if cs.Properties != nil && cs.Properties.CustomCloudEnv != nil {
+				return cs.Properties.CustomCloudEnv.ResourceManagerEndpoint
+			}
+			// Public sovereign clouds (FF/MC) — endpoints are public knowledge so
+			// it's safe to map by cloud name. Public cloud falls through to empty;
+			// scripts default to https://management.azure.com/.
+			switch GetCloudTargetEnv(cs.Location) {
+			case datamodel.AzureUSGovernmentCloud:
+				return "https://management.usgovcloudapi.net/"
+			case datamodel.AzureChinaCloud:
+				return "https://management.chinacloudapi.cn/"
+			}
+			return ""
+		},
 		"IsAKSCustomCloud": func() bool {
 			return cs.IsAKSCustomCloud()
 		},

staging/cse/windows/networkisolatedclusterfunc.ps1

@@ -172,7 +172,8 @@ function Invoke-OrasLogin {
     }
 
     # Get AAD Access Token using Managed Identity Metadata Service
-    $accessUrl = "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https://management.azure.com/&client_id=$ClientID"
+    $armEndpoint = if ([string]::IsNullOrWhiteSpace($ArmResourceEndpoint)) { "https://management.azure.com/" } else { $ArmResourceEndpoint }
+    $accessUrl = "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=${armEndpoint}&client_id=$ClientID"
     try {
         $requestArgs = @{
             Uri     = $accessUrl