Refactor emit_kvp_kayer and otel_layer with improved per-layer filtering, adjust key_buf logic to avoid indexing, refine tracing and function calls, and streamline start/end time declarations as per PR feedback.

peytonr18 · peytonr18 · commit 0536a3b68e27 · 2025-01-03T16:12:19.000-08:00
diff --git a/libazureinit/src/config.rs b/libazureinit/src/config.rs
@@ -292,7 +292,7 @@ impl Config {
             Figment::from(Serialized::defaults(Config::default()));
 
         if base_path.exists() {
-            tracing::debug!(path=?base_path, "Loading base configuration file");
+            tracing::info!(path=?base_path, "Loading base configuration file");
             figment = figment.merge(Toml::file(base_path));
         } else {
             tracing::warn!(
@@ -307,7 +307,7 @@ impl Config {
             if cli_path.is_dir() {
                 figment = Self::merge_toml_directory(figment, cli_path)?;
             } else {
-                tracing::debug!(
+                tracing::info!(
                     "Merging configuration file from CLI: {:?}",
                     cli_path
                 );
@@ -318,7 +318,10 @@ impl Config {
         figment
             .extract::<Config>()
             .map(|config| {
-                tracing::info!("Configuration successfully loaded.");
+                tracing::info!(
+                    target: "kvp",
+                    "Configuration successfully loaded."
+                );
                 config
             })
             .map_err(|e| {
@@ -355,12 +358,12 @@ impl Config {
             entries.sort();
 
             for path_entry in entries {
-                tracing::debug!("Merging configuration file: {:?}", path_entry);
+                tracing::info!("Merging configuration file: {:?}", path_entry);
                 figment = figment.merge(Toml::file(path_entry));
             }
             Ok(figment)
         } else {
-            tracing::debug!("Directory {:?} not found, skipping.", dir_path);
+            tracing::info!("Directory {:?} not found, skipping.", dir_path);
             Ok(figment.clone())
         }
     }
diff --git a/libazureinit/src/http.rs b/libazureinit/src/http.rs
@@ -111,7 +111,7 @@ async fn request(
         let now = std::time::Instant::now();
         let mut attempt =  0_u32;
         loop {
-            let span = tracing::debug_span!("request", attempt, http_status = tracing::field::Empty);
+            let span = tracing::info_span!("request", attempt, http_status = tracing::field::Empty);
             let req = request.try_clone().expect("The request body MUST be clone-able");
             match client
                 .execute(req)
@@ -121,12 +121,12 @@ async fn request(
                         let _enter = span.enter();
                         let statuscode = response.status();
                         span.record("http_status", statuscode.as_u16());
-                        tracing::info!(url=response.url().as_str(), "HTTP response received");
+                        tracing::info!(target: "kvp", url=response.url().as_str(), "HTTP response received");
 
                         match response.error_for_status() {
                             Ok(response) => {
                                 if statuscode == StatusCode::OK {
-                                    tracing::info!("HTTP response succeeded with status {}", statuscode);
+                                    tracing::info!(target: "kvp", "HTTP response succeeded with status {}", statuscode);
                                     return Ok((response, retry_for.saturating_sub(now.elapsed() + retry_interval)));
                                 }
                             },
diff --git a/libazureinit/src/provision/ssh.rs b/libazureinit/src/provision/ssh.rs
@@ -86,7 +86,11 @@ pub(crate) fn provision_ssh(
 
     chown(&ssh_dir, Some(user.uid), Some(user.gid))?;
 
-    tracing::info!("Using authorized_keys path: {:?}", authorized_keys_path);
+    tracing::info!(
+        target: "kvp",
+        "Using authorized_keys path: {:?}",
+        authorized_keys_path
+    );
 
     let mut authorized_keys = File::create(&authorized_keys_path)?;
     authorized_keys.set_permissions(Permissions::from_mode(0o600))?;
@@ -139,6 +143,7 @@ fn run_sshd_command(
     match sshd_config_command_runner() {
         Ok(output) if output.status.success() => {
             info!(
+                target: "kvp",
                 stdout_length = output.stdout.len(),
                 "Executed sshd -G successfully",
             );
@@ -184,6 +189,7 @@ fn extract_authorized_keys_file_path(stdout: &[u8]) -> Option<String> {
         if line.starts_with("authorizedkeysfile") {
             let keypath = line.split_whitespace().nth(1).map(|s| {
                 info!(
+                    target: "kvp",
                     authorizedkeysfile = %s,
                     "Using sshd's authorizedkeysfile path configuration"
                 );
diff --git a/libazureinit/src/provision/user.rs b/libazureinit/src/provision/user.rs
@@ -135,13 +135,15 @@ fn user_exists(username: &str) -> Result<bool, Error> {
 fn useradd(user: &User) -> Result<(), Error> {
     if user_exists(&user.name)? {
         tracing::info!(
+            target: "kvp",
             "User '{}' already exists. Skipping user creation.",
             user.name
         );
 
         let group_list = user.groups.join(",");
 
         tracing::info!(
+            target: "kvp",
             "User '{}' is being added to the following groups: {}",
             user.name,
             group_list
diff --git a/src/kvp.rs b/src/kvp.rs
@@ -13,7 +13,7 @@
 use std::{
     fmt::{self as std_fmt, Write as std_write},
     fs::{File, OpenOptions},
-    io::{self, Error, ErrorKind, Write},
+    io::{self, ErrorKind, Write},
     os::unix::fs::MetadataExt,
     path::Path,
     time::{Duration, Instant, SystemTime, UNIX_EPOCH},
@@ -81,10 +81,10 @@ impl Visit for StringVisitor<'_> {
     /// * `value` - The debug value associated with the field.
     fn record_debug(
         &mut self,
-        _field: &tracing::field::Field,
+        field: &tracing::field::Field,
         value: &dyn std_fmt::Debug,
     ) {
-        write!(self.string, "{}={:?}; ", _field.name(), value)
+        write!(self.string, "{}={:?}; ", field.name(), value)
             .expect("Writing to a string should never fail");
     }
 }
@@ -134,7 +134,7 @@ impl EmitKVPLayer {
     /// # Arguments
     /// * `file_path` - The file path where the KVP data will be stored.
     ///
-    pub fn new(file_path: std::path::PathBuf) -> Result<Self, std::io::Error> {
+    pub fn new(file_path: std::path::PathBuf) -> Result<Self, anyhow::Error> {
         truncate_guest_pool_file(&file_path)?;
 
         let file = OpenOptions::new()
@@ -309,21 +309,11 @@ where
                 let start_time =
                     end_time.checked_sub(elapsed).unwrap_or(UNIX_EPOCH);
 
-                let start_time_dt = DateTime::<Utc>::from(
-                    UNIX_EPOCH
-                        + start_time
-                            .duration_since(UNIX_EPOCH)
-                            .unwrap_or_default(),
-                )
-                .format("%Y-%m-%dT%H:%M:%S%.3fZ");
+                let start_time_dt = DateTime::<Utc>::from(start_time)
+                    .format("%Y-%m-%dT%H:%M:%S%.3fZ");
 
-                let end_time_dt = DateTime::<Utc>::from(
-                    UNIX_EPOCH
-                        + end_time
-                            .duration_since(UNIX_EPOCH)
-                            .unwrap_or_default(),
-                )
-                .format("%Y-%m-%dT%H:%M:%S%.3fZ");
+                let end_time_dt = DateTime::<Utc>::from(end_time)
+                    .format("%Y-%m-%dT%H:%M:%S%.3fZ");
 
                 let event_value = format!(
                     "Start: {} | End: {} | Status: {}",
@@ -362,41 +352,51 @@ fn generate_event_key(
 /// exceeds the allowed size, it is split into multiple slices for encoding.
 /// This is used for logging events to a KVP file.
 ///
+/// # Note
+/// - The key must be zero-padded to `HV_KVP_EXCHANGE_MAX_KEY_SIZE` to meet Hyper-V's expected formatting.
+///
 /// # Arguments
 /// * `key` - The key as a string slice.
 /// * `value` - The value associated with the key.
 fn encode_kvp_item(key: &str, value: &str) -> Vec<Vec<u8>> {
-    let key_bytes = key.as_bytes();
-    let value_bytes = value.as_bytes();
-
-    let key_len = key_bytes.len().min(HV_KVP_EXCHANGE_MAX_KEY_SIZE);
-    let mut key_buf = vec![0u8; HV_KVP_EXCHANGE_MAX_KEY_SIZE];
-    key_buf[..key_len].copy_from_slice(&key_bytes[..key_len]);
-
-    let mut kvp_slices = Vec::new();
+    let key_buf = key
+        .as_bytes()
+        .iter()
+        .take(HV_KVP_EXCHANGE_MAX_KEY_SIZE)
+        .chain(
+            vec![0_u8; HV_KVP_EXCHANGE_MAX_KEY_SIZE.saturating_sub(key.len())]
+                .iter(),
+        )
+        .copied()
+        .collect::<Vec<_>>();
+
+    debug_assert!(key_buf.len() == HV_KVP_EXCHANGE_MAX_KEY_SIZE);
+
+    let kvp_slices = value
+        .as_bytes()
+        .chunks(HV_KVP_AZURE_MAX_VALUE_SIZE)
+        .map(|chunk| {
+            let mut buffer = Vec::with_capacity(
+                HV_KVP_EXCHANGE_MAX_KEY_SIZE + HV_KVP_EXCHANGE_MAX_VALUE_SIZE,
+            );
+            buffer.extend_from_slice(&key_buf);
+            buffer.extend_from_slice(chunk);
+            while buffer.len()
+                < HV_KVP_EXCHANGE_MAX_KEY_SIZE + HV_KVP_EXCHANGE_MAX_VALUE_SIZE
+            {
+                buffer.push(0);
+            }
 
-    for chunk in value_bytes.chunks(HV_KVP_AZURE_MAX_VALUE_SIZE) {
-        let mut value_buf = vec![0u8; HV_KVP_EXCHANGE_MAX_VALUE_SIZE];
-        value_buf[..chunk.len()].copy_from_slice(chunk);
+            buffer
+        })
+        .collect::<Vec<Vec<u8>>>();
 
-        kvp_slices.push(encode_kvp_slice(key_buf.clone(), value_buf));
-    }
+    debug_assert!(kvp_slices.iter().all(|kvp| kvp.len()
+        == HV_KVP_EXCHANGE_MAX_KEY_SIZE + HV_KVP_EXCHANGE_MAX_VALUE_SIZE));
 
     kvp_slices
 }
 
-/// Combines the key and value of a KVP into a single byte slice, ensuring
-/// proper formatting for consumption by hv_kvp_daemon service,
-/// which typically reads from /var/lib/hyperv/.kvp_pool_1.
-fn encode_kvp_slice(key: Vec<u8>, value: Vec<u8>) -> Vec<u8> {
-    let mut buffer = Vec::with_capacity(
-        HV_KVP_EXCHANGE_MAX_KEY_SIZE + HV_KVP_EXCHANGE_MAX_VALUE_SIZE,
-    );
-    buffer.extend_from_slice(&key);
-    buffer.extend_from_slice(&value);
-    buffer
-}
-
 /// Decodes a KVP byte slice into its corresponding key and value strings.
 /// This is useful for inspecting or logging raw KVP data.
 #[cfg(test)]
@@ -412,18 +412,23 @@ pub fn decode_kvp_item(
     }
 
     let key = String::from_utf8(
-        record_data[0..HV_KVP_EXCHANGE_MAX_KEY_SIZE].to_vec(),
+        record_data
+            .iter()
+            .take(HV_KVP_EXCHANGE_MAX_KEY_SIZE)
+            .cloned()
+            .collect::<Vec<_>>(),
     )
     .unwrap_or_else(|_| String::new())
     .trim_end_matches('\x00')
     .to_string();
 
     let value = String::from_utf8(
-        record_data[HV_KVP_EXCHANGE_MAX_KEY_SIZE..]
+        record_data
             .iter()
+            .skip(HV_KVP_EXCHANGE_MAX_KEY_SIZE)
             .take(HV_KVP_AZURE_MAX_VALUE_SIZE)
             .cloned()
-            .collect(),
+            .collect::<Vec<_>>(),
     )
     .unwrap_or_else(|_| String::new())
     .trim_end_matches('\x00')
@@ -435,12 +440,9 @@ pub fn decode_kvp_item(
 /// Truncates the guest pool KVP file if it contains stale data (i.e., data
 /// older than the system's boot time). Logs whether the file was truncated
 /// or no action was needed.
-fn truncate_guest_pool_file(kvp_file: &Path) -> Result<(), Error> {
-    let boot_time = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .map_err(|e| Error::new(std::io::ErrorKind::Other, e))?
-        .as_secs()
-        - get_uptime()?.as_secs();
+fn truncate_guest_pool_file(kvp_file: &Path) -> Result<(), anyhow::Error> {
+    let boot_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs()
+        - get_uptime().as_secs();
 
     match kvp_file.metadata() {
         Ok(metadata) => {
@@ -461,7 +463,8 @@ fn truncate_guest_pool_file(kvp_file: &Path) -> Result<(), Error> {
             return Ok(());
         }
         Err(e) => {
-            return Err(e);
+            return Err(anyhow::Error::from(e)
+                .context("Failed to access file metadata"));
         }
     }
 
@@ -471,12 +474,12 @@ fn truncate_guest_pool_file(kvp_file: &Path) -> Result<(), Error> {
 /// Retrieves the system's uptime using the `sysinfo` crate, returning the duration
 /// since the system booted. This can be useful for time-based calculations or checks,
 /// such as determining whether data is stale or calculating the approximate boot time.
-fn get_uptime() -> Result<Duration, Error> {
+fn get_uptime() -> Duration {
     let mut system = System::new();
     system.refresh_system();
 
     let uptime_seconds = system.uptime();
-    Ok(Duration::from_secs(uptime_seconds))
+    Duration::from_secs(uptime_seconds)
 }
 
 #[cfg(test)]
diff --git a/src/logging.rs b/src/logging.rs
@@ -7,6 +7,7 @@ use opentelemetry_sdk::trace::{
 };
 use tracing::{event, Level};
 use tracing_opentelemetry::OpenTelemetryLayer;
+use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::fmt::format::FmtSpan;
 use tracing_subscriber::{
     fmt, layer::SubscriberExt, EnvFilter, Layer, Registry,
@@ -26,15 +27,17 @@ pub fn initialize_tracing() -> sdktrace::Tracer {
 pub fn setup_layers(
     tracer: sdktrace::Tracer,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let otel_layer = OpenTelemetryLayer::new(tracer);
+    let otel_layer = OpenTelemetryLayer::new(tracer)
+        .with_filter(EnvFilter::from_env("AZURE_INIT_LOG"));
+
+    let kvp_filter = EnvFilter::builder()
+        .with_default_directive(LevelFilter::INFO.into())
+        .parse_lossy("azure_init=INFO,libazureinit=WARN,kvp=INFO");
 
     let emit_kvp_layer = match EmitKVPLayer::new(std::path::PathBuf::from(
         "/var/lib/hyperv/.kvp_pool_1",
     )) {
-        Ok(layer) => {
-            let filter = EnvFilter::new("INFO");
-            Some(layer.with_filter(filter))
-        }
+        Ok(layer) => Some(layer.with_filter(kvp_filter)),
         Err(e) => {
             event!(Level::ERROR, "Failed to initialize EmitKVPLayer: {}. Continuing without KVP logging.", e);
             None
diff --git a/src/main.rs b/src/main.rs
@@ -132,7 +132,7 @@ async fn main() -> ExitCode {
     }
 }
 
-#[instrument(name = "root")]
+#[instrument(name = "root", skip_all)]
 async fn provision(config: Config, opts: Cli) -> Result<(), anyhow::Error> {
     let system = System::new();
     let kernel_version = system
@@ -141,13 +141,12 @@ async fn provision(config: Config, opts: Cli) -> Result<(), anyhow::Error> {
     let os_version = system
         .os_version()
         .unwrap_or("Unknown OS Version".to_string());
-    let azure_init_version = env!("CARGO_PKG_VERSION");
 
     tracing::info!(
         "Kernel Version: {}, OS Version: {}, Azure-Init Version: {}",
         kernel_version,
         os_version,
-        azure_init_version
+        VERSION
     );
 
     let mut default_headers = header::HeaderMap::new();
