diff --git a/libazureinit-kvp/Cargo.toml b/libazureinit-kvp/Cargo.toml index d58f862d..3c0381cd 100644 --- a/libazureinit-kvp/Cargo.toml +++ b/libazureinit-kvp/Cargo.toml @@ -13,13 +13,15 @@ chrono = { version = "0.4", default-features = false, features = ["clock", "std" clap = { version = "4.5.21", features = ["derive"] } csv = "1" libc = "0.2" +serde = { version = "1.0", features = ["derive"] } serde_json = "1.0.96" tracing = "0.1.40" -uuid = "1.3" +uuid = { version = "1.3", features = ["v4"] } [dev-dependencies] rstest = { version = "0.26", default-features = false } tempfile = "3" +uuid = "1.3" [lib] name = "libazureinit_kvp" diff --git a/libazureinit-kvp/src/cli.rs b/libazureinit-kvp/src/cli.rs index fa347cc7..50558293 100644 --- a/libazureinit-kvp/src/cli.rs +++ b/libazureinit-kvp/src/cli.rs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use std::fmt::Write as _; use std::fs; use std::io::{self, Read, Write}; use std::path::PathBuf; @@ -10,8 +11,8 @@ use clap::{Parser, Subcommand, ValueEnum}; use serde_json::json; use crate::{ - write_report, KvpError, KvpPool, KvpPoolStore, PoolMode, - ProvisioningReport, ReportPpsType, + write_report, DiagnosticEvent, DiagnosticRecord, DiagnosticsKvp, KvpError, + KvpPool, KvpPoolStore, PoolMode, ProvisioningReport, ReportPpsType, }; const EXIT_OK: u8 = 0; @@ -89,7 +90,38 @@ enum Command { /// Print store metadata. Info, /// Print every record in insertion order as KEY=VALUE lines. - Dump, + /// + /// With --parse-diagnostics, reassemble chunked diagnostic events and + /// decode each record instead of printing raw KEY=VALUE lines. + Dump { + /// Reassemble chunked diagnostic events and decode each record as + /// an azure-init event, cloud-init event, raw, or malformed entry. + #[arg(long)] + parse_diagnostics: bool, + /// Also print raw (non-event) records such as PROVISIONING_REPORT. + /// Only applies to the unfiltered view; --name/--tail produce an + /// azure-init events-only view where raw records never appear. + #[arg( + long, + requires = "parse_diagnostics", + conflicts_with_all = ["name", "tail"] + )] + include_raw: bool, + /// Only show azure-init events whose name contains this + /// substring. + #[arg(long, requires = "parse_diagnostics")] + name: Option, + /// Print only the last COUNT azure-init events (default 20 when + /// COUNT is omitted). + #[arg( + short = 'n', + long = "tail", + num_args = 0..=1, + default_missing_value = "20", + requires = "parse_diagnostics" + )] + tail: Option, + }, /// Print key=last_value entries sorted by key. Entries, /// Print the last value for KEY (exit 1 if missing). @@ -102,6 +134,23 @@ enum Command { key: String, value: String, }, + /// Emit an azure-init diagnostic event: a structured KVP entry keyed + /// `||||||`. + /// Distinct from the raw `write` command. + Emit { + /// Event name, e.g. user:create_user. + #[arg(long)] + name: String, + /// Event message (stored as the record value). + #[arg(long)] + message: String, + /// VM identifier (defaults to the current VM's ID). + #[arg(long)] + vm_id: Option, + /// Event-key prefix (defaults to the reporting agent identifier). + #[arg(long)] + prefix: Option, + }, /// Replace the pool from KEY=VALUE lines read from --file or stdin. Load { /// Read records from PATH instead of stdin. @@ -121,11 +170,16 @@ enum Command { #[arg(required = true)] keys: Vec, }, - /// Clear the pool. Pass --if-stale to clear only when stale. + /// Clear the pool. Pass --if-stale to clear only when stale, or + /// --diagnostics to remove only diagnostic event keys. Clear { /// Only clear if the store is currently stale. - #[arg(long = "if-stale")] + #[arg(long = "if-stale", conflicts_with = "diagnostics")] if_stale: bool, + /// Remove every diagnostic event key (valid or malformed), + /// leaving raw records such as PROVISIONING_REPORT intact. + #[arg(long)] + diagnostics: bool, }, /// Print whether the pool is stale (exit 0 if stale, 1 otherwise). IsStale, @@ -222,7 +276,19 @@ fn dispatch(cli: Cli, stdout: &mut W) -> Result { match cli.command { Command::Info => info(&store, stdout, output), - Command::Dump => dump(&store, stdout, output), + Command::Dump { + parse_diagnostics, + include_raw, + name, + tail, + } => { + let parse = parse_diagnostics.then_some(ParseDiagnosticsArgs { + include_raw, + name, + tail, + }); + dump(&store, stdout, parse, output) + } Command::Entries => entries(&store, stdout, output), Command::Read { key } => read(&store, stdout, &key, output), Command::Write { append, key, value } => { @@ -233,14 +299,25 @@ fn dispatch(cli: Cli, stdout: &mut W) -> Result { } Ok(EXIT_OK) } + Command::Emit { + name, + message, + vm_id, + prefix, + } => emit(&store, name, message, vm_id, prefix), Command::Load { file } => load(&store, file), Command::AppendMultiple { file } => append_multiple(&store, file), Command::Delete { key } => delete(&store, stdout, &key, output), Command::DeleteMultiple { keys } => { delete_multiple(&store, stdout, keys, output) } - Command::Clear { if_stale } => { - if if_stale { + Command::Clear { + if_stale, + diagnostics, + } => { + if diagnostics { + DiagnosticsKvp::new(store.clone(), "", "").clear()?; + } else if if_stale { store.clear_if_stale()?; } else { store.clear()?; @@ -311,12 +388,27 @@ fn info( } Ok(EXIT_OK) } +struct ParseDiagnosticsArgs { + include_raw: bool, + name: Option, + tail: Option, +} fn dump( store: &KvpPoolStore, stdout: &mut W, + parse: Option, output: OutputMode, ) -> Result { + if let Some(parse) = parse { + if parse.name.is_some() || parse.tail.is_some() { + return diagnostics_events( + store, stdout, parse.name, parse.tail, output, + ); + } + return diagnostics_records(store, stdout, parse.include_raw, output); + } + let records = store.dump()?; match output { OutputMode::Text => { @@ -425,6 +517,161 @@ fn is_stale( Ok(if stale { EXIT_OK } else { EXIT_NOT_FOUND }) } +fn diagnostics_records( + store: &KvpPoolStore, + stdout: &mut W, + include_raw: bool, + output: OutputMode, +) -> Result { + let diagnostics = DiagnosticsKvp::new(store.clone(), "", ""); + let records: Vec<_> = diagnostics + .records()? + .into_iter() + .filter(|record| { + include_raw || !matches!(record, DiagnosticRecord::Raw { .. }) + }) + .collect(); + + match output { + OutputMode::Text => { + for record in &records { + let line = match record { + DiagnosticRecord::Decoded { event, chunks } => { + let mut line = diagnostics_event_text(event); + let _ = write!( + line, + " chunks={chunks} message={}", + event.message + ); + line + } + DiagnosticRecord::Raw { key, value } => { + format!("raw key={key} value={value}") + } + DiagnosticRecord::Malformed { key, value, reason } => { + format!( + "malformed key={key} reason={reason} \ + value={value}" + ) + } + }; + writeln!(stdout, "{line}")?; + } + } + OutputMode::Json => { + let array: Vec<_> = + records.iter().map(diagnostics_record_json).collect(); + writeln_json(stdout, &serde_json::Value::Array(array))?; + } + } + Ok(EXIT_OK) +} + +fn diagnostics_events( + store: &KvpPoolStore, + stdout: &mut W, + name: Option, + tail: Option, + output: OutputMode, +) -> Result { + let diagnostics = DiagnosticsKvp::new(store.clone(), "", ""); + let mut events = diagnostics.events()?; + + if let Some(needle) = name.as_deref() { + events.retain(|event| event.name.contains(needle)); + } + if let Some(count) = tail { + let excess = events.len().saturating_sub(count); + events.drain(..excess); + } + + match output { + OutputMode::Text => { + for event in &events { + let mut line = diagnostics_event_text(event); + let _ = write!(line, " message={}", event.message); + writeln!(stdout, "{line}")?; + } + } + OutputMode::Json => { + let array: Vec<_> = + events.iter().map(diagnostics_event_json).collect(); + writeln_json(stdout, &serde_json::Value::Array(array))?; + } + } + Ok(EXIT_OK) +} + +/// Render a [`DiagnosticEvent`] as a single text line of `key=value` +/// fields, omitting optional fields the source did not provide. +fn diagnostics_event_text(event: &DiagnosticEvent) -> String { + let mut line = format!( + "event kind={} agent={} boot_epoch={}", + event.kind, event.agent, event.boot_epoch + ); + if let Some(vm_id) = &event.vm_id { + let _ = write!(line, " vm_id={vm_id}"); + } + let _ = write!(line, " name={} event_id={}", event.name, event.event_id); + if let Some(ts) = &event.timestamp { + let _ = write!(line, " timestamp={ts}"); + } + if let Some(result) = &event.result { + let _ = write!(line, " result={result}"); + } + if let Some(duration) = event.duration { + let _ = write!(line, " duration={duration}"); + } + line +} + +/// Render a [`DiagnosticRecord`] as a JSON object. +fn diagnostics_record_json(record: &DiagnosticRecord) -> serde_json::Value { + match record { + DiagnosticRecord::Decoded { event, chunks } => { + let mut value = diagnostics_event_json(event); + if let serde_json::Value::Object(map) = &mut value { + map.insert("record".to_string(), json!("event")); + map.insert("chunks".to_string(), json!(chunks)); + } + value + } + DiagnosticRecord::Raw { key, value } => json!({ + "record": "raw", + "key": key, + "value": value, + }), + DiagnosticRecord::Malformed { key, value, reason } => json!({ + "record": "malformed", + "key": key, + "value": value, + "reason": reason, + }), + } +} + +/// Render a [`DiagnosticEvent`] as a JSON object (without chunk count), +/// omitting optional fields the source did not provide. +fn diagnostics_event_json(event: &DiagnosticEvent) -> serde_json::Value { + serde_json::to_value(event) + .expect("DiagnosticEvent always serializes to a JSON object") +} + +/// Emit an azure-init diagnostic event with the given fields. +fn emit( + store: &KvpPoolStore, + name: String, + message: String, + vm_id: Option, + prefix: Option, +) -> Result { + let vm_id = resolve_vm_id(vm_id)?; + let prefix = prefix.unwrap_or_else(|| DEFAULT_AGENT.to_string()); + let diagnostics = DiagnosticsKvp::new(store.clone(), vm_id, prefix); + diagnostics.emit_event(name, message)?; + Ok(EXIT_OK) +} + fn report_success( store: &KvpPoolStore, vm_id: Option, @@ -492,26 +739,15 @@ fn resolve_vm_id_with( #[derive(Clone, Debug, PartialEq, Eq)] struct SupportingData(Vec<(String, String)>); -/// Parse a `--supporting-data` argument into its `key=value` pairs. -/// -/// Fields are comma-separated. A value may be wrapped in matching single or -/// double quotes so it can contain literal commas; the quotes are honored -/// only when they wrap the *entire* value (the opening quote immediately -/// follows `=` and the matching quote ends the field) and are stripped from -/// the stored value. Empty fields (such as a trailing comma) are ignored. -/// -/// Supported (input -> parsed pairs): -/// - `k=v` -> `k`=`v` -/// - `k1=v1,k2=v2` -> `k1`=`v1`, `k2`=`v2` -/// - `k='a,b'` or `k="a,b"` -> `k`=`a,b` (quotes protect the comma) -/// - `k=a'b` -> `k`=`a'b` (a quote not at the value start is literal) -/// - `k=v,` -> `k`=`v` (trailing/empty field ignored) +/// Parse a `--supporting-data` argument into its comma-separated +/// `key=value` pairs. A value wrapped in matching single/double quotes may +/// contain literal commas (the quotes must wrap the whole value and are +/// stripped); empty fields are ignored. /// -/// Rejected: -/// - `novalue` -> missing `=` -/// - `=v` -> empty key -/// - `k='a,b` -> unterminated quote -/// - `k='a,b'x` -> characters after a quoted value +/// Supported: `k=v`; `k1=v1,k2=v2`; `k='a,b'` or `k="a,b"` -> `k`=`a,b`; +/// `k=a'b` -> literal quote; `k=v,` -> trailing field ignored. +/// Rejected: `novalue` (no `=`), `=v` (empty key), `k='a,b` +/// (unterminated quote), `k='a,b'x` (chars after a quoted value). fn parse_supporting_data(raw: &str) -> Result { let mut pairs = Vec::new(); for field in split_supporting_data_fields(raw)? { @@ -771,6 +1007,16 @@ mod tests { (code, String::from_utf8(out).unwrap()) } + /// A plain `dump` command with no diagnostics parsing. + fn dump_cmd() -> Command { + Command::Dump { + parse_diagnostics: false, + include_raw: false, + name: None, + tail: None, + } + } + fn set_mtime_to_epoch(path: &Path) { let c_path = CString::new(path.as_os_str().as_encoded_bytes()).unwrap(); let times = [libc::timeval { @@ -810,7 +1056,7 @@ mod tests { assert_eq!(cli.dir, Some(PathBuf::from("/tmp/kvp"))); assert!(cli.unsafe_mode); assert!(cli.json); - assert!(matches!(cli.command, Command::Dump)); + assert!(matches!(cli.command, Command::Dump { .. })); } #[test] @@ -1306,7 +1552,7 @@ mod tests { }, )); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); assert_eq!(dumped, "k=1\nk=2\n"); } @@ -1330,12 +1576,11 @@ mod tests { store.insert("b", "two").unwrap(); store.insert("a", "one").unwrap(); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); assert!(dumped.contains("a=one")); assert!(dumped.contains("b=two")); let (_, entries) = run_dispatch(cli(&dir, Command::Entries)); - // entries are sorted by key assert_eq!(entries, "a=one\nb=two\n"); } @@ -1370,7 +1615,7 @@ mod tests { .unwrap(); assert_eq!(code, EXIT_OK); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); assert_eq!(dumped, "a=1\na=2\nb=3\n"); } @@ -1395,7 +1640,7 @@ mod tests { assert_eq!(code, EXIT_OK); assert_eq!(out, "2\n"); - let (_, dumped) = run_dispatch(cli(&dir, Command::Dump)); + let (_, dumped) = run_dispatch(cli(&dir, dump_cmd())); assert_eq!(dumped, "b=2\n"); } @@ -1409,7 +1654,13 @@ mod tests { let dir = TempDir::new().unwrap(); store_at(&dir).insert("k", "v").unwrap(); - let (code, _) = run_dispatch(cli(&dir, Command::Clear { if_stale })); + let (code, _) = run_dispatch(cli( + &dir, + Command::Clear { + if_stale, + diagnostics: false, + }, + )); assert_eq!(code, EXIT_OK); assert_eq!(store_at(&dir).is_empty().unwrap(), expect_empty_after); } @@ -1602,7 +1853,7 @@ mod tests { store.append("b", "two-prime").unwrap(); store.insert("a", "one").unwrap(); - let (_, out) = run_dispatch(cli_json(&dir, Command::Dump)); + let (_, out) = run_dispatch(cli_json(&dir, dump_cmd())); let json = parse_json(&out); let array = json.as_array().expect("dump --json returns array"); assert_eq!(array.len(), 3); diff --git a/libazureinit-kvp/src/diagnostics.rs b/libazureinit-kvp/src/diagnostics.rs new file mode 100644 index 00000000..bafe87dd --- /dev/null +++ b/libazureinit-kvp/src/diagnostics.rs @@ -0,0 +1,1321 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Typed diagnostics layer over the raw +//! [`KvpPoolStore`](crate::KvpPoolStore) key/value API. +//! +//! Where [`KvpPoolStore`](crate::KvpPoolStore) treats keys and values as +//! opaque bytes, [`DiagnosticsKvp`] understands the telemetry conventions +//! azure-init writes into the guest pool and decodes cloud-init's +//! reporting entries into the same [`DiagnosticEvent`] shape. +//! +//! - **azure-init keys** encode metadata as a seven-segment, +//! pipe-delimited string +//! (`||||||`). +//! The value is the record's message string, stored verbatim for +//! every kind. +//! - **cloud-init keys** +//! (`CLOUD_INIT||||[|]`) store a +//! JSON value; the reader pulls `ts`/`result`/`duration`/`msg` from it +//! and takes everything else from the key, decoding into the same +//! [`DiagnosticEvent`]. This crate only *reads* cloud-init. +//! - **Chunking**: values longer than [`MAX_CHUNK_BYTES`] are split at +//! UTF-8 codepoint boundaries into multiple records under one lock, +//! each keyed with a unique `|` suffix (`0`, `1`, …) +//! since the Hyper-V host keeps only one record per key. Chunks are +//! regrouped on read. +//! - **Classification**: [`records`](DiagnosticsKvp::records) sorts every +//! stored record into a [`DiagnosticRecord`] — a reassembled +//! [`DiagnosticEvent`] (from either agent), an unstructured +//! [`Raw`](DiagnosticRecord::Raw) record such as `PROVISIONING_REPORT`, +//! or a [`Malformed`](DiagnosticRecord::Malformed) event key. +//! +//! This module is policy only: all locking, size enforcement, and +//! on-disk encoding stay in [`KvpPoolStore`](crate::KvpPoolStore). +//! +//! # Example +//! +//! ``` +//! use libazureinit_kvp::{ +//! DiagnosticsKvp, KvpPool, KvpPoolStore, PoolMode, MAX_CHUNK_BYTES, +//! }; +//! +//! # fn main() -> Result<(), libazureinit_kvp::KvpError> { +//! let dir = std::env::temp_dir() +//! .join(format!("libazureinit-kvp-doc-{}", std::process::id())); +//! std::fs::create_dir_all(&dir)?; +//! let store = KvpPoolStore::new_in(KvpPool::Guest, &dir, PoolMode::Safe)?; +//! store.clear()?; +//! +//! let diagnostics = +//! DiagnosticsKvp::new(store, "vm-1234", "azure-init-doc"); +//! +//! // A short event lands in a single record. +//! diagnostics.emit_event("user:create_user", "Creating user azureuser")?; +//! +//! // A long message is split across records and reassembled on read. +//! let long = "x".repeat(MAX_CHUNK_BYTES * 2 + 10); +//! diagnostics.emit_event("config:dump", &long)?; +//! +//! let events = diagnostics.events()?; +//! assert_eq!(events.len(), 2); +//! assert_eq!(events[1].message.len(), MAX_CHUNK_BYTES * 2 + 10); +//! +//! # std::fs::remove_dir_all(&dir).ok(); +//! # Ok(()) +//! # } +//! ``` + +use chrono::Utc; +use uuid::Uuid; + +use crate::{KvpError, KvpPoolStore}; + +/// Literal prefix identifying a cloud-init reporting KVP key. +const CLOUD_INIT_PREFIX: &str = "CLOUD_INIT"; + +/// Maximum number of value bytes per diagnostic KVP record. +/// +/// [`DiagnosticsKvp::emit_event`] splits messages longer than this into +/// multiple records, regardless of the store's +/// [`PoolMode`](crate::PoolMode). It is the conservative +/// [`Safe`](crate::PoolMode::Safe) limit (2 bytes under the Linux kernel +/// `HV_KVP_EXCHANGE_MAX_VALUE` maximum), so diagnostic records stay +/// readable by the Hyper-V host even on an +/// [`Unsafe`](crate::PoolMode::Unsafe) store — its larger capacity is +/// deliberately not used for diagnostics. +pub const MAX_CHUNK_BYTES: usize = 1022; + +/// Delimiter separating the segments of a diagnostic event key. +const EVENT_KEY_DELIMITER: char = '|'; + +/// The kind of a diagnostic record: a span boundary (`start`/`finish`) +/// or a point `event`. The `start`/`finish` tokens match cloud-init's, +/// so a span's boundaries read the same whichever agent emitted them. +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "lowercase")] +pub enum RecordKind { + /// The opening of a span (a function or stage begins), written + /// `start`. + Start, + /// The closing of a span (a function or stage ends), written + /// `finish`. + Finish, + /// A point-in-time event, written `event`. + Event, +} + +impl std::fmt::Display for RecordKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Start => "start", + Self::Finish => "finish", + Self::Event => "event", + }) + } +} + +/// Parses the on-disk `kind`/`type` token; cloud-init emits only +/// `start`/`finish`. +impl std::str::FromStr for RecordKind { + type Err = (); + + fn from_str(token: &str) -> Result { + match token { + "start" => Ok(Self::Start), + "finish" => Ok(Self::Finish), + "event" => Ok(Self::Event), + _ => Err(()), + } + } +} + +/// The current time as an ISO-8601 UTC timestamp (millisecond precision). +fn now_timestamp() -> String { + Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string() +} + +/// Format an azure-init diagnostic event's *shared* key as its +/// `|`-delimited on-disk string: +/// `||||||`. +/// +/// `boot_epoch` is the Unix epoch second the system booted (see +/// [`KvpPoolStore::boot_epoch`](crate::KvpPoolStore::boot_epoch)); it sits +/// in the same slot as cloud-init's incarnation. `kind` records the +/// span/event shape. [`classify_key`] is the inverse. For example: +/// +/// ```text +/// azure-init-0.1.0|1785187982|3f2504e0-...|event|user:create_user|8f3e9c4a-...|2026-07-27T21:33:24.300Z +/// ``` +fn format_event_key( + agent: &str, + boot_epoch: i64, + vm_id: &str, + kind: RecordKind, + name: &str, + event_id: &str, + timestamp: &str, +) -> String { + let d = EVENT_KEY_DELIMITER; + format!( + "{agent}{d}{boot_epoch}{d}{vm_id}{d}{kind}{d}{name}{d}{event_id}\ + {d}{timestamp}" + ) +} +enum KeyClass<'a> { + Event { + agent: &'a str, + boot_epoch: i64, + vm_id: &'a str, + kind: RecordKind, + name: &'a str, + event_id: &'a str, + timestamp: &'a str, + }, + /// The key is a well-formed cloud-init reporting event key + /// (`CLOUD_INIT||||[|]`). + CloudInit { + boot_epoch: i64, + kind: RecordKind, + name: &'a str, + vm_id: Option<&'a str>, + uuid: &'a str, + }, + Malformed { + reason: String, + }, + Raw, +} + +/// Classify a raw pool key. +fn classify_key(key: &str) -> KeyClass<'_> { + if key.split(EVENT_KEY_DELIMITER).next() == Some(CLOUD_INIT_PREFIX) { + return classify_cloud_init_key(key); + } + + let mut segments = key.split(EVENT_KEY_DELIMITER); + let ( + Some(agent), + Some(boot_epoch), + Some(vm_id), + Some(kind), + Some(name), + Some(event_id), + Some(timestamp), + ) = ( + segments.next(), + segments.next(), + segments.next(), + segments.next(), + segments.next(), + segments.next(), + segments.next(), + ) + else { + return KeyClass::Raw; + }; + if segments.next().is_some() { + return KeyClass::Raw; + } + + let Ok(boot_epoch) = boot_epoch.parse::() else { + return KeyClass::Raw; + }; + + match kind.parse::() { + Ok(kind) => KeyClass::Event { + agent, + boot_epoch, + vm_id, + kind, + name, + event_id, + timestamp, + }, + Err(()) => KeyClass::Malformed { + reason: format!("unrecognized kind {kind:?}"), + }, + } +} + +/// Classify a `CLOUD_INIT`-prefixed key into a [`KeyClass::CloudInit`]. +/// +/// Handles both the current layout +/// (`CLOUD_INIT|||||`) and the +/// older one that predates the `vm_id` segment +/// (`CLOUD_INIT||||`). Any other segment +/// count is [`KeyClass::Raw`]; a right-shaped key with a non-numeric +/// incarnation or unrecognized type is [`KeyClass::Malformed`]. +fn classify_cloud_init_key(key: &str) -> KeyClass<'_> { + let mut segments = key.split(EVENT_KEY_DELIMITER); + let _prefix = segments.next(); + let (Some(incarnation), Some(event_type), Some(name), Some(fourth)) = ( + segments.next(), + segments.next(), + segments.next(), + segments.next(), + ) else { + return KeyClass::Raw; + }; + let (vm_id, uuid) = match (segments.next(), segments.next()) { + (None, None) => (None, fourth), + (Some(uuid), None) => (Some(fourth), uuid), + _ => return KeyClass::Raw, + }; + let Ok(boot_epoch) = incarnation.parse::() else { + return KeyClass::Malformed { + reason: format!( + "non-numeric cloud-init incarnation {incarnation:?}" + ), + }; + }; + let Ok(kind) = event_type.parse::() else { + return KeyClass::Malformed { + reason: format!("unrecognized cloud-init type {event_type:?}"), + }; + }; + KeyClass::CloudInit { + boot_epoch, + kind, + name, + vm_id, + uuid, + } +} + +/// Split `value` into pieces of at most `max_bytes` bytes each, always +/// at UTF-8 codepoint boundaries. +/// +/// An empty input yields a single empty chunk so callers still write one +/// record. A codepoint wider than `max_bytes` (only possible for tiny +/// `max_bytes`, never for [`MAX_CHUNK_BYTES`]) is emitted whole so the +/// split always makes progress. +fn chunk_at_char_boundary(value: &str, max_bytes: usize) -> Vec<&str> { + debug_assert!(max_bytes > 0, "max_bytes must be positive"); + if value.is_empty() { + return vec![""]; + } + + let mut chunks = Vec::new(); + let mut start = 0; + while start < value.len() { + if value.len() - start <= max_bytes { + chunks.push(&value[start..]); + break; + } + + let mut end = start + max_bytes; + while end > start && !value.is_char_boundary(end) { + end -= 1; + } + if end == start { + end = start + max_bytes + 1; + while end < value.len() && !value.is_char_boundary(end) { + end += 1; + } + } + + chunks.push(&value[start..end]); + start = end; + } + chunks +} + +/// Reject the `|` key delimiter in an event field so the formatted key +/// round-trips through [`classify_key`]. +fn reject_delimiter(field: &'static str, value: &str) -> Result<(), KvpError> { + if value.contains(EVENT_KEY_DELIMITER) { + return Err(KvpError::EventFieldContainsDelimiter { field }); + } + Ok(()) +} + +/// A single diagnostic event — the decoded, source-agnostic form of one +/// azure-init or cloud-init KVP entry. +/// +/// Metadata (`agent`, `boot_epoch`, `vm_id`, `kind`, `name`, `event_id`) +/// comes from the record key; the payload (`timestamp`, `result`, +/// `duration`, `message`) from the value. Optional fields are populated +/// only when the source provides them (e.g. cloud-init `finish` records +/// carry `result` and `duration`). +#[derive(Clone, Debug, PartialEq, serde::Serialize)] +#[non_exhaustive] +pub struct DiagnosticEvent { + /// Reporting agent identifier from the key, e.g. `azure-init-0.1.0` + /// or `CLOUD_INIT`; also distinguishes the record's source. + pub agent: String, + /// Unix epoch second the system booted (cloud-init's incarnation), + /// shared by every record of one boot. + pub boot_epoch: i64, + /// VM identifier from the key. Absent in cloud-init builds that + /// predate the `vm_id` key segment. + #[serde(skip_serializing_if = "Option::is_none")] + pub vm_id: Option, + /// Whether this record opens a span, closes a span, or is a point + /// event. + pub kind: RecordKind, + /// Formatted event or span name, e.g. `user:create_user`. + pub name: String, + /// Per-record identifier (azure-init's UUIDv4 / cloud-init's uuid); + /// every chunk of one record shares it, as do a span's start and + /// finish. + pub event_id: String, + /// ISO-8601 timestamp, if the source provides one. + #[serde(skip_serializing_if = "Option::is_none")] + pub timestamp: Option, + /// Result string (e.g. `SUCCESS`), present on cloud-init `finish` + /// records. + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + /// Duration in seconds, present on cloud-init `finish` records. + #[serde(skip_serializing_if = "Option::is_none")] + pub duration: Option, + /// Human-readable message. The diagnostics layer imposes no format + /// on this string. + pub message: String, +} + +/// A single record read back from the pool and classified by +/// [`DiagnosticsKvp::records`]. +#[derive(Clone, Debug, PartialEq)] +#[non_exhaustive] +pub enum DiagnosticRecord { + /// A reassembled diagnostic event, from either agent. + Decoded { + /// The decoded event. + event: DiagnosticEvent, + /// Number of on-disk records the value spanned (1 when short). + chunks: usize, + }, + /// An unstructured record whose key is not an event key, such as + /// `PROVISIONING_REPORT`. + Raw { + /// The record key. + key: String, + /// The reassembled record value. + value: String, + }, + /// A record whose key is event-shaped but is not a valid event (for + /// example, an unrecognized kind or invalid cloud-init JSON). + Malformed { + /// The record key. + key: String, + /// The reassembled record value. + value: String, + /// Why the key failed to parse as an event. + reason: String, + }, +} + +/// A typed diagnostics view over a [`KvpPoolStore`]. +/// +/// Owns the `agent` and `vm_id` stamped into this layer's azure-init +/// event keys. See the module-level documentation for the on-disk +/// format. +#[derive(Clone, Debug)] +pub struct DiagnosticsKvp { + store: KvpPoolStore, + vm_id: String, + agent: String, +} + +impl DiagnosticsKvp { + pub fn new( + store: KvpPoolStore, + vm_id: impl Into, + agent: impl Into, + ) -> Self { + Self { + store, + vm_id: vm_id.into(), + agent: agent.into(), + } + } + pub fn store(&self) -> &KvpPoolStore { + &self.store + } + pub fn vm_id(&self) -> &str { + &self.vm_id + } + pub fn agent(&self) -> &str { + &self.agent + } + + /// Emit an azure-init point event with `name` and `message`. + /// + /// The `message` is stored verbatim as the record value. A fresh + /// `event_id` (UUIDv4) is generated. + pub fn emit_event( + &self, + name: impl Into, + message: impl AsRef, + ) -> Result<(), KvpError> { + let event_id = Uuid::new_v4().to_string(); + self.write_event( + RecordKind::Event, + &event_id, + &name.into(), + message.as_ref(), + ) + } + + /// Format the key + /// `||||||` + /// (stamping `boot_epoch`, `vm_id`, `agent`, and the current + /// `timestamp`) and write `value` as its message. + /// + /// The message is written under a single lock via + /// [`KvpPoolStore::append_multiple`]; values longer than + /// [`MAX_CHUNK_BYTES`] are split at UTF-8 codepoint boundaries into + /// multiple records. Every record is keyed with a `|` + /// suffix (`0`, `1`, …) so it is unique, and the chunks are regrouped + /// by [`records`](Self::records) on read. + /// + /// Returns [`KvpError::EventFieldContainsDelimiter`] if `agent`, + /// `vm_id`, `name`, or `event_id` contains the `|` key delimiter. + fn write_event( + &self, + kind: RecordKind, + event_id: &str, + name: &str, + value: &str, + ) -> Result<(), KvpError> { + reject_delimiter("agent", &self.agent)?; + reject_delimiter("vm_id", &self.vm_id)?; + reject_delimiter("name", name)?; + reject_delimiter("event_id", event_id)?; + + let boot_epoch = self.store.boot_epoch()?; + let timestamp = now_timestamp(); + let key = format_event_key( + &self.agent, + boot_epoch, + &self.vm_id, + kind, + name, + event_id, + ×tamp, + ); + + self.write_chunked(&key, value) + } + + /// Split `value` at [`MAX_CHUNK_BYTES`] and append each chunk in one + /// atomic batch, keyed `|` (`0`, `1`, …) so every + /// record is unique. [`reassemble`] strips the index on read. + fn write_chunked(&self, key: &str, value: &str) -> Result<(), KvpError> { + let records: Vec<(String, &str)> = + chunk_at_char_boundary(value, MAX_CHUNK_BYTES) + .into_iter() + .enumerate() + .map(|(subevent_index, chunk)| { + let chunk_key = + format!("{key}{EVENT_KEY_DELIMITER}{subevent_index}"); + (chunk_key, chunk) + }) + .collect(); + self.store.append_multiple(records) + } + + /// Read every record, reassembling chunked events and classifying + /// each into a [`DiagnosticRecord`]. + /// + /// Records are returned in on-disk order. Consecutive records that + /// share an event key — ignoring the `|` suffix — are + /// one event; because [`emit_event`](Self::emit_event) writes an + /// event's chunks contiguously under a single lock, reassembly is + /// correct even under concurrent writers. + pub fn records(&self) -> Result, KvpError> { + Ok(reassemble(self.store.dump()?)) + } + + /// Read back every decoded [`DiagnosticEvent`], from either agent, in + /// on-disk order. + /// + /// Raw and malformed records are excluded. Use + /// [`records`](Self::records) for the full view that includes them. + pub fn events(&self) -> Result, KvpError> { + Ok(self + .records()? + .into_iter() + .filter_map(|record| match record { + DiagnosticRecord::Decoded { event, .. } => Some(event), + DiagnosticRecord::Raw { .. } + | DiagnosticRecord::Malformed { .. } => None, + }) + .collect()) + } + + /// Remove every diagnostic key: any key that parses as an event key + /// (including every `|` chunk of a multi-record event) + /// or a malformed event key. Raw records such as `PROVISIONING_REPORT` + /// are left intact. + pub fn clear(&self) -> Result<(), KvpError> { + let keys: Vec = self + .store + .dump()? + .into_iter() + .filter_map(|(key, _)| { + let is_diagnostic = !matches!( + classify_key(base_event_key(&key)), + KeyClass::Raw + ); + is_diagnostic.then_some(key) + }) + .collect(); + self.store.delete_multiple(keys)?; + Ok(()) + } +} + +/// The shared event key a chunk belongs to: strips a trailing +/// `|`, or returns the key unchanged if it has none. +fn base_event_key(key: &str) -> &str { + split_subevent_index(key).0 +} + +/// Split a key into its base event key and optional trailing subevent +/// index: `(base, Some(index))` when a trailing numeric segment follows an +/// event-shaped base (valid or malformed), else `(key, None)`. +/// [`reassemble`] uses the index to regroup an event's chunks and restore +/// their write order. +fn split_subevent_index(key: &str) -> (&str, Option) { + if let Some((base, index)) = key.rsplit_once(EVENT_KEY_DELIMITER) { + if let Ok(index) = index.parse::() { + if matches!( + classify_key(base), + KeyClass::Event { .. } + | KeyClass::CloudInit { .. } + | KeyClass::Malformed { .. } + ) { + return (base, Some(index)); + } + } + } + (key, None) +} + +/// Group consecutive records sharing an event key — chunk +/// `|` suffixes stripped — from [`KvpPoolStore::dump`] +/// and classify each group into a [`DiagnosticRecord`]. +fn reassemble(dumped: Vec<(String, String)>) -> Vec { + let mut parsed = dumped + .into_iter() + .map(|(key, value)| { + let (base, index) = split_subevent_index(&key); + (base.to_string(), index, value) + }) + .peekable(); + + let mut records = Vec::new(); + while let Some((base, index, value)) = parsed.next() { + let mut indexed = vec![(index, value)]; + while parsed.peek().is_some_and(|(next, _, _)| *next == base) { + let (_, next_index, next_value) = + parsed.next().expect("peeked value exists"); + indexed.push((next_index, next_value)); + } + // Restore write order by subevent index. Stable, so a single + // record (index `None`) or any equal indices keep on-disk order. + indexed.sort_by_key(|(index, _)| *index); + let chunk_values = + indexed.into_iter().map(|(_, value)| value).collect(); + records.push(classify_record(base, chunk_values)); + } + + records +} + +/// Classify one reassembled group of chunk values (ordered by subevent +/// index, never empty) into a [`DiagnosticRecord`]. azure-init and raw +/// records concatenate their values; cloud-init chunks are stitched and +/// decoded via [`decode_cloud_init_value`]. +fn classify_record(key: String, chunk_values: Vec) -> DiagnosticRecord { + let chunks = chunk_values.len(); + match classify_key(&key) { + KeyClass::Event { + agent, + boot_epoch, + vm_id, + kind, + name, + event_id, + timestamp, + } => { + let message = chunk_values.concat(); + DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: agent.to_string(), + boot_epoch, + vm_id: Some(vm_id.to_string()), + kind, + name: name.to_string(), + event_id: event_id.to_string(), + timestamp: Some(timestamp.to_string()), + result: None, + duration: None, + message, + }, + chunks, + } + } + KeyClass::CloudInit { + boot_epoch, + kind, + name, + vm_id, + uuid, + } => { + // Own the key-derived fields up front so `key` and + // `chunk_values` can move into a `Malformed` record when a + // chunk's value fails to decode. + let name = name.to_string(); + let vm_id = vm_id.map(str::to_string); + let uuid = uuid.to_string(); + match decode_cloud_init_value(&chunk_values) { + Ok((meta, message)) => DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: CLOUD_INIT_PREFIX.to_string(), + boot_epoch, + vm_id, + kind, + name, + event_id: uuid, + timestamp: meta + .get("ts") + .and_then(|t| t.as_str()) + .map(str::to_string), + result: meta + .get("result") + .and_then(|r| r.as_str()) + .map(str::to_string), + duration: meta.get("duration").and_then(|d| d.as_f64()), + message, + }, + chunks, + }, + Err(err) => DiagnosticRecord::Malformed { + key, + value: chunk_values.concat(), + reason: format!("invalid cloud-init JSON value: {err}"), + }, + } + } + KeyClass::Malformed { reason } => DiagnosticRecord::Malformed { + key, + value: chunk_values.concat(), + reason, + }, + KeyClass::Raw => DiagnosticRecord::Raw { + key, + value: chunk_values.concat(), + }, + } +} + +/// Marker preceding a cloud-init value's message field: `"msg":"`. +const CLOUD_INIT_MSG_MARKER: &str = "\"msg\":\""; + +/// Decode a cloud-init event's chunk value(s) into `(metadata, message)`, +/// reading `ts`/`result`/`duration` from an untyped [`serde_json::Value`]. +/// +/// A single record is complete JSON, parsed directly. A multi-record event +/// was split mid-escape by cloud-init's `_break_down` (e.g. a `\n` cut into +/// `\` and `n`), so no chunk is valid JSON alone: recover each chunk's raw +/// escaped `msg` slice, concatenate, and unescape once; metadata comes from +/// the first chunk. +fn decode_cloud_init_value( + chunks: &[String], +) -> Result<(serde_json::Value, String), String> { + if let [only] = chunks { + let value: serde_json::Value = + serde_json::from_str(only).map_err(|e| e.to_string())?; + let message = value + .get("msg") + .and_then(|m| m.as_str()) + .unwrap_or_default() + .to_string(); + return Ok((value, message)); + } + + let mut escaped = String::new(); + for chunk in chunks { + escaped.push_str(cloud_init_escaped_msg_slice(chunk)?); + } + let message: String = serde_json::from_str(&format!("\"{escaped}\"")) + .map_err(|e| e.to_string())?; + + Ok((cloud_init_chunk_metadata(&chunks[0])?, message)) +} + +/// Recover a chunk's raw (still-escaped) `msg` slice — the bytes between +/// the `"msg":"` marker and the closing `"}` — without unescaping. +fn cloud_init_escaped_msg_slice(chunk: &str) -> Result<&str, String> { + let start = chunk + .find(CLOUD_INIT_MSG_MARKER) + .ok_or("chunk is missing a \"msg\" field")? + + CLOUD_INIT_MSG_MARKER.len(); + let end = chunk + .strip_suffix("\"}") + .map(str::len) + .ok_or("chunk does not end with '\"}'")?; + chunk + .get(start..end) + .ok_or_else(|| "chunk \"msg\" field is malformed".to_string()) +} + +/// Parse a chunk's non-`msg` prefix (the portion before its `,"msg":"` +/// field, which is always valid JSON) into a [`serde_json::Value`]. +fn cloud_init_chunk_metadata(chunk: &str) -> Result { + let marker = format!(",{CLOUD_INIT_MSG_MARKER}"); + let end = chunk + .find(&marker) + .ok_or("chunk is missing a \"msg\" field")?; + serde_json::from_str(&format!("{}}}", &chunk[..end])) + .map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{KvpPool, PoolMode}; + use rstest::rstest; + + const AGENT: &str = "azure-init-0.1.0"; + const VM_ID: &str = "3f2504e0-4f89-41d3-9a0c-0305e82c3301"; + const EVENT_ID: &str = "8f3e9c4a-1b2c-4d5e-9f01-234567890abc"; + const BOOT_EPOCH: i64 = 1_700_000_000; + const TIMESTAMP: &str = "2026-07-27T21:33:24.300Z"; + + #[test] + fn event_key_formats_and_classifies() { + let formatted = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Event, + "user:create_user", + EVENT_ID, + TIMESTAMP, + ); + assert_eq!( + formatted, + format!( + "{AGENT}|{BOOT_EPOCH}|{VM_ID}|event|user:create_user|\ + {EVENT_ID}|{TIMESTAMP}" + ) + ); + assert!(matches!( + classify_key(&formatted), + KeyClass::Event { + agent, + boot_epoch, + vm_id, + kind, + name, + event_id, + timestamp, + } if agent == AGENT + && boot_epoch == BOOT_EPOCH + && vm_id == VM_ID + && kind == RecordKind::Event + && name == "user:create_user" + && event_id == EVENT_ID + && timestamp == TIMESTAMP + )); + } + + #[test] + fn classify_round_trips_every_kind() { + for expected in + [RecordKind::Start, RecordKind::Finish, RecordKind::Event] + { + let key = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + expected, + "span:event", + EVENT_ID, + TIMESTAMP, + ); + assert!(matches!( + classify_key(&key), + KeyClass::Event { kind, .. } if kind == expected + )); + } + } + + fn class_of(key: &str) -> &'static str { + match classify_key(key) { + KeyClass::Event { .. } => "event", + KeyClass::CloudInit { .. } => "cloud-init", + KeyClass::Malformed { .. } => "malformed", + KeyClass::Raw => "raw", + } + } + + #[rstest] + #[case::event("a|100|vm|event|name|id|ts", "event")] + #[case::cloud_init( + "CLOUD_INIT|1785187982|finish|name|vmid|uuid", + "cloud-init" + )] + #[case::raw_single_segment("PROVISIONING_REPORT", "raw")] + #[case::raw_too_few_segments("a|100|vm|event|name|id", "raw")] + #[case::raw_too_many_segments("a|100|vm|event|name|id|ts|extra", "raw")] + #[case::raw_non_numeric_boot_epoch("a|notnum|vm|event|name|id|ts", "raw")] + #[case::malformed_bad_kind("a|100|vm|NOTAKIND|name|id|ts", "malformed")] + #[case::malformed_other_kind("a|100|vm|nope|name|id|ts", "malformed")] + #[case::cloud_init_bad_type( + "CLOUD_INIT|100|weird|name|vmid|uuid", + "malformed" + )] + #[case::cloud_init_non_numeric_incarnation( + "CLOUD_INIT|notnum|finish|name|vmid|uuid", + "malformed" + )] + fn classify_key_categorizes(#[case] key: &str, #[case] expected: &str) { + assert_eq!(class_of(key), expected); + } + + #[rstest] + #[case::empty("", 4, vec![""])] + #[case::shorter_than_max("abc", 8, vec!["abc"])] + #[case::exact_multiple("abcdef", 2, vec!["ab", "cd", "ef"])] + #[case::ascii_remainder("abcde", 2, vec!["ab", "cd", "e"])] + #[case::two_byte_boundary("aéb", 2, vec!["a", "é", "b"])] + #[case::oversized_three_byte("€", 1, vec!["€"])] + #[case::oversized_repeated("€€", 1, vec!["€", "€"])] + fn chunk_splits_at_utf8_boundaries( + #[case] input: &str, + #[case] max_bytes: usize, + #[case] expected: Vec<&str>, + ) { + assert_eq!(chunk_at_char_boundary(input, max_bytes), expected); + } + + #[test] + fn chunk_reassembles_multibyte_payload() { + let payload = "🚀".repeat(100); + let chunks = chunk_at_char_boundary(&payload, 7); + assert!(chunks.iter().all(|chunk| chunk.len() <= 7)); + assert_eq!(chunks.concat(), payload); + } + + #[test] + fn reject_delimiter_flags_pipe() { + assert!(reject_delimiter("name", "no pipe here").is_ok()); + let err = reject_delimiter("name", "has|pipe").unwrap_err(); + assert!(matches!( + err, + KvpError::EventFieldContainsDelimiter { field: "name" } + )); + assert_eq!( + err.to_string(), + "event key field 'name' must not contain '|'" + ); + } + + #[test] + fn reassemble_groups_chunks_and_classifies() { + let key = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Start, + "config:dump", + EVENT_ID, + TIMESTAMP, + ); + + let dumped = vec![ + (key.clone(), "part-one/".to_string()), + (key.clone(), "part-two".to_string()), + ( + "PROVISIONING_REPORT".to_string(), + "result=success".to_string(), + ), + ("a|100|vm|NOPE|name|id|ts".to_string(), "junk".to_string()), + ]; + + let records = reassemble(dumped); + assert_eq!(records.len(), 3); + + assert_eq!( + records[0], + DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: AGENT.to_string(), + boot_epoch: BOOT_EPOCH, + vm_id: Some(VM_ID.to_string()), + kind: RecordKind::Start, + name: "config:dump".to_string(), + event_id: EVENT_ID.to_string(), + timestamp: Some(TIMESTAMP.to_string()), + result: None, + duration: None, + message: "part-one/part-two".to_string(), + }, + chunks: 2, + } + ); + assert!(matches!(&records[1], DiagnosticRecord::Raw { key, .. } + if key == "PROVISIONING_REPORT")); + assert!(matches!(&records[2], DiagnosticRecord::Malformed { .. })); + } + + #[test] + fn reassemble_keeps_distinct_adjacent_keys_separate() { + let make = |event_id: &str| { + format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Start, + "span:name", + event_id, + TIMESTAMP, + ) + }; + let dumped = vec![ + (make("id-1"), "first".to_string()), + (make("id-2"), "second".to_string()), + ]; + let records = reassemble(dumped); + assert_eq!(records.len(), 2); + assert!(matches!( + &records[0], + DiagnosticRecord::Decoded { chunks: 1, .. } + )); + assert!(matches!( + &records[1], + DiagnosticRecord::Decoded { chunks: 1, .. } + )); + } + + #[rstest] + #[case::indexed_chunk( + "a|100|vm|event|name|id|ts|0", + "a|100|vm|event|name|id|ts" + )] + #[case::indexed_chunk_multi_digit( + "a|100|vm|event|name|id|ts|12", + "a|100|vm|event|name|id|ts" + )] + #[case::single_event_unchanged( + "a|100|vm|event|name|id|ts", + "a|100|vm|event|name|id|ts" + )] + #[case::cloud_init_indexed_chunk( + "CLOUD_INIT|1785187982|finish|mod|vmid|uuid|0", + "CLOUD_INIT|1785187982|finish|mod|vmid|uuid" + )] + #[case::raw_unchanged("PROVISIONING_REPORT", "PROVISIONING_REPORT")] + #[case::non_event_numeric_tail_unchanged("foo|3", "foo|3")] + #[case::malformed_unchanged( + "a|100|vm|NOPE|name|id|ts", + "a|100|vm|NOPE|name|id|ts" + )] + #[case::malformed_indexed_chunk( + "a|100|vm|NOPE|name|id|ts|0", + "a|100|vm|NOPE|name|id|ts" + )] + fn base_event_key_strips_event_subevent_index( + #[case] key: &str, + #[case] expected: &str, + ) { + assert_eq!(base_event_key(key), expected); + } + + #[test] + fn reassemble_groups_indexed_chunk_keys() { + let base = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Finish, + "config:dump", + EVENT_ID, + TIMESTAMP, + ); + let dumped = vec![ + (format!("{base}|0"), "part-one/".to_string()), + (format!("{base}|1"), "part-two/".to_string()), + (format!("{base}|2"), "part-three".to_string()), + ]; + + let records = reassemble(dumped); + assert_eq!(records.len(), 1); + assert_eq!( + records[0], + DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: AGENT.to_string(), + boot_epoch: BOOT_EPOCH, + vm_id: Some(VM_ID.to_string()), + kind: RecordKind::Finish, + name: "config:dump".to_string(), + event_id: EVENT_ID.to_string(), + timestamp: Some(TIMESTAMP.to_string()), + result: None, + duration: None, + message: "part-one/part-two/part-three".to_string(), + }, + chunks: 3, + } + ); + } + + #[test] + fn azure_event_value_is_full_message() { + let key = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Event, + "user:create_user", + EVENT_ID, + TIMESTAMP, + ); + assert!(matches!( + classify_record(key, vec!["boom".to_string()]), + DiagnosticRecord::Decoded { event, chunks: 1 } + if event.kind == RecordKind::Event + && event.message == "boom" + && event.agent == AGENT + && event.vm_id.as_deref() == Some(VM_ID) + && event.timestamp.as_deref() == Some(TIMESTAMP) + )); + } + + #[test] + fn azure_span_value_is_full_message() { + let key = format_event_key( + AGENT, + BOOT_EPOCH, + VM_ID, + RecordKind::Finish, + "config:write", + EVENT_ID, + TIMESTAMP, + ); + assert!(matches!( + classify_record(key, vec!["write_config completed".to_string()]), + DiagnosticRecord::Decoded { event, chunks: 1 } + if event.kind == RecordKind::Finish + && event.message == "write_config completed" + )); + } + + #[test] + fn azure_span_start_finish_pair_round_trips_through_writer() { + let dir = tempfile::TempDir::new().unwrap(); + let store = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap(); + let diag = DiagnosticsKvp::new(store, VM_ID, AGENT); + + // A span emits a start and a finish sharing one event_id. + diag.write_event( + RecordKind::Start, + EVENT_ID, + "provision:run", + "starting provision", + ) + .unwrap(); + diag.write_event( + RecordKind::Finish, + EVENT_ID, + "provision:run", + "provision completed", + ) + .unwrap(); + + let events = diag.events().unwrap(); + assert_eq!(events.len(), 2); + + assert_eq!(events[0].kind, RecordKind::Start); + assert_eq!(events[0].message, "starting provision"); + assert_eq!(events[1].kind, RecordKind::Finish); + assert_eq!(events[1].message, "provision completed"); + + for event in &events { + assert_eq!(event.event_id, EVENT_ID); + assert_eq!(event.name, "provision:run"); + assert_eq!(event.agent, AGENT); + assert_eq!(event.vm_id.as_deref(), Some(VM_ID)); + } + } + + const CLOUD_INIT_VM_ID: &str = "0e5e179d-5341-478b-8456-fbb90621bdf8"; + const CLOUD_INIT_KEY_FINISH: &str = "CLOUD_INIT|1785187982|finish|modules-final/config-scripts_user|0e5e179d-5341-478b-8456-fbb90621bdf8|e5f01809-a7a3-4279-aa64-1f18e21eda6e"; + const CLOUD_INIT_VALUE_FINISH: &str = r#"{"name":"modules-final/config-scripts_user","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.0006448590000012189,"msg":"config-scripts_user ran successfully and took 0.001 seconds"}"#; + + #[test] + fn cloud_init_key_classifies() { + assert!(matches!( + classify_key(CLOUD_INIT_KEY_FINISH), + KeyClass::CloudInit { boot_epoch, kind, name, vm_id, uuid } + if boot_epoch == 1785187982 + && kind == RecordKind::Finish + && name == "modules-final/config-scripts_user" + && vm_id == Some(CLOUD_INIT_VM_ID) + && uuid == "e5f01809-a7a3-4279-aa64-1f18e21eda6e" + )); + assert!(matches!( + classify_key( + "CLOUD_INIT|1785187982|start|modules-config/foo|\ + c4d4a08d-fe93-4c7a-9be6-9a38c212e212" + ), + KeyClass::CloudInit { boot_epoch, kind, name, vm_id, uuid } + if boot_epoch == 1785187982 + && kind == RecordKind::Start + && name == "modules-config/foo" + && vm_id.is_none() + && uuid == "c4d4a08d-fe93-4c7a-9be6-9a38c212e212" + )); + } + + #[rstest] + #[case::too_many("CLOUD_INIT|a|b|c|d|e|f", "raw")] + #[case::too_few("CLOUD_INIT|a|b|c", "raw")] + #[case::prefix_only("CLOUD_INIT", "raw")] + fn cloud_init_bad_shapes_are_raw( + #[case] key: &str, + #[case] expected: &str, + ) { + assert_eq!(class_of(key), expected); + } + + #[test] + fn cloud_init_finish_record_decodes_all_fields() { + assert!(matches!( + classify_record( + CLOUD_INIT_KEY_FINISH.to_string(), + vec![CLOUD_INIT_VALUE_FINISH.to_string()], + ), + DiagnosticRecord::Decoded { event, chunks: 1 } + if event.agent == "CLOUD_INIT" + && event.boot_epoch == 1785187982 + && event.kind == RecordKind::Finish + && event.name == "modules-final/config-scripts_user" + && event.vm_id.as_deref() == Some(CLOUD_INIT_VM_ID) + && event.event_id == "e5f01809-a7a3-4279-aa64-1f18e21eda6e" + && event.timestamp.as_deref() + == Some("2026-07-27T21:33:24.339006+00:00") + && event.result.as_deref() == Some("SUCCESS") + && event.duration.is_some_and(|d| { + (d - 0.000_644_859_000_001_218_9).abs() < 1e-12 + }) + && event.message + == "config-scripts_user ran successfully and took \ + 0.001 seconds" + )); + } + + #[test] + fn cloud_init_start_record_has_no_result_or_duration() { + let value = r#"{"name":"modules-final/config-keys_to_console","type":"start","ts":"2026-07-27T21:33:24.344349+00:00","msg":"running config-keys_to_console with frequency once-per-instance"}"#; + let key = "CLOUD_INIT|1785187982|start|modules-final/config-keys_to_console|0e5e179d-5341-478b-8456-fbb90621bdf8|7792621b-b339-4274-8b71-2a3dcbd2db4e"; + assert_eq!( + classify_record(key.to_string(), vec![value.to_string()]), + DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: "CLOUD_INIT".to_string(), + boot_epoch: 1785187982, + vm_id: Some(CLOUD_INIT_VM_ID.to_string()), + kind: RecordKind::Start, + name: "modules-final/config-keys_to_console".to_string(), + event_id: "7792621b-b339-4274-8b71-2a3dcbd2db4e" + .to_string(), + timestamp: Some( + "2026-07-27T21:33:24.344349+00:00".to_string() + ), + result: None, + duration: None, + message: "running config-keys_to_console with frequency \ + once-per-instance" + .to_string(), + }, + chunks: 1, + } + ); + } + + #[test] + fn cloud_init_key_with_invalid_json_is_malformed() { + assert!(matches!( + classify_record( + CLOUD_INIT_KEY_FINISH.to_string(), + vec!["not json".to_string()], + ), + DiagnosticRecord::Malformed { reason, .. } + if reason.contains("cloud-init") + )); + } + + #[test] + fn cloud_init_chunks_reassemble_by_subevent_index() { + let base = "CLOUD_INIT|1785187982|finish|modules-final/long|0e5e179d-5341-478b-8456-fbb90621bdf8|abc12345-1111-2222-3333-444455556666"; + let chunk = |i: u32, msg: &str| { + format!( + r#"{{"name":"modules-final/long","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.5,"msg_i":{i},"msg":"{msg}"}}"# + ) + }; + let dumped = vec![ + (format!("{base}|1"), chunk(1, "two ")), + (format!("{base}|0"), chunk(0, "one ")), + (format!("{base}|2"), chunk(2, "three")), + ]; + + let records = reassemble(dumped); + assert_eq!( + records, + vec![DiagnosticRecord::Decoded { + event: DiagnosticEvent { + agent: "CLOUD_INIT".to_string(), + boot_epoch: 1785187982, + vm_id: Some(CLOUD_INIT_VM_ID.to_string()), + kind: RecordKind::Finish, + name: "modules-final/long".to_string(), + event_id: "abc12345-1111-2222-3333-444455556666" + .to_string(), + timestamp: Some( + "2026-07-27T21:33:24.339006+00:00".to_string() + ), + result: Some("SUCCESS".to_string()), + duration: Some(0.5), + message: "one two three".to_string(), + }, + chunks: 3, + }] + ); + } + + #[test] + fn cloud_init_chunks_reassemble_split_json_escape() { + let base = "CLOUD_INIT|1785187982|finish|modules-final/x|0e5e179d-5341-478b-8456-fbb90621bdf8|abc12345-1111-2222-3333-444455556666"; + let dumped = vec![ + ( + format!("{base}|0"), + r#"{"name":"modules-final/x","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.5,"msg_i":0,"msg":"line1\"}"# + .to_string(), + ), + ( + format!("{base}|1"), + r#"{"name":"modules-final/x","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.5,"msg_i":1,"msg":"nline2"}"# + .to_string(), + ), + ]; + + let records = reassemble(dumped); + assert!(matches!( + &records[..], + [DiagnosticRecord::Decoded { event, chunks: 2 }] + if event.message == "line1\nline2" + && event.result.as_deref() == Some("SUCCESS") + && event.duration == Some(0.5) + )); + } +} diff --git a/libazureinit-kvp/src/error.rs b/libazureinit-kvp/src/error.rs index 02b9ec9c..0fb171d0 100644 --- a/libazureinit-kvp/src/error.rs +++ b/libazureinit-kvp/src/error.rs @@ -11,6 +11,10 @@ pub enum KvpError { EmptyKey, /// An underlying I/O error. Io(io::Error), + /// An event key field (`agent`, `vm_id`, `name`, or `event_id`) + /// contained the `|` delimiter, which would make the formatted event + /// key ambiguous to parse back. + EventFieldContainsDelimiter { field: &'static str }, /// The key contains a null byte, which is incompatible with the /// on-disk format (null-padded fixed-width fields). KeyContainsNull, @@ -29,6 +33,9 @@ impl fmt::Display for KvpError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::EmptyKey => write!(f, "KVP key must not be empty"), + Self::EventFieldContainsDelimiter { field } => { + write!(f, "event key field '{field}' must not contain '|'") + } Self::Io(e) => write!(f, "{e}"), Self::KeyContainsNull => { write!(f, "KVP key must not contain null bytes") diff --git a/libazureinit-kvp/src/lib.rs b/libazureinit-kvp/src/lib.rs index f17b66f7..197093f0 100644 --- a/libazureinit-kvp/src/lib.rs +++ b/libazureinit-kvp/src/lib.rs @@ -9,14 +9,21 @@ //! - [`ProvisioningReport`]: structured provisioning health report that //! is persisted as the single `PROVISIONING_REPORT` record with //! [`write_report`]. +//! - [`DiagnosticsKvp`]: typed view over [`KvpPoolStore`] that formats, +//! chunks, and reassembles azure-init diagnostic events. mod cli; +mod diagnostics; mod error; mod report; mod store; mod vm_id; pub use cli::run; +pub use diagnostics::{ + DiagnosticEvent, DiagnosticRecord, DiagnosticsKvp, RecordKind, + MAX_CHUNK_BYTES, +}; pub use error::KvpError; pub use report::{ write_report, ProvisioningReport, ReportPpsType, PROVISIONING_REPORT_KEY, diff --git a/libazureinit-kvp/src/report.rs b/libazureinit-kvp/src/report.rs index 98274588..43b3a829 100644 --- a/libazureinit-kvp/src/report.rs +++ b/libazureinit-kvp/src/report.rs @@ -55,9 +55,7 @@ impl std::fmt::Display for ReportResult { /// Pre-provisioning (PPS) type reported in the `pps_type` field. /// /// Mirrors the values cloud-init reports for the platform's -/// `PreprovisionedVMType` / IMDS `ppsType`. Marked `#[non_exhaustive]` -/// so new platform PPS types can be added without breaking downstream -/// `match` statements. +/// `PreprovisionedVMType` / IMDS `ppsType`. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ReportPpsType { /// Not pre-provisioned (`None`). diff --git a/libazureinit-kvp/src/store.rs b/libazureinit-kvp/src/store.rs index a067c980..64152a28 100644 --- a/libazureinit-kvp/src/store.rs +++ b/libazureinit-kvp/src/store.rs @@ -133,7 +133,7 @@ impl KvpPoolStore { /// checking for an existing key. /// /// This preserves any existing records, including duplicate keys, - /// and does not enforce [`MAX_UNIQUE_KEYS`]. Use + /// and does not enforce the `MAX_UNIQUE_KEYS` cap. Use /// [`insert`](Self::insert) when callers need upsert semantics. pub fn append(&self, key: &str, value: &str) -> Result<(), KvpError> { validate_key(key, self.mode.max_key_size())?; @@ -153,8 +153,8 @@ impl KvpPoolStore { /// interleaving with the batch. /// /// Existing records are kept and duplicate keys are preserved. - /// Like [`append`](Self::append), this does not enforce - /// [`MAX_UNIQUE_KEYS`]; use [`load`](Self::load) when the + /// Like [`append`](Self::append), this does not enforce the + /// `MAX_UNIQUE_KEYS` cap; use [`load`](Self::load) when the /// caller is replacing the entire pool and wants the unique-key cap /// enforced. Validation happens before the file is opened; empty /// input is a no-op and does not create the pool file. If an I/O @@ -211,7 +211,7 @@ impl KvpPoolStore { let boot_time = boot_time(&*self.ops)?; lock_for_writing(&mut *handle)?; - if handle.metadata()?.mtime <= boot_time { + if handle.metadata()?.mtime < boot_time { handle.set_len(0)?; } } @@ -399,7 +399,18 @@ impl KvpPoolStore { Err(ref e) if e.kind() == ErrorKind::NotFound => return Ok(false), Err(e) => return Err(e.into()), }; - Ok(metadata.mtime <= boot_time(&*self.ops)?) + Ok(metadata.mtime < boot_time(&*self.ops)?) + } + + /// The system boot time as a Unix epoch timestamp in seconds, read + /// from `/proc/stat` `btime`. + /// + /// Diagnostic event keys stamp this value so records can be + /// attributed to a specific boot (mirroring cloud-init's + /// incarnation), letting readers tell this boot's telemetry from a + /// previous boot's. + pub fn boot_epoch(&self) -> Result { + boot_time(&*self.ops) } /// Variant of [`is_stale`](Self::is_stale) that takes an explicit @@ -412,7 +423,7 @@ impl KvpPoolStore { Err(ref e) if e.kind() == ErrorKind::NotFound => return Ok(false), Err(e) => return Err(e.into()), }; - Ok(metadata.mtime <= boot_time) + Ok(metadata.mtime < boot_time) } fn iter(&self) -> Result { @@ -522,8 +533,6 @@ impl KvpPoolStore { iter.flush()?; Ok(()) } - - /// Return a reference to the pool file path. pub fn path(&self) -> &Path { &self.path } @@ -538,7 +547,7 @@ impl KvpPoolStore { /// /// This is the inverse of [`dump`](Self::dump): duplicate keys are /// preserved exactly as provided, but the number of unique keys is - /// capped at [`MAX_UNIQUE_KEYS`]. Existing records are discarded; + /// capped at `MAX_UNIQUE_KEYS`. Existing records are discarded; /// empty input clears the pool. Use /// [`append_multiple`](Self::append_multiple) when callers need to /// extend the pool instead. @@ -1359,7 +1368,6 @@ mod tests { assert_eq!(store.read(&long_key).unwrap(), Some("val".to_string())); - // Read is not size-capped: an oversized key simply misses. let too_long = "k".repeat(513); assert_eq!(store.read(&too_long).unwrap(), None); } @@ -1660,7 +1668,6 @@ mod tests { assert!(store.delete("k4").unwrap()); - // k9 takes k4's slot; the rest stays put. assert_eq!( store.dump().unwrap(), pairs([ @@ -1754,7 +1761,6 @@ mod tests { let store = safe_store(dir.path()); store.load(pairs([("keep", "me")])).unwrap(); - // Bad record mid-batch: file must be untouched on rejection. let bad_value = "v".repeat(1023); let err = store .load(vec![ @@ -1779,7 +1785,6 @@ mod tests { let err = store.load(too_many).unwrap_err(); assert!(is_max_keys(&err), "got {err:?}"); - // Cap is checked pre-lock; the file is never opened. assert!(!store.path().exists()); } @@ -1788,7 +1793,6 @@ mod tests { let dir = TempDir::new().unwrap(); let store = safe_store(dir.path()); - // 2 * MAX_UNIQUE_KEYS records, MAX_UNIQUE_KEYS unique keys. let mut records: Vec<(String, String)> = (0..MAX_UNIQUE_KEYS) .map(|i| (format!("k{i}"), "a".to_string())) .collect(); @@ -1813,7 +1817,6 @@ mod tests { let err = store.insert("overflow", "v").unwrap_err(); assert!(is_max_keys(&err), "got {err:?}"); - // Overwriting an existing key at the cap still works. store.insert("k0", "updated").unwrap(); assert_eq!(store.read("k0").unwrap(), Some("updated".to_string())); } @@ -1849,14 +1852,11 @@ mod tests { let dir = TempDir::new().unwrap(); let store = safe_store(dir.path()); - // Mirrors the chunked-event use case: many records sharing a - // single key, written atomically. let records = pairs([("chunk", "part1"), ("chunk", "part2"), ("chunk", "part3")]); store.append_multiple(records.clone()).unwrap(); assert_eq!(store.dump().unwrap(), records); - // `read` returns last-write-wins; entries() collapses to 1. assert_eq!(store.entries().unwrap().len(), 1); } @@ -1869,7 +1869,6 @@ mod tests { .append_multiple(Vec::<(String, String)>::new()) .unwrap(); - // No file created when the input is empty. assert!(!store.path().exists()); } @@ -1889,8 +1888,6 @@ mod tests { .unwrap_err(); assert!(is_value_too_large(&err), "got {err:?}"); - // Rejection is all-or-nothing: previously-written records - // are untouched, none of the new batch lands. assert_eq!(store.dump().unwrap(), pairs([("keep", "me")])); } @@ -1919,16 +1916,11 @@ mod tests { #[test] fn test_append_multiple_does_not_enforce_unique_key_cap() { - // Matches `append`'s contract: the bulk variant deliberately - // skips the unique-key cap so chunked writes (many records - // sharing one key) cannot accidentally trip it. let dir = TempDir::new().unwrap(); let store = safe_store(dir.path()); seed_unique_keys(&store, MAX_UNIQUE_KEYS); - // Adding a new unique key via append_multiple is allowed even - // when the pool is already at the cap. store.append_multiple(pairs([("extra", "v")])).unwrap(); assert_eq!(store.len().unwrap(), MAX_UNIQUE_KEYS + 1); } @@ -1958,7 +1950,6 @@ mod tests { let dir = TempDir::new().unwrap(); let store = safe_store(dir.path()); - // Two unique keys, three matching records. store .load(pairs([ ("k", "v1"), @@ -2007,7 +1998,6 @@ mod tests { let removed = store.delete_multiple(Vec::::new()).unwrap(); assert_eq!(removed, 0); - // Empty input never opens or creates the file. assert!(!store.path().exists()); } @@ -2028,8 +2018,6 @@ mod tests { store.load(pairs([("a", "1"), ("b", "2")])).unwrap(); - // Listing the same key twice still removes the (one) record - // exactly once. let removed = store.delete_multiple(vec!["a", "a", "a"]).unwrap(); assert_eq!(removed, 1); assert_eq!(store.dump().unwrap(), pairs([("b", "2")])); @@ -2046,7 +2034,6 @@ mod tests { .unwrap_err(); assert!(matches!(err, KvpError::EmptyKey), "got {err:?}"); - // Validation runs before any record is removed. assert_eq!(store.dump().unwrap(), pairs([("a", "1")])); } @@ -2063,8 +2050,6 @@ mod tests { #[test] fn test_delete_multiple_size_independent_of_mode() { - // Mirrors `delete`'s contract: keys longer than the safe-mode - // cap can be removed from a safe-mode store. let dir = TempDir::new().unwrap(); let store_unsafe = unsafe_store(dir.path()); let long_key = "k".repeat(SAFE_MAX_KEY_BYTES + 1); @@ -3063,21 +3048,14 @@ mod tests { let store = safe_store(dir.path()); store.load(pairs([("a", "1"), ("b", "2")])).unwrap(); - // Open a mutable iterator (exclusive lock) so we can manipulate the file. let mut iter = store.iter_mut().unwrap(); assert_eq!(iter.record_count(), 2); - // Read the first record successfully. let (k, _) = iter.next().unwrap().unwrap(); assert_eq!(k, "a"); - // Truncate via the iterator's own handle to remove the second record. - // The iterator still thinks record_count == 2, so the next - // read_exact will hit an unexpected EOF. iter.handle.set_len(0).unwrap(); - // The iterator's cached record_count (2) > current_index (1), - // so it attempts read_exact, which fails. let err = iter.next().unwrap().unwrap_err(); assert_eq!(err.kind(), ErrorKind::UnexpectedEof); } @@ -3345,12 +3323,15 @@ mod tests { fn is_io(e: &KvpError) -> bool { matches!(e, KvpError::Io(_)) } + fn is_max_keys(e: &KvpError) -> bool { matches!(e, KvpError::MaxUniqueKeysExceeded { .. }) } + fn is_value_too_large(e: &KvpError) -> bool { matches!(e, KvpError::ValueTooLarge { .. }) } + fn is_key_too_large(e: &KvpError) -> bool { matches!(e, KvpError::KeyTooLarge { .. }) } @@ -3440,7 +3421,6 @@ mod tests { #[test] fn test_clear_if_stale_truncates_when_stale() { - // mtime (0) <= boot_time (10) → triggers set_len branch. let (store, ops, p) = mock_store(PoolMode::Safe); preload(&ops, &p, &[("a", "1")]); ops.set_boot_time(10); @@ -3468,6 +3448,17 @@ mod tests { assert_eq!(ops.lock().files.get(&p).unwrap().len(), RECORD_SIZE); } + #[test] + fn test_clear_if_stale_keeps_file_written_in_boot_second() { + let (store, ops, p) = mock_store(PoolMode::Safe); + ops.put_file(&p, vec![0u8; RECORD_SIZE], 10); + ops.set_boot_time(10); + + assert!(!store.is_stale().unwrap()); + store.clear_if_stale().unwrap(); + assert_eq!(ops.lock().files.get(&p).unwrap().len(), RECORD_SIZE); + } + #[test] fn test_delete_fails_when_iter_read_fails() { let (store, ops, p) = mock_store(PoolMode::Safe); diff --git a/libazureinit-kvp/src/vm_id.rs b/libazureinit-kvp/src/vm_id.rs index a8c1edbd..7b1e5a85 100644 --- a/libazureinit-kvp/src/vm_id.rs +++ b/libazureinit-kvp/src/vm_id.rs @@ -67,7 +67,7 @@ fn is_vm_gen1( let sysfs_efi = sysfs_efi_path.unwrap_or("/sys/firmware/efi"); let dev_efi = dev_efi_path.unwrap_or("/dev/efi"); - // If *either* efi path exists, this is Gen2; if *neither* exist, Gen1. + // If either efi path exists, this is Gen2; if neither exist, Gen1. !Path::new(sysfs_efi).exists() && !Path::new(dev_efi).exists() } @@ -176,8 +176,6 @@ mod tests { let path = dir.path().join("product_uuid"); fs::write(&path, "not-a-uuid").unwrap(); - // Gen1 (no EFI paths) but the content cannot be parsed as a UUID, - // so the raw lowercased value is returned unchanged. let actual = private_get_vm_id( Some(path.to_str().unwrap()), Some("/nonexistent_sysfs_efi"), @@ -190,9 +188,6 @@ mod tests { #[test] fn get_vm_id_public_wrapper_is_callable() { - // Exercises the public entry point. It reads the host's - // product_uuid if present, so the result is environment dependent; - // we only assert that invoking it does not panic. let _ = get_vm_id(); } diff --git a/libazureinit-kvp/tests/cli.rs b/libazureinit-kvp/tests/cli.rs index ef485314..6c6afd03 100644 --- a/libazureinit-kvp/tests/cli.rs +++ b/libazureinit-kvp/tests/cli.rs @@ -236,9 +236,6 @@ fn validation_errors_exit_two() { #[test] fn json_read_round_trips_value_with_equals_and_newline() { let dir = TempDir::new().unwrap(); - // A value containing both '=' and an embedded newline would be - // ambiguous in the default key=value text output but must survive - // round-tripping through JSON unchanged. let raw_value = "https://example.test/q=1\nline2"; let status = std::process::Command::new(env!("CARGO_BIN_EXE_libazureinit-kvp")) @@ -298,3 +295,321 @@ fn report_failure_rejects_invalid_supporting_data() { .unwrap() .contains("key=value")); } + +#[test] +fn dump_parse_diagnostics_json_reassembles_and_classifies() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + "azure-init-x|100|vm|event|a:b|id1|ts", + "one/", + ], + ))); + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + "azure-init-x|100|vm|event|a:b|id1|ts", + "two", + ], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "PROVISIONING_REPORT", "result=success"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics"], + ))); + assert!(out.contains("\"kind\":\"event\"")); + assert!(out.contains("\"chunks\":2")); + assert!(out.contains("\"message\":\"one/two\"")); + assert!(!out.contains("PROVISIONING_REPORT")); + + let out_raw = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics", "--include-raw"], + ))); + assert!(out_raw.contains("\"record\":\"raw\"")); + assert!(out_raw.contains("PROVISIONING_REPORT")); +} + +#[test] +fn dump_parse_diagnostics_text_renders_cloud_init_event() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + "CLOUD_INIT|1785187982|finish|modules-final/config-scripts_user|0e5e179d-5341-478b-8456-fbb90621bdf8|e5f01809-a7a3-4279-aa64-1f18e21eda6e", + r#"{"name":"modules-final/config-scripts_user","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.5,"msg":"scripts ran"}"#, + ], + ))); + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + "CLOUD_INIT|1785187982|start|modules-final/config-keys_to_console|0e5e179d-5341-478b-8456-fbb90621bdf8|7792621b-b339-4274-8b71-2a3dcbd2db4e", + r#"{"name":"modules-final/config-keys_to_console","type":"start","ts":"2026-07-27T21:33:24.344349+00:00","msg":"running keys_to_console"}"#, + ], + ))); + + let out = + assert_success(kvp(&with_dir(&dir, &["dump", "--parse-diagnostics"]))); + assert!(out.contains("event kind=finish")); + assert!(out.contains("agent=CLOUD_INIT")); + assert!(out.contains("boot_epoch=1785187982")); + assert!(out.contains("name=modules-final/config-scripts_user")); + assert!(out.contains("vm_id=0e5e179d-5341-478b-8456-fbb90621bdf8")); + assert!(out.contains("result=SUCCESS")); + assert!(out.contains("timestamp=2026-07-27T21:33:24.339006+00:00")); + assert!(out.contains("duration=0.5")); + assert!(out.contains("chunks=1")); + assert!(out.contains("message=scripts ran")); + assert!(out.contains("event kind=start")); +} + +#[test] +fn dump_parse_diagnostics_json_renders_cloud_init_event() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + "CLOUD_INIT|1785187982|finish|modules-final/config-scripts_user|0e5e179d-5341-478b-8456-fbb90621bdf8|e5f01809-a7a3-4279-aa64-1f18e21eda6e", + r#"{"name":"modules-final/config-scripts_user","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.5,"msg":"scripts ran"}"#, + ], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics"], + ))); + assert!(out.contains("\"record\":\"event\"")); + assert!(out.contains("\"kind\":\"finish\"")); + assert!(out.contains("\"agent\":\"CLOUD_INIT\"")); + assert!(out.contains("\"boot_epoch\":1785187982")); + assert!(out.contains("\"name\":\"modules-final/config-scripts_user\"")); + assert!(out.contains("\"vm_id\":\"0e5e179d-5341-478b-8456-fbb90621bdf8\"")); + assert!( + out.contains("\"event_id\":\"e5f01809-a7a3-4279-aa64-1f18e21eda6e\"") + ); + assert!(out.contains("\"timestamp\":\"2026-07-27T21:33:24.339006+00:00\"")); + assert!(out.contains("\"result\":\"SUCCESS\"")); + assert!(out.contains("\"duration\":0.5")); + assert!(out.contains("\"chunks\":1")); + assert!(out.contains("\"message\":\"scripts ran\"")); +} + +#[test] +fn clear_diagnostics_removes_events_and_malformed_keeps_raw() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|a:b|i1|ts", "msg"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|NOPE|c:d|i2|ts", "junk"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "PROVISIONING_REPORT", "result=success"], + ))); + + assert_success(kvp(&with_dir(&dir, &["clear", "--diagnostics"]))); + + let out = assert_success(kvp(&with_dir(&dir, &["dump"]))); + assert_eq!(out, "PROVISIONING_REPORT=result=success\n"); +} + +#[test] +fn clear_diagnostics_conflicts_with_if_stale() { + let dir = TempDir::new().unwrap(); + let output = + kvp(&with_dir(&dir, &["clear", "--diagnostics", "--if-stale"])); + assert_eq!(output.status.code(), Some(2)); +} + +#[test] +fn dump_parse_diagnostics_text_renders_all_record_kinds() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|a:b|id1|ts", "one/"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|a:b|id1|ts", "two"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "PROVISIONING_REPORT", "result=success"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|NOPE|c:d|id2|ts", "junk"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["dump", "--parse-diagnostics", "--include-raw"], + ))); + assert!(out.contains( + "event kind=event agent=a boot_epoch=100 vm_id=vm \ + name=a:b event_id=id1 timestamp=ts chunks=2 message=one/two" + )); + assert!(out.contains("raw key=PROVISIONING_REPORT value=result=success")); + assert!(out.contains("malformed key=a|100|vm|NOPE|c:d|id2|ts")); + assert!(out.contains("value=junk")); +} + +#[test] +fn dump_parse_diagnostics_tail_limits_to_last_events() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|a:b|i1|ts", "first"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|c:d|i2|ts", "second"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["dump", "--parse-diagnostics", "-n", "1"], + ))); + assert!(out.contains("second")); + assert!(!out.contains("first")); +} + +#[test] +fn dump_parse_diagnostics_tail_defaults_to_20_when_count_omitted() { + let dir = TempDir::new().unwrap(); + for i in 1..=25 { + assert_success(kvp(&with_dir( + &dir, + &[ + "write", + "--append", + &format!("a|100|vm|event|n:{i}|id{i}|ts"), + &format!("msg{i}"), + ], + ))); + } + + let out = assert_success(kvp(&with_dir( + &dir, + &["dump", "--parse-diagnostics", "--tail"], + ))); + assert_eq!(out.lines().count(), 20); + assert!(out.contains("msg25")); + assert!(out.contains("msg6")); + assert!(!out.contains("msg5")); +} + +#[test] +fn dump_parse_diagnostics_filters_by_name_substring() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|user:add|i1|ts", "u"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|ssh:key|i2|ts", "s"], + ))); + + let out = assert_success(kvp(&with_dir( + &dir, + &["dump", "--parse-diagnostics", "--name", "ssh"], + ))); + assert!(out.contains("ssh:key")); + assert!(!out.contains("user:add")); +} + +#[test] +fn dump_parse_diagnostics_include_raw_conflicts_with_filters() { + let dir = TempDir::new().unwrap(); + let output = kvp(&with_dir( + &dir, + &[ + "dump", + "--parse-diagnostics", + "--include-raw", + "--name", + "a:b", + ], + )); + assert_eq!(output.status.code(), Some(2)); +} + +#[test] +fn dump_parse_diagnostics_json_covers_events_and_malformed() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|event|a:b|i1|ts", "hello"], + ))); + assert_success(kvp(&with_dir( + &dir, + &["write", "--append", "a|100|vm|NOPE|c:d|i2|ts", "junk"], + ))); + + let dump = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics"], + ))); + assert!(dump.contains("\"kind\":\"event\"")); + assert!(dump.contains("\"record\":\"malformed\"")); + assert!(dump.contains("\"reason\":")); + + let events = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics", "--name", "a:b"], + ))); + assert!(events.contains("\"message\":\"hello\"")); + assert!(!events.contains("malformed")); +} + +#[test] +fn emit_writes_event_readable_by_dump() { + let dir = TempDir::new().unwrap(); + assert_success(kvp(&with_dir( + &dir, + &[ + "emit", + "--name", + "user:create_user", + "--message", + "created azureuser", + "--vm-id", + "vm-emit", + "--prefix", + "azure-init-test", + ], + ))); + + let out = + assert_success(kvp(&with_dir(&dir, &["dump", "--parse-diagnostics"]))); + assert!(out.contains("vm_id=vm-emit")); + assert!(out.contains("name=user:create_user")); + assert!(out.contains("message=created azureuser")); + + let json = assert_success(kvp(&with_dir( + &dir, + &["--json", "dump", "--parse-diagnostics"], + ))); + assert!(json.contains("\"kind\":\"event\"")); + assert!(json.contains("\"vm_id\":\"vm-emit\"")); + assert!(json.contains("\"name\":\"user:create_user\"")); +} diff --git a/libazureinit-kvp/tests/diagnostics.rs b/libazureinit-kvp/tests/diagnostics.rs new file mode 100644 index 00000000..01cd1a5a --- /dev/null +++ b/libazureinit-kvp/tests/diagnostics.rs @@ -0,0 +1,322 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Integration tests for the [`DiagnosticsKvp`] layer: emit/read +//! round-trips, chunk reassembly, classification, scoped clearing, and +//! the concurrent-write atomicity guarantee that keeps chunked events +//! from interleaving. + +use std::thread; + +use libazureinit_kvp::{ + DiagnosticRecord, DiagnosticsKvp, KvpPool, KvpPoolStore, PoolMode, + RecordKind, MAX_CHUNK_BYTES, +}; +use tempfile::TempDir; + +const PREFIX: &str = "azure-init-test"; +const VM_ID: &str = "vm-abc"; + +fn diagnostics(dir: &TempDir) -> DiagnosticsKvp { + let store = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap(); + DiagnosticsKvp::new(store, VM_ID, PREFIX) +} + +/// Real cloud-init reporting entries captured from a guest pool 1 file. +/// Each tuple is one record's `(key, JSON value)`. +const CLOUD_INIT_RECORDS: &[(&str, &str)] = &[ + ( + "CLOUD_INIT|1785187982|finish|modules-final/config-scripts_user|0e5e179d-5341-478b-8456-fbb90621bdf8|e5f01809-a7a3-4279-aa64-1f18e21eda6e", + r#"{"name":"modules-final/config-scripts_user","type":"finish","ts":"2026-07-27T21:33:24.339006+00:00","result":"SUCCESS","duration":0.0006448590000012189,"msg":"config-scripts_user ran successfully and took 0.001 seconds"}"#, + ), + ( + "CLOUD_INIT|1785187982|start|modules-final/config-ssh_authkey_fingerprints|0e5e179d-5341-478b-8456-fbb90621bdf8|c4d4a08d-fe93-4c7a-9be6-9a38c212e212", + r#"{"name":"modules-final/config-ssh_authkey_fingerprints","type":"start","ts":"2026-07-27T21:33:24.339170+00:00","msg":"running config-ssh_authkey_fingerprints with frequency once-per-instance"}"#, + ), + ( + "CLOUD_INIT|1785187982|finish|modules-final|0e5e179d-5341-478b-8456-fbb90621bdf8|126f969f-13fd-4b4b-a136-b7114518491f", + r#"{"name":"modules-final","type":"finish","ts":"2026-07-27T21:33:24.431885+00:00","result":"SUCCESS","duration":0.340712044,"msg":"running modules for final"}"#, + ), +]; + +#[test] +fn reads_and_parses_real_cloud_init_pool() { + let dir = TempDir::new().unwrap(); + let store = + KvpPoolStore::new_in(KvpPool::Guest, dir.path(), PoolMode::Safe) + .unwrap(); + for &(key, value) in CLOUD_INIT_RECORDS { + store.append(key, value).unwrap(); + } + + let diagnostics = DiagnosticsKvp::new(store, "", ""); + let records = diagnostics.records().unwrap(); + assert_eq!(records.len(), CLOUD_INIT_RECORDS.len()); + + for record in &records { + assert!(matches!(record, DiagnosticRecord::Decoded { .. })); + } + + match &records[0] { + DiagnosticRecord::Decoded { event, chunks } => { + assert_eq!(*chunks, 1); + assert_eq!(event.agent, "CLOUD_INIT"); + assert_eq!(event.kind, RecordKind::Finish); + assert_eq!(event.name, "modules-final/config-scripts_user"); + assert_eq!( + event.vm_id.as_deref(), + Some("0e5e179d-5341-478b-8456-fbb90621bdf8") + ); + assert_eq!(event.result.as_deref(), Some("SUCCESS")); + assert_eq!( + event.message, + "config-scripts_user ran successfully and took 0.001 seconds" + ); + } + other => panic!("expected event, got {other:?}"), + } + + match &records[1] { + DiagnosticRecord::Decoded { event, .. } => { + assert_eq!(event.kind, RecordKind::Start); + assert!(event.result.is_none()); + assert!(event.duration.is_none()); + } + other => panic!("expected event, got {other:?}"), + } +} + +#[test] +fn short_event_round_trips_as_single_record() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + assert_eq!(diag.vm_id(), VM_ID); + assert_eq!(diag.agent(), PREFIX); + + diag.emit_event("user:create_user", "created").unwrap(); + + let dumped = diag.store().dump().unwrap(); + assert_eq!(dumped.len(), 1); + assert!(dumped[0].0.ends_with("|0"), "key: {}", dumped[0].0); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + match &records[0] { + DiagnosticRecord::Decoded { + event: decoded, + chunks, + } => { + assert_eq!(*chunks, 1); + assert_eq!(decoded.kind, RecordKind::Event); + assert_eq!(decoded.vm_id.as_deref(), Some(VM_ID)); + assert_eq!(decoded.boot_epoch, diag.store().boot_epoch().unwrap()); + assert_eq!(decoded.name, "user:create_user"); + let event_id = uuid::Uuid::parse_str(&decoded.event_id) + .expect("event_id should be a valid UUID"); + assert_eq!( + event_id.get_version_num(), + 4, + "event_id should be a UUIDv4" + ); + assert_eq!(decoded.message, "created"); + } + other => panic!("expected event, got {other:?}"), + } +} + +#[test] +fn long_event_splits_across_records_and_reassembles() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + let message = "x".repeat(MAX_CHUNK_BYTES * 3 + 50); + diag.emit_event("config:dump", &message).unwrap(); + + let dumped = diag.store().dump().unwrap(); + assert_eq!(dumped.len(), 4); + let base_of = |k: &str| k.rsplit_once('|').unwrap().0.to_string(); + let base = base_of(&dumped[0].0); + assert!( + dumped.iter().all(|(k, _)| base_of(k) == base), + "all chunks share one event-key base" + ); + let mut keys: Vec = dumped.iter().map(|(k, _)| k.clone()).collect(); + keys.sort(); + keys.dedup(); + assert_eq!(keys.len(), 4, "each chunk must have a unique key"); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + match &records[0] { + DiagnosticRecord::Decoded { + event: decoded, + chunks, + } => { + assert_eq!(*chunks, 4); + assert_eq!(decoded.message, message); + } + other => panic!("expected event, got {other:?}"), + } +} + +#[test] +fn multi_chunk_event_uses_unique_keys_so_host_keeps_all() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + let message = "z".repeat(MAX_CHUNK_BYTES * 2 + 1); + diag.emit_event("big:event", &message).unwrap(); + + let dumped = diag.store().dump().unwrap(); + assert_eq!(dumped.len(), 3); + let total = dumped.len(); + let mut keys: Vec = dumped.into_iter().map(|(k, _)| k).collect(); + keys.sort(); + keys.dedup(); + assert_eq!(keys.len(), total, "chunk keys must be unique"); + + let events = diag.events().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].message, message); +} + +#[test] +fn injected_malformed_key_is_classified() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.store() + .append(&format!("{PREFIX}|100|{VM_ID}|NOPE|bad:kind|id|ts"), "junk") + .unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + assert!(matches!( + &records[0], + DiagnosticRecord::Malformed { reason, .. } if reason.contains("NOPE") + )); +} + +#[test] +fn mixed_records_round_trip_together() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit_event("a:b", "short").unwrap(); + diag.emit_event("c:d", "y".repeat(MAX_CHUNK_BYTES + 5)) + .unwrap(); + diag.store() + .append("PROVISIONING_REPORT", "result=success") + .unwrap(); + diag.store() + .append(&format!("{PREFIX}|100|{VM_ID}|NOPE|e:f|id|ts"), "junk") + .unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 4); + assert_eq!(diag.events().unwrap().len(), 2); +} + +#[test] +fn clear_removes_events_but_keeps_raw() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit_event("a:b", "e1").unwrap(); + diag.emit_event("c:d", "z".repeat(MAX_CHUNK_BYTES * 2)) + .unwrap(); + diag.store() + .append("PROVISIONING_REPORT", "result=success") + .unwrap(); + + diag.clear().unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + assert!(matches!( + &records[0], + DiagnosticRecord::Raw { key, .. } if key == "PROVISIONING_REPORT" + )); + assert!(diag.events().unwrap().is_empty()); +} + +#[test] +fn clear_removes_all_diagnostics_regardless_of_scope() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + diag.emit_event("a:b", "mine").unwrap(); + diag.store() + .append("other-agent|100|other-vm|event|x:y|id|ts", "theirs") + .unwrap(); + diag.store() + .append("p|100|vm|NOPE|c:d|id|ts", "junk") + .unwrap(); + diag.store() + .append("p|100|vm|NOPE|c:d|id|ts|0", "junk-0") + .unwrap(); + diag.store() + .append("p|100|vm|NOPE|c:d|id|ts|1", "junk-1") + .unwrap(); + diag.store() + .append("PROVISIONING_REPORT", "result=success") + .unwrap(); + + diag.clear().unwrap(); + + let records = diag.records().unwrap(); + assert_eq!(records.len(), 1); + assert!(matches!( + &records[0], + DiagnosticRecord::Raw { key, .. } if key == "PROVISIONING_REPORT" + )); + assert!(diag.events().unwrap().is_empty()); +} + +#[test] +fn emit_rejects_delimiter_in_event_fields() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + assert!(diag.emit_event("a|b", "msg").is_err()); + assert!(diag.store().dump().unwrap().is_empty()); +} + +#[test] +fn concurrent_multichunk_emits_reassemble_without_interleaving() { + let dir = TempDir::new().unwrap(); + let diag = diagnostics(&dir); + + const THREADS: usize = 5; + const PER_THREAD: usize = 8; + let len = MAX_CHUNK_BYTES * 2 + 7; + + let handles: Vec<_> = (0..THREADS) + .map(|t| { + let diag = diag.clone(); + let marker = (b'a' + t as u8) as char; + thread::spawn(move || { + for _ in 0..PER_THREAD { + let message = marker.to_string().repeat(len); + diag.emit_event(format!("thread:{marker}"), message) + .unwrap(); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let events = diag.events().unwrap(); + assert_eq!(events.len(), THREADS * PER_THREAD); + for event in &events { + assert_eq!(event.message.len(), len); + let first = event.message.chars().next().unwrap(); + assert!(event.message.chars().all(|c| c == first)); + assert_eq!(event.name, format!("thread:{first}")); + } +}