fix: regenerate localdns corefile if it goes missing

e2e/node_config.go

@@ -238,7 +238,7 @@ func nbcToAKSNodeConfigV1(nbc *datamodel.NodeBootstrappingConfiguration) *aksnod
 				"testdomain456.com": {
 					QueryLogging:                "Log",
 					Protocol:                    "PreferUDP",
-					ForwardDestination:          "ClusterCoreDNS",
+					ForwardDestination:          "VnetDNS",
 					ForwardPolicy:               "Random",
 					MaxConcurrent:               to.Ptr(int32(1000)),
 					CacheDurationInSeconds:      to.Ptr(int32(3600)),

e2e/scenario_test.go

@@ -1715,6 +1715,47 @@ func Test_AzureLinuxV2_LocalDns_Disabled_Scriptless(t *testing.T) {
 	})
 }
 
+// func Test_Ubuntu2204_LocalDns_Corefile_Regeneration_Scriptless(t *testing.T) {
+// 	RunScenario(t, &Scenario{
+// 		Description: "Tests that localdns corefile can be regenerated after deletion using systemd environment variable",
+// 		Tags: Tags{
+// 			Scriptless: true,
+// 		},
+// 		Config: Config{
+// 			Cluster: ClusterAzureNetwork,
+// 			VHD:     config.VHDUbuntu2204Gen2Containerd,
+// 			BootstrapConfigMutator: func(nbc *datamodel.NodeBootstrappingConfiguration) {
+// 				nbc.ContainerService.Properties.AgentPoolProfiles[0].Distro = "aks-ubuntu-containerd-22.04-gen2"
+// 				nbc.AgentPoolProfile.Distro = "aks-ubuntu-containerd-22.04-gen2"
+// 			},
+// 			Validator: func(ctx context.Context, s *Scenario) {
+// 				validateLocalDNSCorefileRegeneration(ctx, s)
+// 			},
+// 		},
+// 	})
+// }
+
+// func Test_AzureLinuxV2_LocalDns_Corefile_Regeneration_Scriptless(t *testing.T) {
+// 	RunScenario(t, &Scenario{
+// 		Description: "Tests that localdns corefile can be regenerated after deletion on Azure Linux V2",
+// 		Tags: Tags{
+// 			Scriptless: true,
+// 		},
+// 		Config: Config{
+// 			Cluster: ClusterAzureNetwork,
+// 			VHD:     config.VHDAzureLinuxV2Gen2,
+// 			AKSNodeConfigMutator: func(config *aksnodeconfigv1.Configuration) {
+// 				config.LocalDnsProfile = &aksnodeconfigv1.LocalDnsProfile{
+// 					EnableLocalDns: true,
+// 				}
+// 			},
+// 			Validator: func(ctx context.Context, s *Scenario) {
+// 				validateLocalDNSCorefileRegeneration(ctx, s)
+// 			},
+// 		},
+// 	})
+// }
+
 func Test_Ubuntu2204_KubeletCustomConfig(t *testing.T) {
 	RunScenario(t, &Scenario{
 		Tags: Tags{

e2e/validation.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/Azure/agentbaker/e2e/config"
 	"github.com/Azure/agentbaker/e2e/toolkit"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/stretchr/testify/assert"
@@ -78,7 +77,7 @@ func ValidateCommonLinux(ctx context.Context, s *Scenario) {
 	}
 
 	// localdns is not supported on scriptless, privatekube and VHDUbuntu2204Gen2ContainerdAirgappedK8sNotCached.
-	if s.Tags.Scriptless != true && s.VHD != config.VHDUbuntu2204Gen2ContainerdPrivateKubePkg && s.VHD != config.VHDUbuntu2204Gen2ContainerdAirgappedK8sNotCached && !s.VHD.UnsupportedLocalDns {
+	if !s.VHD.UnsupportedLocalDns {
 		ValidateLocalDNSService(ctx, s, "enabled")
 		ValidateLocalDNSResolution(ctx, s, "169.254.10.10")
 	}

e2e/validators.go

@@ -1178,6 +1178,58 @@ func ValidateLocalDNSResolution(ctx context.Context, s *Scenario, server string)
 	assert.Contains(s.T, execResult.stdout.String(), fmt.Sprintf("SERVER: %s", server))
 }
 
+// func validateLocalDNSCorefileRegeneration(ctx context.Context, s *Scenario) {
+// 	ValidateLocalDNSService(ctx, s, "enabled")
+
+// 	// Validate systemd drop-in file exists with correct environment variable
+// 	dropinFile := "/etc/systemd/system/localdns.service.d/10-localdns-environment.conf"
+// 	ValidateFileExists(ctx, s, dropinFile)
+// 	ValidateFileHasContent(ctx, s, dropinFile, "[Service]")
+// 	ValidateFileHasContent(ctx, s, dropinFile, "Environment=\"LOCALDNS_BASE64_ENCODED_COREFILE=")
+
+// 	// Validate original corefile exists
+// 	corefilePath := "/opt/azure/containers/localdns/localdns.corefile"
+// 	ValidateFileExists(ctx, s, corefilePath)
+
+// 	// Test corefile regeneration: delete corefile and restart service
+// 	s.T.Log("Testing corefile regeneration after deletion...")
+// 	script := `set -euo pipefail
+// # Delete the corefile
+// sudo rm -f /opt/azure/containers/localdns/localdns.corefile
+
+// # Verify file is deleted
+// if [ -f /opt/azure/containers/localdns/localdns.corefile ]; then
+//     echo "ERROR: Corefile still exists after deletion"
+//     exit 1
+// fi
+
+// # Restart localdns service
+// sudo systemctl restart localdns
+
+// # Wait for service to be active
+// sleep 2
+// systemctl is-active localdns
+
+// # Verify corefile was regenerated
+// if [ ! -f /opt/azure/containers/localdns/localdns.corefile ]; then
+//     echo "ERROR: Corefile was not regenerated"
+//     exit 1
+// fi
+
+// echo "SUCCESS: Corefile regenerated successfully"
+// exit 0
+// `
+// 	execScriptOnVMForScenarioValidateExitCode(ctx, s, script, 0, "corefile regeneration test failed")
+
+// 	// Validate localdns service is still running after regeneration
+// 	ValidateLocalDNSService(ctx, s, "enabled")
+
+// 	// Validate DNS resolution still works
+// 	ValidateLocalDNSResolution(ctx, s, "169.254.10.10")
+
+// 	s.T.Log("All localdns corefile regeneration tests passed successfully")
+// }
+
 // ValidateJournalctlOutput checks if specific content exists in the systemd service logs
 func ValidateJournalctlOutput(ctx context.Context, s *Scenario, serviceName string, expectedContent string) {
 	s.T.Helper()

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

@@ -1152,6 +1152,15 @@ shouldEnableLocalDns() {
     chmod 0644 "${LOCALDNS_COREFILE}"
     echo "${LOCALDNS_GENERATED_COREFILE}" | base64 -d > "${LOCALDNS_COREFILE}" || exit $ERR_LOCALDNS_FAIL
 
+    # Create environment file for corefile regeneration.
+    # This file will be referenced by localdns.service using EnvironmentFile directive.
+    LOCALDNS_ENV_FILE="/etc/localdns/environment"
+    mkdir -p "$(dirname "${LOCALDNS_ENV_FILE}")"
+    cat > "${LOCALDNS_ENV_FILE}" <<EOF
+LOCALDNS_BASE64_ENCODED_COREFILE=${LOCALDNS_GENERATED_COREFILE}
+EOF
+    chmod 0644 "${LOCALDNS_ENV_FILE}"
+
 	mkdir -p "$(dirname "${LOCALDNS_SLICEFILE}")"
     touch "${LOCALDNS_SLICEFILE}"
     chmod 0644 "${LOCALDNS_SLICEFILE}"

parts/linux/cloud-init/artifacts/localdns.service

@@ -16,7 +16,8 @@ Restart=on-failure
 KillMode=mixed
 TimeoutStopSec=30
 Slice=localdns.slice
+EnvironmentFile=-/etc/localdns/environment
 ExecStart=/opt/azure/containers/localdns/localdns.sh
 
 [Install]
-WantedBy=multi-user.target
+WantedBy=multi-user.target

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

@@ -74,9 +74,18 @@ verify_localdns_corefile() {
         return 1
     fi
 
+    # Check if corefile exists, is not empty, and has valid coredns server block.
     if [ ! -f "${LOCALDNS_CORE_FILE}" ] || [ ! -s "${LOCALDNS_CORE_FILE}" ]; then
         echo "Localdns corefile either does not exist or is empty at ${LOCALDNS_CORE_FILE}."
-        return 1
+
+        echo "Attempting to regenerate localdns corefile..."
+        if regenerate_localdns_corefile; then
+            echo "Localdns corefile regenerated successfully."
+            return 0
+        else
+            echo "Failed to regenerate localdns corefile."
+            return 1
+        fi
     fi
     return 0
 }
@@ -114,6 +123,31 @@ verify_localdns_binary() {
     return 0
 }
 
+# Regenerate the localdns corefile from base64 encoded content.
+# This is used when the corefile goes missing.
+regenerate_localdns_corefile() {
+    if [ -z "${LOCALDNS_BASE64_ENCODED_COREFILE:-}" ]; then
+        echo "LOCALDNS_BASE64_ENCODED_COREFILE is not set. Cannot regenerate corefile."
+        return 1
+    fi
+    echo "Regenerating localdns corefile at ${LOCALDNS_CORE_FILE}"
+
+    mkdir -p "$(dirname "${LOCALDNS_CORE_FILE}")"
+    # Decode base64 corefile content and write to corefile.
+    if ! echo "${LOCALDNS_BASE64_ENCODED_COREFILE}" | base64 -d > "${LOCALDNS_CORE_FILE}"; then
+        echo "Failed to decode and write corefile."
+        return 1
+    fi
+
+    chmod 0644 "${LOCALDNS_CORE_FILE}" || {
+        echo "Failed to set permissions on ${LOCALDNS_CORE_FILE}"
+        return 1
+    }
+
+    echo "Successfully regenerated localdns corefile."
+    return 0
+}
+
 # Replace AzureDNSIP in corefile with VNET DNS ServerIPs if necessary.
 replace_azurednsip_in_corefile() {
     if [ -z "${RESOLV_CONF:-}" ]; then