Add baked-In test scenario to ensure no proxy if MSP at disabled state

e2etest/GuestProxyAgentTest/LinuxScripts/IMDSPingTest.sh

@@ -16,18 +16,24 @@ for i in {1..10}; do
     fi
     sleep 1
 
+    authorizationHeader=$(curl -s -I -H "Metadata:True" $url | grep -Fi "x-ms-azure-host-authorization")
     if [ "${imdsSecureChannelEnabled,,}" = "true" ]  # case insensitive comparison
     then
-        authorizationHeader=$(curl -s -I -H "Metadata:True" $url | grep -Fi "x-ms-azure-host-authorization")
         if [ "$authorizationHeader" = "" ]; then
             echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - Response authorization header not exist"
             exit -1
         else
-            echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - Response authorization header exists"
+            echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - Response authorization header exists as expected"
         fi
         sleep 1
     else
-		echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - IMDS secure channel is not enabled. Skipping x-ms-azure-host-authorization header validation"
+        if [ "$authorizationHeader" = "" ]; then
+            echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - Response authorization header not exist as expected"
+        else
+            echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ") - Response authorization header exists"
+            exit -1
+        fi
+        sleep 1
 	fi
 done
 

e2etest/GuestProxyAgentTest/Models/TestMap.cs

@@ -19,6 +19,7 @@ public class TestGroupDetails
         public string VmImageOffer { get; set; } = null!;
         public string VmImageSku { get; set; } = null!;
         public string VmImageVersion { get; set; } = null!;
+        public string SharedGalleryImageUniqueId { get; set; } = null!;
         public List<TestScenarioConfig> Scenarios { get; set; } = null!;
     }
 

e2etest/GuestProxyAgentTest/Scripts/IMDSPingTest.ps1

@@ -23,19 +23,24 @@ while ($i -lt 10) {
             exit -1
         }
 
+        $responseHeaders = $response.Headers
         if ("$imdsSecureChannelEnabled" -ieq "true") { # case insensitive comparison
-            $responseHeaders = $response.Headers
             if ($null -eq $responseHeaders["x-ms-azure-host-authorization"]) {
                 Write-Error "$((Get-Date).ToUniversalTime()) - Ping test failed. Response does not contain x-ms-azure-host-authorization header"
                 exit -1
             }
             else {
                 Write-Output "$((Get-Date).ToUniversalTime()) - Ping test passed. Response contains x-ms-azure-host-authorization header"
             }
-
         }
         else {
-            Write-Output "$((Get-Date).ToUniversalTime()) - IMDS secure channel is not enabled. Skipping x-ms-azure-host-authorization header validation"
+            if ($null -eq $responseHeaders["x-ms-azure-host-authorization"]) {
+                Write-Output "$((Get-Date).ToUniversalTime()) - Ping test passed. Response does not contain x-ms-azure-host-authorization header as expected"
+            }
+            else {
+                Write-Error "$((Get-Date).ToUniversalTime()) - Ping test failed. Response contains x-ms-azure-host-authorization header"
+                exit -1
+            }
         }
 
         $webRequest.Abort()

e2etest/GuestProxyAgentTest/Settings/TestScenarioSetting.cs

@@ -13,6 +13,7 @@ public class TestScenarioSetting
         internal string vmImageOffer = "";
         internal string vmImageSku = "";
         internal string vmImageVersion = "";
+        internal string sharedGalleryImageUniqueId = "";
         internal string suffixName = new Random().Next(1000).ToString();
         internal string testScenarioClassName = "GuestProxyAgentTest.TestScenarios.BVTScenario";
         internal int testScenarioTimeoutMilliseconds = 1000 * 60 * 120;
@@ -26,7 +27,8 @@ internal VMImageDetails VMImageDetails
                     Publisher = vmImagePublisher,
                     Offer = vmImageOffer,
                     Sku = vmImageSku,
-                    Version = vmImageVersion
+                    Version = vmImageVersion,
+                    SharedGalleryImageUniqueId = sharedGalleryImageUniqueId
                 };
             }
         }
@@ -54,11 +56,13 @@ public class VMImageDetails
         public string Offer { get; set; } = null!;
         public string Sku { get; set; } = null!;
         public string Version { get; set; } = null!;
