fix: isolate additional Linux SSH lifecycle safeguards

aks-node-controller/parser/parser.go

@@ -47,6 +47,8 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 	containerdVersion, _ := detectContainerdVersion(ctx)
 	cloudProviderSettings := getCloudProviderSettings(config)
 	isMIGNode := getIsMIGNode(config.GetGpuConfig().GetGpuInstanceProfile(), config.GetGpuConfig().GetMigProfileLayout())
+	// Explicit SSH disablement wins; otherwise allow explicit SSH or Entra-only access.
+	shouldRestoreSSH := !getDisableSSH(config) && (config.GetEnableSsh() || config.GetDisablePubkeyAuth())
 	env := map[string]string{
 		"PROVISION_OUTPUT":                                     "/var/log/azure/cluster-provision-cse-output.log",
 		"CUSTOM_DATA_ONLY_PROVISIONING":                        fmt.Sprintf("%t", config.GetCustomDataOnlyProvisioning()),
@@ -112,6 +114,7 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 		"RUNC_PACKAGE_URL":                                     config.GetRuncConfig().GetRuncPackageUrl(),
 		"ENABLE_HOSTS_CONFIG_AGENT":                            fmt.Sprintf("%v", config.GetEnableHostsConfigAgent()),
 		"DISABLE_SSH":                                          fmt.Sprintf("%v", getDisableSSH(config)),
+		"SHOULD_RESTORE_SSH":                                   fmt.Sprintf("%v", shouldRestoreSSH),
 		"SHOULD_CONFIGURE_HTTP_PROXY":                          fmt.Sprintf("%v", getShouldConfigureHTTPProxy(config.GetHttpProxyConfig())),
 		"SHOULD_CONFIGURE_HTTP_PROXY_CA":                       fmt.Sprintf("%v", getShouldConfigureHTTPProxyCA(config.GetHttpProxyConfig())),
 		"HTTP_PROXY_TRUSTED_CA":                                removeNewlines(config.GetHttpProxyConfig().GetProxyTrustedCa()),

aks-node-controller/parser/parser_test.go

@@ -21,6 +21,32 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+func TestBuildCSECmdSSHPolicy(t *testing.T) {
+	for _, test := range []struct {
+		name      string
+		enable    *bool
+		pubkeyOff bool
+		disable   string
+		restore   string
+	}{
+		{name: "unspecified", disable: "false", restore: "false"},
+		{name: "disabled", enable: proto.Bool(false), disable: "true", restore: "false"},
+		{name: "explicitly enabled", enable: proto.Bool(true), disable: "false", restore: "true"},
+		{name: "explicit Entra authentication", pubkeyOff: true, disable: "false", restore: "true"},
+		{name: "disabled takes precedence", enable: proto.Bool(false), pubkeyOff: true, disable: "true", restore: "false"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			cmd, err := BuildCSECmd(context.Background(), &aksnodeconfigv1.Configuration{
+				EnableSsh: test.enable, DisablePubkeyAuth: proto.Bool(test.pubkeyOff),
+			}, nil)
+			require.NoError(t, err)
+			env := environToMap(cmd.Env)
+			require.Equal(t, test.disable, env["DISABLE_SSH"])
+			require.Equal(t, test.restore, env["SHOULD_RESTORE_SSH"])
+		})
+	}
+}
+
 func TestBuildCSECmdContainerdImagePullSettings(t *testing.T) {
 	for _, version := range []string{"1.7.22", "2.2.4"} {
 		t.Run(version, func(t *testing.T) {

e2e/README.md

@@ -30,12 +30,14 @@ and require an updated CSE bundle or VHD to receive this fix.
 
 Linux SSH-disabled scenarios use RunCommand to verify the persistent SSH block,
 restart the same VM through ARM, confirm its boot ID changed, and check the block
-again after kubelet becomes active. CSE stops SSH listeners and uses an AKS-owned
-systemd condition without changing distribution enablement or administrator masks.
+again after kubelet becomes active. CSE records activation, stops and disables SSH
+listeners, and uses an AKS-owned systemd condition without removing administrator masks.
 SSH policy runs only in `nodePrep`, so image preparation leaves SSH defaults
 unchanged. This delays SSH policy enforcement until `nodePrep` on ordinary nodes
-as well. There is no activation history or automatic re-enable path; a code
-rollback alone does not remove a persisted block.
+as well. The follow-up retains activation and public-key authentication snapshots
+for explicit restoration. These are separate safeguards, not a requirement of the
+normal node-pool SSH configuration change, which takes effect through reimage.
+A code rollback alone does not remove a persisted block.
 
 ## Writing and extending scenarios
 

parts/linux/cloud-init/artifacts/README.MD

@@ -7,18 +7,18 @@ For CA refresh logging and Linux opt-in certificate diagnostics, see
 
 When `DISABLE_SSH=true` and `nodePrep` runs, `disableSSH` emits structured
 guest-agent events. A persistent systemd condition blocks future starts; existing
-listeners are stopped without changing unit enablement. Stop retry parameters and
-missing-unit handling are unchanged. Policy-installation and stop failures surface
-as CSE exit code 172 after every SSH unit has been attempted. A successful stop
-does not have to clear a unit's existing `failed` state.
+listeners are stopped and disabled after their activation policy is recorded.
+Stop retry parameters and missing-unit handling are unchanged. Policy-installation,
+inspection, stop, disable, or inactive-state verification failures surface as CSE
+exit code 172 after every SSH unit has been attempted.
 Telemetry failures remain warnings, not new
 provisioning failures. Other `systemctlDisableAndStop` callers and
 `logs_to_events` keep their existing telemetry behavior.
 
 | TaskName | Scope |
 | --- | --- |
 | `AKS.CSE.disableSSH` | One `Started` and one `Completed` event per SSH-disable attempt |
-| `AKS.CSE.disableSSH.unit` | `Started` and `Completed` for each unit reached: `ssh.socket`, `sshd.socket`, `ssh`, `sshd` |
+| `AKS.CSE.disableSSH.unit` | `Started` and `Completed` for each unit handled by `systemctlDisableAndStop`: `ssh.socket`, `sshd.socket`, `ssh`, `sshd` |
 
 The event envelope's `Message` is a JSON string, available as `Context1` in
 `GuestAgentGenericLogs`. Parse it with `parse_json(Context1)`.
@@ -31,13 +31,14 @@ The event envelope's `Message` is a JSON string, available as `Context1` in
 | `ExitCode` | Null for starts; 0/1 for completed units; 0/172 for the completed attempt |
 | `CatExitCode`, `Skipped` | Whether the unit was readable; nonzero `systemctl cat` preserves the existing skip behavior |
 | `StopExitCode` | Stop retry-helper result; null when skipped, not the raw systemctl exit code |
-| `DisableExitCode` | Always null for this stop-only SSH path; retained for schema compatibility |
+| `DisableExitCode` | Disable retry-helper result; null when skipped |
 
 The last three rows apply to completed unit events. The envelope's `Timestamp` and
 `OperationId` contain start and emission times. The schema is kept compact to fit
 the custom-data size limit. A skipped unit can be absent or unreadable;
-it is not evidence of a successful stop. Readable masked units are stopped and
-emit per-unit events without removing their masks. Command exit 0 alone does not prove port
+it is not evidence of a successful stop. Externally masked units are stopped without
+removing their masks and emit only the attempt-level events. Other readable units
+emit per-unit events. Command exit 0 alone does not prove port
 22 is closed: inspect the unit states and use the reboot E2E validation for
 listener checks.
 

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

@@ -94,6 +94,7 @@ RUNC_PACKAGE_URL={{GetParameter "runcPackageURL"}}
 ENABLE_HOSTS_CONFIG_AGENT="{{EnableHostsConfigAgent}}"
 DISABLE_SSH="{{ShouldDisableSSH}}"
 DISABLE_PUBKEY_AUTH="{{ShouldTurnOffPubkeyAuthSSH}}"
+SHOULD_RESTORE_SSH="{{ShouldRestoreSSH}}"
 SHOULD_CONFIGURE_HTTP_PROXY="{{ShouldConfigureHTTPProxy}}"
 SHOULD_CONFIGURE_HTTP_PROXY_CA="{{ShouldConfigureHTTPProxyCA}}"
 HTTP_PROXY_TRUSTED_CA="{{GetHTTPProxyCA}}"