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/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"
@@ -99,7 +98,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")
 	}

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

@@ -1145,12 +1145,21 @@ LOCALDNS_SLICEFILE="/etc/systemd/system/localdns.slice"
 # It creates the localdns corefile and slicefile, then enables and starts localdns.
 # In this function, generated base64 encoded localdns corefile is decoded and written to the corefile path.
 # This function also creates the localdns slice file with memory and cpu limits, that will be used by localdns systemd unit.
-shouldEnableLocalDns() {
+enableLocalDNSForScriptless() {
     mkdir -p "$(dirname "${LOCALDNS_COREFILE}")"
     touch "${LOCALDNS_COREFILE}"
     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/cse_main.sh

@@ -295,7 +295,7 @@ EOF
 
     # This is to enable localdns using scriptless.
     if [ "${SHOULD_ENABLE_LOCALDNS}" = "true" ]; then
-        logs_to_events "AKS.CSE.shouldEnableLocalDns" shouldEnableLocalDns || exit $ERR_LOCALDNS_FAIL
+        logs_to_events "AKS.CSE.enableLocalDNSForScriptless" enableLocalDNSForScriptless || exit $ERR_LOCALDNS_FAIL
     fi
 
     if [ "${ID}" != "mariner" ] && [ "${ID}" != "azurelinux" ]; then

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

@@ -75,9 +75,18 @@ verify_localdns_corefile() {
         return 1
     fi
 
+    # Check if corefile exists, is not empty, else attempt to regenerate it.
     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
 }
@@ -115,6 +124,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

spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh

@@ -548,15 +548,15 @@ Describe 'cse_config.sh'
 
         # Success case.
         It 'should enable localdns successfully'
-            When call shouldEnableLocalDns
+            When call enableLocalDNSForScriptless
             The status should be success
             The output should include "localdns should be enabled."
             The output should include "Enable localdns succeeded."
         End
 
         # Corefile file creation.
         It 'should create localdns.corefile with correct data'
-            When call shouldEnableLocalDns
+            When call enableLocalDNSForScriptless
             The status should be success
             The output should include "localdns should be enabled."
             The path "$LOCALDNS_COREFILE" should be file
@@ -568,7 +568,7 @@ Describe 'cse_config.sh'
         # Corefile already exists (idempotency).
         It 'should overwrite existing localdns.corefile'
             echo "wrong data" > "$LOCALDNS_COREFILE"
-            When call shouldEnableLocalDns
+            When call enableLocalDNSForScriptless
             The status should be success
             The path "$LOCALDNS_COREFILE" should be file
             The contents of file "$LOCALDNS_COREFILE" should include "localdns corefile"
@@ -578,7 +578,7 @@ Describe 'cse_config.sh'
 
         # Slice file creation.
         It 'should create localdns.slice with correct CPU and Memory limits'
-            When call shouldEnableLocalDns
+            When call enableLocalDNSForScriptless
             The status should be success
             The output should include "localdns should be enabled."
             The path "$LOCALDNS_SLICEFILE" should be file

spec/parts/linux/cloud-init/artifacts/localdns_spec.sh

@@ -10,6 +10,13 @@ Describe 'localdns.sh'
 # These functions are defined in parts/linux/cloud-init/artifacts/localdns.sh file.
 #------------------------------------------------------------------------------------------------------------------------------------
     Describe 'verify_localdns_files'
+        check_file_permissions() {
+            local file=$1
+            local expected_perms=$2
+            local actual_perms=$(stat -c '%a' "$file" 2>/dev/null)
+            [ "$actual_perms" = "$expected_perms" ]
+        }
+
         setup() {
             Include "./parts/linux/cloud-init/artifacts/localdns.sh"
 
@@ -18,8 +25,12 @@ Describe 'localdns.sh'
             LOCALDNS_CORE_FILE="${LOCALDNS_SCRIPT_PATH}/localdns.corefile"
             UPDATED_LOCALDNS_CORE_FILE="${LOCALDNS_SCRIPT_PATH}/updated.localdns.corefile"
             mkdir -p "$LOCALDNS_SCRIPT_PATH"
-            echo "forward . 168.63.129.16" >> "$LOCALDNS_CORE_FILE"
-            echo "forward . 168.63.129.16" >> "$UPDATED_LOCALDNS_CORE_FILE"
+            echo ".:5353 {" >> "$LOCALDNS_CORE_FILE"
+            echo "    forward . 168.63.129.16" >> "$LOCALDNS_CORE_FILE"
+            echo "}" >> "$LOCALDNS_CORE_FILE"
+            echo ".:5353 {" >> "$UPDATED_LOCALDNS_CORE_FILE"
+            echo "    forward . 168.63.129.16" >> "$UPDATED_LOCALDNS_CORE_FILE"
+            echo "}" >> "$UPDATED_LOCALDNS_CORE_FILE"
 
             LOCALDNS_SLICE_PATH="${TEST_DIR}/etc/systemd/system"
             LOCALDNS_SLICE_FILE="${LOCALDNS_SLICE_PATH}/localdns.slice"
@@ -54,24 +65,78 @@ EOF
         }
         BeforeEach 'setup'
         AfterEach 'cleanup'
+        #------------------------ regenerate_localdns_corefile ---------------------------------------------
+        It 'should regenerate corefile successfully when LOCALDNS_BASE64_ENCODED_COREFILE is set'
+            rm -f "$LOCALDNS_CORE_FILE"
+            LOCALDNS_BASE64_ENCODED_COREFILE=$(echo ".:5353 {
+    forward . 168.63.129.16
+}" | base64)
+            When run regenerate_localdns_corefile
+            The status should be success
+            The stdout should include "Regenerating localdns corefile at $LOCALDNS_CORE_FILE"
+            The stdout should include "Successfully regenerated localdns corefile."
+            The path "$LOCALDNS_CORE_FILE" should be file
+        End
+
+        It 'should fail to regenerate when LOCALDNS_BASE64_ENCODED_COREFILE is not set'
+            rm -f "$LOCALDNS_CORE_FILE"
+            unset LOCALDNS_BASE64_ENCODED_COREFILE
+            When run regenerate_localdns_corefile
+            The status should be failure
+            The stdout should include "LOCALDNS_BASE64_ENCODED_COREFILE is not set. Cannot regenerate corefile."
+        End
+
+        It 'should set correct permissions on regenerated corefile'
+            rm -f "$LOCALDNS_CORE_FILE"
+            LOCALDNS_BASE64_ENCODED_COREFILE=$(echo ".:5353 {
+    forward . 168.63.129.16
+}" | base64)
+            When run regenerate_localdns_corefile
+            The status should be success
+            The stdout should include "Successfully regenerated localdns corefile."
+            The path "$LOCALDNS_CORE_FILE" should be file
+        End
+
         #------------------------ verify_localdns_corefile -------------------------------------------------
         It 'should return success if localdns corefile exists and is not empty'
             When run verify_localdns_corefile
             The status should be success
         End
 
-        It 'should return failure if localdns corefile does not exist'
+        It 'should succeed without regeneration if corefile exists and is not empty'
+            echo ".:5353 {
+    forward . 168.63.129.16
+}" > "$LOCALDNS_CORE_FILE"
+            When run verify_localdns_corefile
+            The status should be success
+        End
+
+        It 'should regenerate and succeed if corefile is missing and LOCALDNS_BASE64_ENCODED_COREFILE is set'
+            rm -f "$LOCALDNS_CORE_FILE"
+            LOCALDNS_BASE64_ENCODED_COREFILE=$(echo ".:5353 {
+    forward . 168.63.129.16
+}" | base64)
+            When run verify_localdns_corefile
+            The status should be success
+            The stdout should include "Attempting to regenerate localdns corefile..."
+            The stdout should include "Localdns corefile regenerated successfully."
+        End
+
+        It 'should return failure if localdns corefile does not exist and regeneration fails'
             rm -r "$LOCALDNS_CORE_FILE"
             When run verify_localdns_corefile
             The status should be failure
             The stdout should include "Localdns corefile either does not exist or is empty at $LOCALDNS_CORE_FILE."
+            The stdout should include "Attempting to regenerate localdns corefile..."
+            The stdout should include "LOCALDNS_BASE64_ENCODED_COREFILE is not set. Cannot regenerate corefile."
         End
 
-        It 'should return failure if localdns corefile is empty'
+        It 'should return failure if localdns corefile is empty and regeneration fails'
             > "$LOCALDNS_CORE_FILE"
             When run verify_localdns_corefile
             The status should be failure
             The stdout should include "Localdns corefile either does not exist or is empty at $LOCALDNS_CORE_FILE."
+            The stdout should include "Attempting to regenerate localdns corefile..."
         End
 
         It 'should return failure if LOCALDNS_CORE_FILE is unset'