+        public string SharedGalleryImageUniqueId { get; set; } = null!;
 
         public bool IsArm64
         {
             get
             {
+                // TODO: SharedGalleryImageUniqueId also contains architecture info, need to parse it when it's available
                 return (Offer == null ? false : Offer.Contains("arm64", StringComparison.OrdinalIgnoreCase)) ||
                   (Sku == null ? false : Sku.Contains("arm64", StringComparison.OrdinalIgnoreCase));
             }

e2etest/GuestProxyAgentTest/TestCases/EnableProxyAgentCase.cs

@@ -52,14 +52,34 @@ public override async Task StartAsync(TestCaseExecutionContext context)
             if (EnableProxyAgent)
             {
                 // property 'inVMAccessControlProfileReferenceId' cannot be used together with property 'mode'
-                patch.SecurityProfile.ProxyAgentSettings.WireServer = new HostEndpointSettings
+                if (string.IsNullOrEmpty(TestSetting.Instance.InVmWireServerAccessControlProfileReferenceId))
                 {
-                    InVmAccessControlProfileReferenceId = TestSetting.Instance.InVmWireServerAccessControlProfileReferenceId,
-                };
-                patch.SecurityProfile.ProxyAgentSettings.Imds = new HostEndpointSettings
+                    patch.SecurityProfile.ProxyAgentSettings.WireServer = new HostEndpointSettings
+                    {
+                        Mode = HostEndpointSettingsMode.Enforce,
+                    };
+                }
+                else
+                {
+                    patch.SecurityProfile.ProxyAgentSettings.WireServer = new HostEndpointSettings
+                    {
+                        InVmAccessControlProfileReferenceId = TestSetting.Instance.InVmWireServerAccessControlProfileReferenceId,
+                    };
+                }
+                if (string.IsNullOrEmpty(TestSetting.Instance.InVmIMDSAccessControlProfileReferenceId))
                 {
-                    InVmAccessControlProfileReferenceId = TestSetting.Instance.InVmIMDSAccessControlProfileReferenceId,
-                };
+                    patch.SecurityProfile.ProxyAgentSettings.Imds = new HostEndpointSettings
+                    {
+                        Mode = HostEndpointSettingsMode.Enforce,
+                    };
+                }
+                else
+                {
+                    patch.SecurityProfile.ProxyAgentSettings.Imds = new HostEndpointSettings
+                    {
+                        InVmAccessControlProfileReferenceId = TestSetting.Instance.InVmIMDSAccessControlProfileReferenceId,
+                    };
+                }
             }
 
             await vmr.UpdateAsync(Azure.WaitUntil.Completed, patch, cancellationToken: context.CancellationToken);

e2etest/GuestProxyAgentTest/TestMap/Test-Map-Linux.yml

@@ -5,4 +5,5 @@ testGroupList:
   - include: Ubuntu20-TestGroup.yml
   - include: Redhat90-TestGroup.yml
   - include: Suse15SP4-TestGroup.yml
-  - include: Rocky9-TestGroup.yml
+  - include: Rocky9-TestGroup.yml
+  - include: Ubuntu24-SGI-TestGroup.yml

e2etest/GuestProxyAgentTest/TestMap/Ubuntu24-SGI-TestGroup.yml

@@ -0,0 +1,9 @@
+groupName: Ubuntu24-SGI
+vmImagePublisher:  
+vmImageOffer: 
+vmImageSku: 
+vmImageVersion: 
+sharedGalleryImageUniqueId: /SharedGalleries/0a2c89a7-a44e-4cd0-b6ec-868432ad1d13-proxyagent/images/ubuntu-2404-gen2/versions/latest
+scenarios:
+  - className: GuestProxyAgentTest.TestScenarios.BakedInScenario
+    name: BakedInScenario

e2etest/GuestProxyAgentTest/TestScenarios/BakedInScenario.cs

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+using GuestProxyAgentTest.TestCases;
+using GuestProxyAgentTest.Utilities;
+using System.Threading.Channels;
+
+namespace GuestProxyAgentTest.TestScenarios
+{
+    public class BakedInScenario : TestScenarioBase
+    {
+        public override void TestScenarioSetup()
+        {
+            if (Constants.IS_WINDOWS())
+            {
+                throw new InvalidOperationException("GPA BakedIn scenario can only run on Linux VMs.");
+            }
+
+            var secureChannelEnabled = false;
+            EnableProxyAgentForNewVM = false;
+            AddTestCase(new GuestProxyAgentValidationCase("GuestProxyAgentValidationWithoutMSP", "disabled"));
+            AddTestCase(new IMDSPingTestCase("IMDSPingTestBeforeEnableMSP", secureChannelEnabled));
+
+            // enable secure channel after validation to test IMDS connectivity with secure channel enabled, 
+            AddTestCase(new EnableProxyAgentCase());
+            secureChannelEnabled = true;
+            AddTestCase(new GuestProxyAgentValidationCase("GuestProxyAgentValidationWithSecureChannelEnabled", "WireServer Enforce -  IMDS Enforce - HostGA Enforce"));
+            AddTestCase(new IMDSPingTestCase("IMDSPingTestBeforeReboot", secureChannelEnabled));
+
+            // then reboot to verify the secure channel state is preserved across reboots
+            AddTestCase(new RebootVMCase("RebootVMCaseAfterEnableMSP"));
+            AddTestCase(new IMDSPingTestCase("IMDSPingTestAfterReboot", secureChannelEnabled));
+        }
+    }
+}

