Get username either by IMDS or from a local OVF file

libazureinit/src/error.rs

@@ -27,6 +27,10 @@ pub enum Error {
     Nix(#[from] nix::Error),
     #[error("The user {user} does not exist")]
     UserMissing { user: String },
+    #[error("failed to get username from IMDS or local OVF files")]
+    UsernameFailure,
+    #[error("failed to get instance metadata from IMDS")]
+    InstanceMetadataFailure,
     #[error("Provisioning a user with a non-empty password is not supported")]
     NonEmptyPassword,
     #[error("Unable to get list of block devices")]

src/main.rs

@@ -34,25 +34,28 @@ fn get_environment() -> Result<Environment, anyhow::Error> {
 }
 
 fn get_username(
-    instance_metadata: &InstanceMetadata,
-    environment: &Environment,
+    instance_metadata: Option<&InstanceMetadata>,
+    environment: Option<&Environment>,
 ) -> Result<String, anyhow::Error> {
-    if instance_metadata
-        .compute
-        .os_profile
-        .disable_password_authentication
-    {
-        // password authentication is disabled
-        Ok(instance_metadata.compute.os_profile.admin_username.clone())
-    } else {
-        // password authentication is enabled
-
-        Ok(environment
-            .clone()
-            .provisioning_section
-            .linux_prov_conf_set
-            .username)
+    if let Some(metadata) = instance_metadata {
+        if metadata.compute.os_profile.disable_password_authentication {
+            // If password authentication is disabled,
+            // simply read from IMDS metadata if available.
+            return Ok(metadata.compute.os_profile.admin_username.clone());
+        }
+        // If password authentication is enabled,
+        // fall back to reading from OVF environment file.
     }
+
+    // Read username from OVF environment via mounted local device.
+    environment
+        .map(|env| {
+            env.clone()
+                .provisioning_section
+                .linux_prov_conf_set
+                .username
+        })
+        .ok_or(LibError::UsernameFailure.into())
 }
 
 #[tokio::main]
@@ -86,11 +89,27 @@ async fn provision() -> Result<(), anyhow::Error> {
         .default_headers(default_headers)
         .build()?;
 
-    let instance_metadata = imds::query(&client).await?;
-    let username = get_username(&instance_metadata, &get_environment()?)?;
-    let user = User::new(username, instance_metadata.compute.public_keys);
+    // Username can be obtained either via fetching instance metadata from IMDS
+    // or mounting a local device for OVF environment file. It should not fail
+    // immediately in a single failure, instead it should fall back to the other
+    // mechanism. So it is not a good idea to use `?` for query() or
+    // get_environment().
+    let instance_metadata = imds::query(&client).await.ok();
+
+    let environment = get_environment().ok();
+
+    let username =
+        get_username(instance_metadata.as_ref(), environment.as_ref())?;
+
+    // It is necessary to get the actual instance metadata after getting username,
+    // as it is not desirable to immediately return error before get_username.
+    let im = instance_metadata
+        .clone()
+        .ok_or::<LibError>(LibError::InstanceMetadataFailure)?;
+
+    let user = User::new(username, im.compute.public_keys);
 
-    Provision::new(instance_metadata.compute.os_profile.computer_name, user)
+    Provision::new(im.compute.os_profile.computer_name, user)
         .hostname_provisioners([
             #[cfg(feature = "hostnamectl")]
             HostnameProvisioner::Hostnamectl,

tests/functional_tests.sh

@@ -14,6 +14,7 @@ VM_SIZE="${VM_SIZE:-Standard_D2lds_v5}"
 VM_ADMIN_USERNAME="${VM_ADMIN_USERNAME:-azureuser}"
 AZURE_SSH_KEY_NAME="${AZURE_SSH_KEY_NAME:-azure-ssh-key}"
 VM_NAME_WITH_TIMESTAMP=$VM_NAME-$EPOCH
+VM_SECURITY_TYPE="${VM_SECURITY_TYPE:-TrustedLaunch}"
 
 set -e
 
@@ -40,32 +41,33 @@ else
 fi
 
 # Set the subscription you want to use
-az account set --subscription $SUBSCRIPTION_ID
+az account set --subscription "$SUBSCRIPTION_ID"
 
 # Create resource group
-az group create -g $RG -l $LOCATION
+az group create -g "$RG" -l "$LOCATION"
 
 echo "Creating VM..."
-az vm create -n $VM_NAME_WITH_TIMESTAMP \
--g $RG \
---image $VM_IMAGE \
---size $VM_SIZE \
---admin-username $VM_ADMIN_USERNAME \
---ssh-key-value $PATH_TO_PUBLIC_SSH_KEY \
---public-ip-sku Standard
+az vm create -n "$VM_NAME_WITH_TIMESTAMP" \
+-g "$RG" \
+--image "$VM_IMAGE" \
+--size "$VM_SIZE" \
+--admin-username "$VM_ADMIN_USERNAME" \
+--ssh-key-value "$PATH_TO_PUBLIC_SSH_KEY" \
+--public-ip-sku Standard \
+--security-type "$VM_SECURITY_TYPE"
 echo "VM successfully created"
 
 echo "Sleeping to ensure SSH access set up"
 sleep 15
 
 echo "Getting VM Public IP Address..."
-PUBLIC_IP=$(az vm show -d -g $RG -n $VM_NAME_WITH_TIMESTAMP --query publicIps -o tsv)
-echo $PUBLIC_IP
+PUBLIC_IP=$(az vm show -d -g "$RG" -n "$VM_NAME_WITH_TIMESTAMP" --query publicIps -o tsv)
+echo "$PUBLIC_IP"
 
-scp -o StrictHostKeyChecking=no -i $PATH_TO_PRIVATE_SSH_KEY ./target/debug/functional_tests $VM_ADMIN_USERNAME@$PUBLIC_IP:~
+scp -o StrictHostKeyChecking=no -i "$PATH_TO_PRIVATE_SSH_KEY" ./target/debug/functional_tests "$VM_ADMIN_USERNAME"@"$PUBLIC_IP":~
 
 echo "Logging into VM..."
-ssh -o StrictHostKeyChecking=no -i $PATH_TO_PRIVATE_SSH_KEY $VM_ADMIN_USERNAME@$PUBLIC_IP 'sudo ./functional_tests test_user' 
+ssh -o StrictHostKeyChecking=no -i "$PATH_TO_PRIVATE_SSH_KEY" "$VM_ADMIN_USERNAME"@"$PUBLIC_IP" 'sudo ./functional_tests test_user'
 
 # Delete the resource group
-az group delete -g $RG --yes --no-wait
+az group delete -g "$RG" --yes --no-wait