feat: enable localdns hosts plugin to cache critical AKS FQDNs

.pipelines/scripts/verify_shell.sh

@@ -30,6 +30,7 @@ filesToCheck=$(find . -type f -name "*.sh" -not -path './pkg/agent/testdata/*' -
 # Known bash-only scripts that intentionally use bash specific syntax.
 BASH_ONLY_LIST=$(cat <<'EOF'
 ./vhdbuilder/packer/install-ig.sh
+./parts/linux/cloud-init/artifacts/aks-hosts-setup.sh
 EOF
 )
 

aks-node-controller/parser/helper.go

@@ -719,11 +719,17 @@ func getFuncMapForLocalDnsCorefileTemplate() template.FuncMap {
 	}
 }
 
-// getLocalDnsCorefileBase64 returns the base64 encoded LocalDns corefile.
-// base64 encoded corefile returned from this function will decoded and written
-// to /opt/azure/containers/localdns/localdns.corefile in cse_config.sh
-// and then used by localdns systemd unit to start localdns systemd unit.
-func getLocalDnsCorefileBase64(aksnodeconfig *aksnodeconfigv1.Configuration) string {
+// getLocalDnsCorefileBase64WithHostsPlugin generates and returns the base64-encoded LocalDns corefile
+// with or without the hosts plugin, depending on the includeHostsPlugin parameter.
+//
+// The generated content is returned as a base64-encoded string and stored in environment variables:
+//   - LOCALDNS_GENERATED_COREFILE (with hosts plugin)
+//   - LOCALDNS_GENERATED_COREFILE_NO_HOSTS (without hosts plugin)
+//
+// The actual file writing happens in shell scripts (cse_config.sh) which decode and write
+// the selected variant to /opt/azure/containers/localdns/localdns.corefile.
+// Runtime selection between variants happens in cse_main.sh based on the availability of /etc/localdns/hosts.
+func getLocalDnsCorefileBase64WithHostsPlugin(aksnodeconfig *aksnodeconfigv1.Configuration, includeHostsPlugin bool) string {
 	if aksnodeconfig == nil {
 		return ""
 	}
@@ -737,17 +743,33 @@ func getLocalDnsCorefileBase64(aksnodeconfig *aksnodeconfigv1.Configuration) str
 		return ""
 	}
 
-	localDnsConfig, err := generateLocalDnsCorefileFromAKSNodeConfig(aksnodeconfig)
+	variant := "with hosts plugin"
+	if !includeHostsPlugin {
+		variant = "without hosts plugin"
+	}
+
+	localDnsConfig, err := generateLocalDnsCorefileFromAKSNodeConfig(aksnodeconfig, includeHostsPlugin)
 	if err != nil {
-		return fmt.Sprintf("error getting localdns corfile from aks node config: %v", err)
+		return fmt.Sprintf("error getting localdns corefile (%s) from aks node config: %v", variant, err)
 	}
 	return base64.StdEncoding.EncodeToString([]byte(localDnsConfig))
 }
 
+// localDnsCorefileTemplateData wraps the AKS node config with additional template control flags.
+type localDnsCorefileTemplateData struct {
+	Config             *aksnodeconfigv1.Configuration
+	IncludeHostsPlugin bool
+}
+
 // Corefile is created using localdns.toml.gtpl template and aksnodeconfig values.