e2etest/GuestProxyAgentTest/Utilities/TestMapReader.cs

@@ -49,6 +49,7 @@ public static List<TestScenarioSetting> ReadFlattenTestScenarioSettingFromTestMa
                    vmImagePublisher = group.VmImagePublisher,
                    vmImageSku = group.VmImageSku,
                    vmImageVersion = group.VmImageVersion,
+                   sharedGalleryImageUniqueId = group.SharedGalleryImageUniqueId,
                    testGroupName = group.GroupName,
                    testScenarioClassName = ele.ClassName,
                    testScenarioName = ele.Name,

e2etest/GuestProxyAgentTest/Utilities/VMBuilder.cs

@@ -226,6 +226,16 @@ private async Task<VirtualMachineData> DoCreateVMData(TestLogger logger, Resourc
                 NetworkProfile = await DoCreateVMNetWorkProfile(logger, rgr),
             };
 
+            var useSGI = !string.IsNullOrEmpty(this.testScenarioSetting.VMImageDetails.SharedGalleryImageUniqueId);
+            // Use Shared Gallery Image if SharedGalleryImageUniqueId is provided in the test setting
+            if (useSGI)
+            {
+                vmData.StorageProfile.ImageReference = new ImageReference()
+                {
+                    SharedGalleryImageUniqueId = new ResourceIdentifier(this.testScenarioSetting.VMImageDetails.SharedGalleryImageUniqueId)
+                };
+            }
+
             if (enableProxyAgent)
             {
                 vmData.SecurityProfile = new SecurityProfile()
@@ -243,12 +253,23 @@ private async Task<VirtualMachineData> DoCreateVMData(TestLogger logger, Resourc
                         },
                     }
                 };
-                if (!Constants.IS_WINDOWS())
+                if (!Constants.IS_WINDOWS() && !useSGI)
                 {
                     // Only Linux VMs support flag 'AddProxyAgentExtension',
                     // Windows VMs always have the GPA VM Extension installed when ProxyAgentSettings.Enabled is true.
                     vmData.SecurityProfile.ProxyAgentSettings.AddProxyAgentExtension = true;
                 }
+
+                // If the access control profile reference id is not set, set the mode to Enforce to make sure the proxy agent is working in expected way.
+                // This is for test images which don't have the access control profile set up, or for shared gallery images which can't reference the access control profile in other subscription.
+                if (string.IsNullOrEmpty(vmData.SecurityProfile.ProxyAgentSettings.WireServer.InVmAccessControlProfileReferenceId))
+                {
+                    vmData.SecurityProfile.ProxyAgentSettings.WireServer.Mode = HostEndpointSettingsMode.Enforce;
+                }
+                if (string.IsNullOrEmpty(vmData.SecurityProfile.ProxyAgentSettings.Imds.InVmAccessControlProfileReferenceId))
+                {
+                    vmData.SecurityProfile.ProxyAgentSettings.Imds.Mode = HostEndpointSettingsMode.Enforce;
+                }
             }
 
             if (Constants.IS_WINDOWS())

proxy_agent/src/key_keeper/key.rs

@@ -42,6 +42,7 @@ use std::{ffi::OsString, time::Duration};
 
 const AUDIT_MODE: &str = "audit";
 const ENFORCE_MODE: &str = "enforce";
+const DISABLED_MODE: &str = "disabled";
 //const ALLOW_DEFAULT_ACCESS: &str = "allow";
 //const DENY_DEFAULT_ACCESS: &str = "deny";
 
@@ -63,7 +64,7 @@ pub struct KeyStatus {
     // specifies what keys are expected for telemetry purposes.
     // Exact values are TBD, but could include things like user id.
     requiredClaimsHeaderPairs: Option<Vec<String>>,
-    // One of Disabled, Wireserver, WireserverAndImds. valid at version 1.0
+    // One of Disabled, Audit, Wireserver, WireserverAndImds. valid at version 1.0
     #[serde(skip_serializing_if = "Option::is_none")]
     pub secureChannelState: Option<String>,
     // Indicates if the secure channel is enabled. valid at version 2.0
@@ -605,19 +606,23 @@ impl KeyStatus {
             match &self.authorizationRules {
                 Some(rules) => match &rules.wireserver {
                     Some(item) => item.mode.to_lowercase(),
-                    None => "disabled".to_string(),
+                    None => DISABLED_MODE.to_string(),
                 },
-                None => "disabled".to_string(),
+                None => DISABLED_MODE.to_string(),
             }
         } else {
+            // in older version: secureChannelState indicates what endpoints have secure channel protections enabled.
+            // One of Disabled, Audit, Wireserver, WireserverAndImds.
             let state = match &self.secureChannelState {
                 Some(s) => s.to_lowercase(),
-                None => "disabled".to_string(),
+                None => DISABLED_MODE.to_string(),
             };
             if state == "wireserver" || state == "wireserverandimds" {
                 ENFORCE_MODE.to_string()
-            } else {
+            } else if state == "audit" {
                 AUDIT_MODE.to_string()
+            } else {
+                DISABLED_MODE.to_string()
             }
         }
     }
