feat: add OSSKU-based IsCgroupV2 detection and CustomizedImageTrustedLaunch distro (#8252)

Signed-off-by: Aadhar Agarwal <aadagarwal@microsoft.com>

Author: aadhar-agarwal

pkg/agent/baker.go

@@ -639,7 +639,8 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration
 			return config.GetOrderedKubeproxyConfigStringForPowershell()
 		},
 		"IsCgroupV2": func() bool {
-			return profile.Is2204VHDDistro() || profile.IsAzureLinuxCgroupV2VHDDistro() || profile.Is2404VHDDistro() || profile.IsFlatcar() || profile.IsACL()
+			return profile.Is2204VHDDistro() || profile.Is2404VHDDistro() ||
+				config.IsAzureLinux() || config.IsFlatcar() || config.IsACL()
 		},
 		"GetKubeProxyFeatureGatesPsh": func() string {
 			return cs.Properties.GetKubeProxyFeatureGatesWindowsArguments()
@@ -735,7 +736,8 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration
 		"IsCustomImage": func() bool {
 			return profile.Distro == datamodel.CustomizedImage ||
 				profile.Distro == datamodel.CustomizedImageKata ||
-				profile.Distro == datamodel.CustomizedImageLinuxGuard
+				profile.Distro == datamodel.CustomizedImageLinuxGuard ||
+				profile.Distro == datamodel.CustomizedImageTrustedLaunch
 		},
 		"EnableHostsConfigAgent": func() bool {
 			return cs.Properties.OrchestratorProfile.KubernetesConfig != nil &&

pkg/agent/baker_test.go

@@ -1154,6 +1154,57 @@ var _ = Describe("getLinuxNodeCSECommand", func() {
 		Expect(vars).To(HaveKeyWithValue("NEEDS_CGROUPV2", "true"))
 	})
 
+	It("should set NEEDS_CGROUPV2 for CustomizedImage with AzureLinux OSSKU", func() {
+		config, err := deepcopy.Anything(baseConfig)
+		Expect(err).To(BeNil())
+		typedConfig, ok := config.(*datamodel.NodeBootstrappingConfiguration)
+		Expect(ok).To(BeTrue())
+		typedConfig.AgentPoolProfile.Distro = datamodel.CustomizedImage
+		typedConfig.OSSKU = datamodel.OSSKUAzureLinux
+
+		cseCmd := templateGenerator.getLinuxNodeCSECommand(typedConfig)
+
+		Expect(cseCmd).NotTo(BeEmpty())
+		Expect(strings.Contains(cseCmd, "\n")).To(BeFalse())
+
+		vars := decodeCSEVars(cseCmd)
+		Expect(vars).To(HaveKeyWithValue("NEEDS_CGROUPV2", "true"))
+	})
+
+	It("should set NEEDS_CGROUPV2 for CustomizedImage with Flatcar OSSKU", func() {
+		config, err := deepcopy.Anything(baseConfig)
+		Expect(err).To(BeNil())
+		typedConfig, ok := config.(*datamodel.NodeBootstrappingConfiguration)
+		Expect(ok).To(BeTrue())
+		typedConfig.AgentPoolProfile.Distro = datamodel.CustomizedImage
+		typedConfig.OSSKU = datamodel.OSSKUFlatcar
+
+		cseCmd := templateGenerator.getLinuxNodeCSECommand(typedConfig)
+
+		Expect(cseCmd).NotTo(BeEmpty())
+		Expect(strings.Contains(cseCmd, "\n")).To(BeFalse())
+
+		vars := decodeCSEVars(cseCmd)
+		Expect(vars).To(HaveKeyWithValue("NEEDS_CGROUPV2", "true"))
+	})
+
+	It("should set NEEDS_CGROUPV2 for CustomizedImageTrustedLaunch with AzureContainerLinux OSSKU", func() {
+		config, err := deepcopy.Anything(baseConfig)
+		Expect(err).To(BeNil())
+		typedConfig, ok := config.(*datamodel.NodeBootstrappingConfiguration)
+		Expect(ok).To(BeTrue())
+		typedConfig.AgentPoolProfile.Distro = datamodel.CustomizedImageTrustedLaunch
+		typedConfig.OSSKU = datamodel.OSSKUAzureContainerLinux
+
+		cseCmd := templateGenerator.getLinuxNodeCSECommand(typedConfig)
+
+		Expect(cseCmd).NotTo(BeEmpty())
+		Expect(strings.Contains(cseCmd, "\n")).To(BeFalse())
+
+		vars := decodeCSEVars(cseCmd)
+		Expect(vars).To(HaveKeyWithValue("NEEDS_CGROUPV2", "true"))
+	})
+
 	It("should panic when template processing fails", func() {
 		// Create invalid config that will cause template processing to fail
 		invalidConfig := &datamodel.NodeBootstrappingConfiguration{

pkg/agent/bakerapi.go

@@ -53,7 +53,8 @@ func (agentBaker *agentBakerImpl) GetNodeBootstrapping(ctx context.Context, conf
 	if distro == datamodel.CustomizedWindowsOSImage ||
 		distro == datamodel.CustomizedImage ||
 		distro == datamodel.CustomizedImageKata ||
-		distro == datamodel.CustomizedImageLinuxGuard {
+		distro == datamodel.CustomizedImageLinuxGuard ||
+		distro == datamodel.CustomizedImageTrustedLaunch {
 		return nodeBootstrapping, nil
 	}
 

pkg/agent/bakerapi_test.go

@@ -254,6 +254,15 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 			Expect(err).NotTo(HaveOccurred())
 		})
 
+		It("should not return an error for customized trusted launch image", func() {
+			config.AgentPoolProfile.Distro = datamodel.CustomizedImageTrustedLaunch
+			agentBaker, err := NewAgentBaker()
+			Expect(err).NotTo(HaveOccurred())
+
+			_, err = agentBaker.GetNodeBootstrapping(context.Background(), config)
+			Expect(err).NotTo(HaveOccurred())
+		})
+
 		It("should not return an error for customized windows image", func() {
 			config.AgentPoolProfile.Distro = datamodel.CustomizedWindowsOSImage
 			agentBaker, err := NewAgentBaker()

pkg/agent/datamodel/types.go

@@ -213,11 +213,12 @@ const (
 	// AKSWindows2025Gen2 stands for distro for windows server 2025 Gen 2 SIG image.
 	AKSWindows2025Gen2 Distro = "aks-windows-2025-gen2"
 	// AKSWindows2019PIR stands for distro of windows server 2019 PIR image with docker.
-	AKSWindows2019PIR         Distro = "aks-windows-2019-pir"
-	CustomizedImage           Distro = "CustomizedImage"
-	CustomizedImageKata       Distro = "CustomizedImageKata"
-	CustomizedImageLinuxGuard Distro = "CustomizedImageLinuxGuard"
-	CustomizedWindowsOSImage  Distro = "CustomizedWindowsOSImage"
+	AKSWindows2019PIR            Distro = "aks-windows-2019-pir"
+	CustomizedImage              Distro = "CustomizedImage"
+	CustomizedImageKata          Distro = "CustomizedImageKata"
+	CustomizedImageLinuxGuard    Distro = "CustomizedImageLinuxGuard"
+	CustomizedImageTrustedLaunch Distro = "CustomizedImageTrustedLaunch"
+	CustomizedWindowsOSImage     Distro = "CustomizedWindowsOSImage"
 
 	// USNatCloud is a const string reference identifier for USNat.
 	USNatCloud = "USNatCloud"
@@ -1798,6 +1799,10 @@ type NodeBootstrappingConfiguration struct {
 	EnableScriptlessCSECmd bool
 }
 
+func (config *NodeBootstrappingConfiguration) IsAzureLinux() bool {
+	return config.OSSKU == OSSKUAzureLinux || config.AgentPoolProfile.IsAzureLinuxCgroupV2VHDDistro()
+}
+
 func (config *NodeBootstrappingConfiguration) IsFlatcar() bool {
 	return config.OSSKU == OSSKUFlatcar || config.AgentPoolProfile.IsFlatcar()
 }