From 5171131d549c29db7d465368dc9d407024b27830 Mon Sep 17 00:00:00 2001
From: Dongsu Park <dpark@linux.microsoft.com>
Date: Fri, 5 Jul 2024 12:11:58 +0200
Subject: [PATCH] Get username either by IMDS or from a local OVF file

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().

Explicitly get instance metadata to handle error of missing instance metadata
before dealing with setting ssh keys or hostname.
---
 libazureinit/src/error.rs |  4 +++
 src/main.rs               | 66 ++++++++++++++++++++++++++-------------
 2 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/libazureinit/src/error.rs b/libazureinit/src/error.rs
index f0ebd23..5bf6a8a 100644
--- a/libazureinit/src/error.rs
+++ b/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")]
diff --git a/src/main.rs b/src/main.rs
index 66bb4f6..3526484 100644
--- a/src/main.rs
+++ b/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,8 +89,23 @@ 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()?)?;
+    // 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 = match imds::query(&client).await {
+        Ok(m) => Some(m),
+        Err(_) => None,
+    };
+
+    let environment = match get_environment() {
+        Ok(env) => Some(env),
+        Err(_) => None,
+    };
+
+    let username =
+        get_username(instance_metadata.as_ref(), environment.as_ref())?;
 
     let mut file_path = "/home/".to_string();
     file_path.push_str(username.as_str());
@@ -99,11 +117,17 @@ async fn provision() -> Result<(), anyhow::Error> {
         || format!("Unabled to set an empty password for user '{username}'"),
     )?;
 
-    user::set_ssh_keys(instance_metadata.compute.public_keys, &username)
+    // 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)?;
+
+    user::set_ssh_keys(im.compute.public_keys, &username)
         .with_context(|| "Failed to write ssh public keys.")?;
 
     distro::set_hostname_with_hostnamectl(
-        instance_metadata.compute.os_profile.computer_name.as_str(),
+        im.compute.os_profile.computer_name.as_str(),
     )
     .with_context(|| "Failed to set hostname.")?;