@@ -627,18 +632,25 @@ impl KeyStatus {
             match &self.authorizationRules {
                 Some(rules) => match &rules.imds {
                     Some(item) => item.mode.to_lowercase(),
-                    None => "disabled".to_string(),
+                    None => DISABLED_MODE.to_string(),
                 },
-                None => "disabled".to_string(),
+                None => DISABLED_MODE.to_string(),
             }
         } else {
+            // in older version: secureChannelState indicates what endpoints have secure channel protections enabled.
+            // One of Disabled, Audit, Wireserver, WireserverAndImds.
             let state = match &self.secureChannelState {
                 Some(s) => s.to_lowercase(),
-                None => "disabled".to_string(),
+                None => DISABLED_MODE.to_string(),
             };
-            if state == "wireserverandimds" {
+
+            if state == DISABLED_MODE {
+                DISABLED_MODE.to_string()
+            } else if state == "wireserverandimds" {
                 ENFORCE_MODE.to_string()
             } else {
+                // audit mode when secureChannelState is audit or wireserver,
+                // because in both cases IMDS has some level of protection, just not enforce.
                 AUDIT_MODE.to_string()
             }
         }
@@ -658,15 +670,17 @@ impl KeyStatus {
 impl Display for KeyStatus {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f,
-            "authorizationScheme: {}, keyDeliveryMethod: {}, keyGuid: {}, secureChannelState: {}, version: {}",
+            "authorizationScheme: {}, keyDeliveryMethod: {}, keyGuid: {}, secureChannelState: {}, version: {}, WireServerMode: {}, IMDSMode: {}",
             self.authorizationScheme,
             self.keyDeliveryMethod,
             match &self.keyGuid {
                 Some(s) => s.to_string(),
                 None => "None".to_string(),
             },
             self.get_secure_channel_state(),
-            self.version)
+            self.version,
+            self.get_wire_server_mode(),
+            self.get_imds_mode())
     }
 }
 
@@ -935,10 +949,44 @@ mod tests {
             "WireServer mode mismatch"
         );
         assert_eq!(status_v1.get_imds_mode(), "audit", "IMDS mode mismatch");
+
+        // Test the case when secureChannelState is Disabled, both WireServer and IMDS should be in disabled mode.
+        let status_response_v1 = r#"{
+            "authorizationScheme": "Azure-HMAC-SHA256",
+            "keyDeliveryMethod": "http",
+            "keyGuid": null,
+            "requiredClaimsHeaderPairs": null,
+            "secureChannelState": "Disabled",
+            "version": "1.0"
+        }"#;
+        let status_v1: KeyStatus = serde_json::from_str(status_response_v1).unwrap();
+        assert_eq!(
+            super::DISABLED_MODE.to_string(),
+            status_v1.get_wire_server_mode(),
+            "WireServer mode mismatch when secureChannelState is Disabled"
+        );
+        assert_eq!(
+            super::DISABLED_MODE.to_string(),
+            status_v1.get_imds_mode(),
+            "IMDS mode mismatch when secureChannelState is Disabled"
+        );
     }
 
     #[test]
     fn key_status_v2_test() {
+        let status_response = r#"{"authorizationRules":null,"authorizationScheme":"Azure-HMAC-SHA256","keyDeliveryMethod":"http","keyGuid":null,"keyIncarnationId":null,"requiredClaimsHeaderPairs":[],"secureChannelEnabled":false,"version":"2.0"}"#;
+        let status: KeyStatus = serde_json::from_str(status_response).unwrap();
+        assert_eq!(
+            "disabled",
+            status.get_wire_server_mode(),
+            "WireServer mode mismatch when secureChannelEnabled is false"
+        );
+        assert_eq!(
+            "disabled",
+            status.get_imds_mode(),
+            "IMDS mode mismatch when secureChannelEnabled is false"
+        );
+
         let status_response = r#"{
             "authorizationScheme": "Azure-HMAC-SHA256",
             "keyDeliveryMethod": "http",