-func generateLocalDnsCorefileFromAKSNodeConfig(aksnodeconfig *aksnodeconfigv1.Configuration) (string, error) {
+// includeHostsPlugin controls whether the hosts plugin block is included in the generated Corefile.
+func generateLocalDnsCorefileFromAKSNodeConfig(aksnodeconfig *aksnodeconfigv1.Configuration, includeHostsPlugin bool) (string, error) {
 	var corefileBuffer bytes.Buffer
-	if err := localDnsCorefileTemplate.Execute(&corefileBuffer, aksnodeconfig); err != nil {
+	templateData := localDnsCorefileTemplateData{
+		Config:             aksnodeconfig,
+		IncludeHostsPlugin: includeHostsPlugin,
+	}
+	if err := localDnsCorefileTemplate.Execute(&corefileBuffer, templateData); err != nil {
 		return "", fmt.Errorf("failed to execute localdns corefile template: %w", err)
 	}
 	return corefileBuffer.String(), nil
@@ -785,6 +807,13 @@ func shouldEnableLocalDns(aksnodeconfig *aksnodeconfigv1.Configuration) string {
 	return fmt.Sprintf("%v", aksnodeconfig != nil && aksnodeconfig.GetLocalDnsProfile() != nil && aksnodeconfig.GetLocalDnsProfile().GetEnableLocalDns())
 }
 
+// shouldEnableHostsPlugin returns true if LocalDNS is enabled and the hosts plugin
+// is explicitly enabled. When true, the localdns Corefile will include a hosts plugin
+// block that serves cached DNS entries from /etc/localdns/hosts for critical AKS FQDNs.
+func shouldEnableHostsPlugin(aksnodeconfig *aksnodeconfigv1.Configuration) string {
+	return fmt.Sprintf("%v", shouldEnableLocalDns(aksnodeconfig) == "true" && aksnodeconfig.GetLocalDnsProfile().GetEnableHostsPlugin())
+}
+
 // getLocalDnsCpuLimitInPercentage returns CPU limit in percentage unit that will be used in localdns systemd unit.
 func getLocalDnsCpuLimitInPercentage(aksnodeconfig *aksnodeconfigv1.Configuration) string {
 	if shouldEnableLocalDns(aksnodeconfig) == "true" && aksnodeconfig.GetLocalDnsProfile().GetCpuLimitInMilliCores() != 0 {

aks-node-controller/parser/helper_test.go

@@ -1446,6 +1446,10 @@ health-check.localdns.local:53 {
 .:53 {
     log
     bind 169.254.10.10
+    # Check /etc/localdns/hosts first for critical AKS FQDNs (mcr.microsoft.com, packages.aks.azure.com, etc.)
+    hosts /etc/localdns/hosts {
+        fallthrough
+    }
     forward . 168.63.129.16 {
         policy sequential
         max_concurrent 1000
@@ -1509,6 +1513,10 @@ testdomain456.com:53 {
 .:53 {
     errors
     bind 169.254.10.11
+    # Check /etc/localdns/hosts first for critical AKS FQDNs (mcr.microsoft.com, packages.aks.azure.com, etc.)
+    hosts /etc/localdns/hosts {
+        fallthrough
+    }
     forward . 10.0.0.10 {
         policy sequential
         max_concurrent 2000
@@ -1627,7 +1635,7 @@ func Test_getLocalDNSCorefileBase64(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := getLocalDnsCorefileBase64(tt.args.aksnodeconfig)
+			got := getLocalDnsCorefileBase64WithHostsPlugin(tt.args.aksnodeconfig, true)
 
 			if tt.wantContains == "" && got != "" {
 				t.Errorf("expected empty string, got %q", got)
@@ -1711,6 +1719,71 @@ func Test_shouldEnableLocalDns(t *testing.T) {
 	}
 }
 
+func Test_shouldEnableHostsPlugin(t *testing.T) {
+	type args struct {
+		aksnodeconfig *aksnodeconfigv1.Configuration
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "nil config",
+			args: args{aksnodeconfig: nil},
+			want: "false",
+		},
+		{
+			name: "nil LocalDnsProfile",
+			args: args{aksnodeconfig: &aksnodeconfigv1.Configuration{}},
+			want: "false",
+		},
+		{
+			name: "LocalDns disabled, HostsPlugin enabled",
+			args: args{aksnodeconfig: &aksnodeconfigv1.Configuration{
+				LocalDnsProfile: &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    false,
+					EnableHostsPlugin: true},
+			}},
+			want: "false",
+		},
+		{
+			name: "LocalDns enabled, HostsPlugin disabled",
+			args: args{aksnodeconfig: &aksnodeconfigv1.Configuration{
+				LocalDnsProfile: &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    true,
+					EnableHostsPlugin: false},
+			}},
+			want: "false",
+		},
+		{
+			name: "both LocalDns and HostsPlugin enabled",
+			args: args{aksnodeconfig: &aksnodeconfigv1.Configuration{
+				LocalDnsProfile: &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    true,
+					EnableHostsPlugin: true},
+			}},
+			want: "true",
+		},
+		{
+			name: "both disabled",
+			args: args{aksnodeconfig: &aksnodeconfigv1.Configuration{
+				LocalDnsProfile: &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    false,
+					EnableHostsPlugin: false},
+			}},
+			want: "false",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := shouldEnableHostsPlugin(tt.args.aksnodeconfig); got != tt.want {
+				t.Errorf("shouldEnableHostsPlugin() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func Test_getLocalDnsCpuLimitInPercentage(t *testing.T) {
 	type args struct {
 		aksnodeconfig *aksnodeconfigv1.Configuration

aks-node-controller/parser/parser.go

@@ -170,9 +170,11 @@ func getCSEEnv(config *aksnodeconfigv1.Configuration) map[string]string {
 		"INSERT_IMDS_RESTRICTION_RULE_TO_MANGLE_TABLE":       fmt.Sprintf("%v", config.GetImdsRestrictionConfig().GetInsertImdsRestrictionRuleToMangleTable()),
 		"PRE_PROVISION_ONLY":                                 fmt.Sprintf("%v", config.GetPreProvisionOnly()),
 		"SHOULD_ENABLE_LOCALDNS":                             shouldEnableLocalDns(config),
+		"SHOULD_ENABLE_HOSTS_PLUGIN":                         shouldEnableHostsPlugin(config),
 		"LOCALDNS_CPU_LIMIT":                                 getLocalDnsCpuLimitInPercentage(config),
 		"LOCALDNS_MEMORY_LIMIT":                              getLocalDnsMemoryLimitInMb(config),
-		"LOCALDNS_GENERATED_COREFILE":                        getLocalDnsCorefileBase64(config),
+		"LOCALDNS_GENERATED_COREFILE":                        getLocalDnsCorefileBase64WithHostsPlugin(config, true),
+		"LOCALDNS_GENERATED_COREFILE_NO_HOSTS":               getLocalDnsCorefileBase64WithHostsPlugin(config, false),
 		"DISABLE_PUBKEY_AUTH":                                fmt.Sprintf("%v", config.GetDisablePubkeyAuth()),
 		"SERVICE_ACCOUNT_IMAGE_PULL_ENABLED":                 fmt.Sprintf("%v", config.GetServiceAccountImagePullProfile().GetEnabled()),
 		"SERVICE_ACCOUNT_IMAGE_PULL_DEFAULT_CLIENT_ID":       config.GetServiceAccountImagePullProfile().GetDefaultClientId(),

aks-node-controller/parser/parser_test.go

@@ -229,6 +229,38 @@ oom_score = -999
 				assert.Equal(t, "true", vars["NEEDS_CGROUPV2"])
 			},
 		},
+		{
+			name:       "AKSUbuntu2204 with LocalDNS and hosts plugin enabled",
+			folder:     "AKSUbuntu2204+LocalDNS+HostsPlugin",
+			k8sVersion: "1.24.2",
+			aksNodeConfigUpdator: func(aksNodeConfig *aksnodeconfigv1.Configuration) {
+				aksNodeConfig.LocalDnsProfile = &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    true,
+					EnableHostsPlugin: true,
+				}
+			},
+			validator: func(cmd *exec.Cmd) {
+				vars := environToMap(cmd.Env)
+				assert.Equal(t, "true", vars["SHOULD_ENABLE_LOCALDNS"])
+				assert.Equal(t, "true", vars["SHOULD_ENABLE_HOSTS_PLUGIN"])
+			},
+		},
+		{
+			name:       "AKSUbuntu2204 with LocalDNS enabled but hosts plugin disabled",
+			folder:     "AKSUbuntu2204+LocalDNS",
+			k8sVersion: "1.24.2",
+			aksNodeConfigUpdator: func(aksNodeConfig *aksnodeconfigv1.Configuration) {
+				aksNodeConfig.LocalDnsProfile = &aksnodeconfigv1.LocalDnsProfile{
+					EnableLocalDns:    true,
+					EnableHostsPlugin: false,
+				}
+			},
+			validator: func(cmd *exec.Cmd) {
+				vars := environToMap(cmd.Env)
+				assert.Equal(t, "true", vars["SHOULD_ENABLE_LOCALDNS"])
+				assert.Equal(t, "false", vars["SHOULD_ENABLE_HOSTS_PLUGIN"])
+			},
+		},
 	}
 
 	for _, tt := range tests {

aks-node-controller/parser/templates/localdns.toml.gtpl

@@ -7,7 +7,7 @@ health-check.localdns.local:53 {
     whoami
 }
 # VnetDNS overrides apply to DNS traffic from pods with dnsPolicy:default or kubelet (referred to as VnetDNS traffic).
-{{- range $domain, $override := $.LocalDnsProfile.VnetDnsOverrides -}}
+{{- range $domain, $override := $.Config.LocalDnsProfile.VnetDnsOverrides -}}
 {{- $isRootDomain := eq $domain "." -}}
 {{- $fwdToClusterCoreDNS := or (hasSuffix $domain "cluster.local") (eq $override.ForwardDestination "ClusterCoreDNS")}}
 {{- $forwardPolicy := "sequential" -}}
@@ -23,11 +23,17 @@ health-check.localdns.local:53 {
     log
     {{- end }}
     bind {{getLocalDnsNodeListenerIp}}
+    {{- if and $isRootDomain $.IncludeHostsPlugin}}
+    # Check /etc/localdns/hosts first for critical AKS FQDNs (mcr.microsoft.com, packages.aks.azure.com, etc.)
+    hosts /etc/localdns/hosts {
+        fallthrough
+    }
+    {{- end}}
     {{- if $isRootDomain}}
     forward . {{getAzureDnsIp}} {
     {{- else}}
     {{- if $fwdToClusterCoreDNS}}
-    forward . {{getCoreDnsServiceIp $}} {
+    forward . {{getCoreDnsServiceIp $.Config}} {
     {{- else}}
     forward . {{getAzureDnsIp}} {
     {{- end}}
@@ -67,7 +73,7 @@ health-check.localdns.local:53 {
 }
 {{- end}}
 # KubeDNS overrides apply to DNS traffic from pods with dnsPolicy:ClusterFirst (referred to as KubeDNS traffic).
-{{- range $domain, $override := $.LocalDnsProfile.KubeDnsOverrides}}
+{{- range $domain, $override := $.Config.LocalDnsProfile.KubeDnsOverrides}}
 {{- $isRootDomain := eq $domain "." -}}
 {{- $fwdToClusterCoreDNS := or (hasSuffix $domain "cluster.local") (eq $override.ForwardDestination "ClusterCoreDNS")}}
 {{- $forwardPolicy := "" }}
@@ -84,8 +90,14 @@ health-check.localdns.local:53 {
     log
     {{- end }}
     bind {{getLocalDnsClusterListenerIp}}
+    {{- if and $isRootDomain $.IncludeHostsPlugin}}
+    # Check /etc/localdns/hosts first for critical AKS FQDNs (mcr.microsoft.com, packages.aks.azure.com, etc.)
+    hosts /etc/localdns/hosts {
+        fallthrough
+    }
+    {{- end}}
     {{- if $fwdToClusterCoreDNS}}
-    forward . {{getCoreDnsServiceIp $}} {
+    forward . {{getCoreDnsServiceIp $.Config}} {
     {{- else}}
     forward . {{getAzureDnsIp}} {
     {{- end}}

aks-node-controller/parser/testdata/AKSUbuntu2204+LocalDNS+HostsPlugin/generatedCSECommand

@@ -0,0 +1 @@
+/bin/bash -c echo $(date),$(hostname) > ${PROVISION_OUTPUT};  CLOUD_INIT_STATUS_SCRIPT="/opt/azure/containers/cloud-init-status-check.sh"; cloudInitExitCode=0; if [ -f "${CLOUD_INIT_STATUS_SCRIPT}" ]; then 	/bin/bash -c "source ${CLOUD_INIT_STATUS_SCRIPT}; handleCloudInitStatus \"${PROVISION_OUTPUT}\"; returnStatus=\$?; echo \"Cloud init status check exit code: \$returnStatus\" >> ${PROVISION_OUTPUT}; exit \$returnStatus" >> ${PROVISION_OUTPUT} 2>&1; else     cloud-init status --wait > /dev/null 2>&1; fi; cloudInitExitCode=$?; if [ "$cloudInitExitCode" -eq 0 ]; then 	echo "cloud-init succeeded" >> ${PROVISION_OUTPUT}; else 	echo "cloud-init failed with exit code ${cloudInitExitCode}" >> ${PROVISION_OUTPUT}; 	cat ${PROVISION_OUTPUT} 	exit ${cloudInitExitCode}; fi;   /usr/bin/nohup /bin/bash -c "/bin/bash /opt/azure/containers/provision_start.sh" 

aks-node-controller/parser/testdata/AKSUbuntu2204+LocalDNS/generatedCSECommand

@@ -0,0 +1 @@
+/bin/bash -c echo $(date),$(hostname) > ${PROVISION_OUTPUT};  CLOUD_INIT_STATUS_SCRIPT="/opt/azure/containers/cloud-init-status-check.sh"; cloudInitExitCode=0; if [ -f "${CLOUD_INIT_STATUS_SCRIPT}" ]; then 	/bin/bash -c "source ${CLOUD_INIT_STATUS_SCRIPT}; handleCloudInitStatus \"${PROVISION_OUTPUT}\"; returnStatus=\$?; echo \"Cloud init status check exit code: \$returnStatus\" >> ${PROVISION_OUTPUT}; exit \$returnStatus" >> ${PROVISION_OUTPUT} 2>&1; else     cloud-init status --wait > /dev/null 2>&1; fi; cloudInitExitCode=$?; if [ "$cloudInitExitCode" -eq 0 ]; then 	echo "cloud-init succeeded" >> ${PROVISION_OUTPUT}; else 	echo "cloud-init failed with exit code ${cloudInitExitCode}" >> ${PROVISION_OUTPUT}; 	cat ${PROVISION_OUTPUT} 	exit ${cloudInitExitCode}; fi;   /usr/bin/nohup /bin/bash -c "/bin/bash /opt/azure/containers/provision_start.sh" 

aks-node-controller/proto/aksnodeconfig/v1/localdns_config.proto

@@ -19,6 +19,11 @@ message LocalDnsProfile {
 
   // KubeDns overrides apply to DNS traffic from pods with dnsPolicy:ClusterFirst (referred to as KubeDns traffic).
   map<string, LocalDnsOverrides> kube_dns_overrides = 5;
+
+  // Specifies whether the hosts plugin should be enabled in the localdns Corefile.
+  // When true and LocalDNS is enabled, the Corefile will include a hosts plugin block
+  // that serves cached DNS entries from /etc/localdns/hosts for critical AKS FQDNs.
+  bool enable_hosts_plugin = 6;
 }
 
 // Represents DNS override settings for both VnetDNS and KubeDNS traffic.